diff --git a/lib/python3.12/site-packages/dill/__pycache__/__diff.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/__diff.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3dc3f6c7c3e9907d4724c9e381ef535635daaccf Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/__diff.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/__info__.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/__info__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52f99246cadf4080fec5b14bd0998e82ed99665a Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/__info__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d48ec894b9108f46ca48248dad0dcf6f939551a4 Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/_objects.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/_objects.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6c55ba090cadcfa749852bd0ba77634395abf873 Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/_objects.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/_shims.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/_shims.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9954087b3e836f2400ccf0e3770e15ade526911 Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/_shims.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/detect.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/detect.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57d110c1e7be5bc1dd5328c9b09be9e3bd208298 Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/detect.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/logger.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9a95962d76d932948ab1efefdac607852b6414c5 Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/objtypes.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/objtypes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94af5f4fd24667320c29f08bbb10842897c67a99 Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/objtypes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/pointers.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/pointers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..165b77345a51604e42e5ae17b25d674ad05b2d4a Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/pointers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/session.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/session.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c107be96b5291ccf29648ba35f7d6988244eafe Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/session.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/settings.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/settings.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2aa3e1ba35d98808d6cbb8314652e2b14e98e87 Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/settings.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/source.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/source.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9018002b3279adc3a8ed1b4e82efaf18521b4f3d Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/source.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/__pycache__/temp.cpython-312.pyc b/lib/python3.12/site-packages/dill/__pycache__/temp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c0ceed4a23fa19c2c96069518df02538bfccb0e Binary files /dev/null and b/lib/python3.12/site-packages/dill/__pycache__/temp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..108b9eee6c54a57643d2ee61641e59275c410c46 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/__main__.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/__main__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff066af6d77b36a994c3c0b93e8a7870ec43747a Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/__main__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_abc.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_abc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14dd1f3b2afd64cfa85e08df13f83a8637bdbd85 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_abc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_check.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_check.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e44c2a31373fe0ed2c0dd2ebe769736230300a2a Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_check.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_classdef.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_classdef.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23d7398165a52086457472d32e9a4e14caa96271 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_classdef.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_dataclasses.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_dataclasses.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21a0a5bfb5a93ebcc229ebe9a4d69240daf21472 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_dataclasses.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_detect.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_detect.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2503cbb174a899ec0dcc790887194c6e9d9bb516 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_detect.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_dictviews.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_dictviews.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f7bdd1b2c3864721976e528ad747f4d22e8e7a6 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_dictviews.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_diff.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_diff.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e433a23a26d2d7f378a8bf3ba84c970901c5c159 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_diff.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_extendpickle.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_extendpickle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c4a13561e495be788ee5d87666aff544827ce0d Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_extendpickle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_fglobals.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_fglobals.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2c48239b8f6d2b911d17ab539d0f8dd669b1e42 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_fglobals.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_file.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_file.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..15b0683feef7d3eaea744ed2a0c21d9acf9b817f Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_file.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_functions.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_functions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c525cf2ed64ba95882db8fcbda6daf40a0c52d1 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_functions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_functors.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_functors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9dad24eb16df109f2c2a89f8628d4673d26a8c44 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_functors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_logger.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9317b0161fbdea6c09c4f0f21ff7d7398c2b38f0 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_mixins.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_mixins.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..27cb663ed087692d8dce6e28306bd08295b52564 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_mixins.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_module.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6727e0aa7d51e04750bd338385d974db27f7397 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_moduledict.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_moduledict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..580456973ff0e2fc78ed333381d50e30fbfe3d36 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_moduledict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_nested.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_nested.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b0abd101f799548f26890ea8fc08ff0fe0731a3 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_nested.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_objects.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_objects.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e0d062008adc6086dbf64dc34c9c440beb13faa Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_objects.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_properties.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_properties.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6ec4817dbe0d70baba63057a3c30c6f5238aae7 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_properties.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_pycapsule.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_pycapsule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b0b07c71ff1c00153c319912a819854cf94d41e Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_pycapsule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_recursive.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_recursive.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..524aa4bf35eb60e8ff47d7290113b2bc374bbed2 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_recursive.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_registered.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_registered.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..502c5e24230a639660d190e19bdbe0fe9948568c Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_registered.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_restricted.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_restricted.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e07ba2ba7ba1fe19ecd7e3b4937c308afce3118 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_restricted.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_selected.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_selected.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0f10d4db683644193033b1ff35b53187f139219a Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_selected.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_session.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_session.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc763a57393148ae74ae84738fd4c25dc936e54b Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_session.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_source.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_source.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..782ff6ea5c5dc8eeeb1e39f36f059051ec910724 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_source.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_temp.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_temp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..27b8d2ac4170e949bc2de7701bbdb48d5b60dda8 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_temp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/__pycache__/test_weakref.cpython-312.pyc b/lib/python3.12/site-packages/dill/tests/__pycache__/test_weakref.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..255079bf89a1f819a456e0bcf9c60b235c5f7521 Binary files /dev/null and b/lib/python3.12/site-packages/dill/tests/__pycache__/test_weakref.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/dill/tests/test_classdef.py b/lib/python3.12/site-packages/dill/tests/test_classdef.py new file mode 100644 index 0000000000000000000000000000000000000000..b09df50ca9632a972aa784676771a239d3a7fad4 --- /dev/null +++ b/lib/python3.12/site-packages/dill/tests/test_classdef.py @@ -0,0 +1,340 @@ +#!/usr/bin/env python +# +# Author: Mike McKerns (mmckerns @caltech and @uqfoundation) +# Copyright (c) 2008-2016 California Institute of Technology. +# Copyright (c) 2016-2024 The Uncertainty Quantification Foundation. +# License: 3-clause BSD. The full license text is available at: +# - https://github.com/uqfoundation/dill/blob/master/LICENSE + +import dill +from enum import EnumMeta +import sys +dill.settings['recurse'] = True + +# test classdefs +class _class: + def _method(self): + pass + def ok(self): + return True + +class _class2: + def __call__(self): + pass + def ok(self): + return True + +class _newclass(object): + def _method(self): + pass + def ok(self): + return True + +class _newclass2(object): + def __call__(self): + pass + def ok(self): + return True + +class _meta(type): + pass + +def __call__(self): + pass +def ok(self): + return True + +_mclass = _meta("_mclass", (object,), {"__call__": __call__, "ok": ok}) + +del __call__ +del ok + +o = _class() +oc = _class2() +n = _newclass() +nc = _newclass2() +m = _mclass() + +if sys.hexversion < 0x03090000: + import typing + class customIntList(typing.List[int]): + pass +else: + class customIntList(list[int]): + pass + +# test pickles for class instances +def test_class_instances(): + assert dill.pickles(o) + assert dill.pickles(oc) + assert dill.pickles(n) + assert dill.pickles(nc) + assert dill.pickles(m) + +def test_class_objects(): + clslist = [_class,_class2,_newclass,_newclass2,_mclass] + objlist = [o,oc,n,nc,m] + _clslist = [dill.dumps(obj) for obj in clslist] + _objlist = [dill.dumps(obj) for obj in objlist] + + for obj in clslist: + globals().pop(obj.__name__) + del clslist + for obj in ['o','oc','n','nc']: + globals().pop(obj) + del objlist + del obj + + for obj,cls in zip(_objlist,_clslist): + _cls = dill.loads(cls) + _obj = dill.loads(obj) + assert _obj.ok() + assert _cls.ok(_cls()) + if _cls.__name__ == "_mclass": + assert type(_cls).__name__ == "_meta" + +# test NoneType +def test_specialtypes(): + assert dill.pickles(type(None)) + assert dill.pickles(type(NotImplemented)) + assert dill.pickles(type(Ellipsis)) + assert dill.pickles(type(EnumMeta)) + +from collections import namedtuple +Z = namedtuple("Z", ['a','b']) +Zi = Z(0,1) +X = namedtuple("Y", ['a','b']) +X.__name__ = "X" +X.__qualname__ = "X" #XXX: name must 'match' or fails to pickle +Xi = X(0,1) +Bad = namedtuple("FakeName", ['a','b']) +Badi = Bad(0,1) +Defaults = namedtuple('Defaults', ['x', 'y'], defaults=[1]) +Defaultsi = Defaults(2) + +# test namedtuple +def test_namedtuple(): + assert Z is dill.loads(dill.dumps(Z)) + assert Zi == dill.loads(dill.dumps(Zi)) + assert X is dill.loads(dill.dumps(X)) + assert Xi == dill.loads(dill.dumps(Xi)) + assert Defaults is dill.loads(dill.dumps(Defaults)) + assert Defaultsi == dill.loads(dill.dumps(Defaultsi)) + assert Bad is not dill.loads(dill.dumps(Bad)) + assert Bad._fields == dill.loads(dill.dumps(Bad))._fields + assert tuple(Badi) == tuple(dill.loads(dill.dumps(Badi))) + + class A: + class B(namedtuple("C", ["one", "two"])): + '''docstring''' + B.__module__ = 'testing' + + a = A() + assert dill.copy(a) + + assert dill.copy(A.B).__name__ == 'B' + assert dill.copy(A.B).__qualname__.endswith('..A.B') + assert dill.copy(A.B).__doc__ == 'docstring' + assert dill.copy(A.B).__module__ == 'testing' + + from typing import NamedTuple + + def A(): + class B(NamedTuple): + x: int + return B + + assert type(dill.copy(A()(8))).__qualname__ == type(A()(8)).__qualname__ + +def test_dtype(): + try: + import numpy as np + + dti = np.dtype('int') + assert np.dtype == dill.copy(np.dtype) + assert dti == dill.copy(dti) + except ImportError: pass + + +def test_array_nested(): + try: + import numpy as np + + x = np.array([1]) + y = (x,) + assert y == dill.copy(y) + + except ImportError: pass + + +def test_array_subclass(): + try: + import numpy as np + + class TestArray(np.ndarray): + def __new__(cls, input_array, color): + obj = np.asarray(input_array).view(cls) + obj.color = color + return obj + def __array_finalize__(self, obj): + if obj is None: + return + if isinstance(obj, type(self)): + self.color = obj.color + def __getnewargs__(self): + return np.asarray(self), self.color + + a1 = TestArray(np.zeros(100), color='green') + if not dill._dill.IS_PYPY: + assert dill.pickles(a1) + assert a1.__dict__ == dill.copy(a1).__dict__ + + a2 = a1[0:9] + if not dill._dill.IS_PYPY: + assert dill.pickles(a2) + assert a2.__dict__ == dill.copy(a2).__dict__ + + class TestArray2(np.ndarray): + color = 'blue' + + a3 = TestArray2([1,2,3,4,5]) + a3.color = 'green' + if not dill._dill.IS_PYPY: + assert dill.pickles(a3) + assert a3.__dict__ == dill.copy(a3).__dict__ + + except ImportError: pass + + +def test_method_decorator(): + class A(object): + @classmethod + def test(cls): + pass + + a = A() + + res = dill.dumps(a) + new_obj = dill.loads(res) + new_obj.__class__.test() + +# test slots +class Y(object): + __slots__ = ('y', '__weakref__') + def __init__(self, y): + self.y = y + +value = 123 +y = Y(value) + +class Y2(object): + __slots__ = 'y' + def __init__(self, y): + self.y = y + +def test_slots(): + assert dill.pickles(Y) + assert dill.pickles(y) + assert dill.pickles(Y.y) + assert dill.copy(y).y == value + assert dill.copy(Y2(value)).y == value + +def test_origbases(): + assert dill.copy(customIntList).__orig_bases__ == customIntList.__orig_bases__ + +def test_attr(): + import attr + @attr.s + class A: + a = attr.ib() + + v = A(1) + assert dill.copy(v) == v + +def test_metaclass(): + class metaclass_with_new(type): + def __new__(mcls, name, bases, ns, **kwds): + cls = super().__new__(mcls, name, bases, ns, **kwds) + assert mcls is not None + assert cls.method(mcls) + return cls + def method(cls, mcls): + return isinstance(cls, mcls) + + l = locals() + exec("""class subclass_with_new(metaclass=metaclass_with_new): + def __new__(cls): + self = super().__new__(cls) + return self""", None, l) + subclass_with_new = l['subclass_with_new'] + + assert dill.copy(subclass_with_new()) + +def test_enummeta(): + from http import HTTPStatus + import enum + assert dill.copy(HTTPStatus.OK) is HTTPStatus.OK + assert dill.copy(enum.EnumMeta) is enum.EnumMeta + +def test_inherit(): #NOTE: see issue #612 + class Foo: + w = 0 + x = 1 + y = 1.1 + a = () + b = (1,) + n = None + + class Bar(Foo): + w = 2 + x = 1 + y = 1.1 + z = 0.2 + a = () + b = (1,) + c = (2,) + n = None + + Baz = dill.copy(Bar) + + import platform + is_pypy = platform.python_implementation() == 'PyPy' + assert Bar.__dict__ == Baz.__dict__ + # ints + assert 'w' in Bar.__dict__ and 'w' in Baz.__dict__ + assert Bar.__dict__['w'] is Baz.__dict__['w'] + assert 'x' in Bar.__dict__ and 'x' in Baz.__dict__ + assert Bar.__dict__['x'] is Baz.__dict__['x'] + # floats + assert 'y' in Bar.__dict__ and 'y' in Baz.__dict__ + same = Bar.__dict__['y'] is Baz.__dict__['y'] + assert same if is_pypy else not same + assert 'z' in Bar.__dict__ and 'z' in Baz.__dict__ + same = Bar.__dict__['z'] is Baz.__dict__['z'] + assert same if is_pypy else not same + # tuples + assert 'a' in Bar.__dict__ and 'a' in Baz.__dict__ + assert Bar.__dict__['a'] is Baz.__dict__['a'] + assert 'b' in Bar.__dict__ and 'b' in Baz.__dict__ + assert Bar.__dict__['b'] is not Baz.__dict__['b'] + assert 'c' in Bar.__dict__ and 'c' in Baz.__dict__ + assert Bar.__dict__['c'] is not Baz.__dict__['c'] + # None + assert 'n' in Bar.__dict__ and 'n' in Baz.__dict__ + assert Bar.__dict__['n'] is Baz.__dict__['n'] + + +if __name__ == '__main__': + test_class_instances() + test_class_objects() + test_specialtypes() + test_namedtuple() + test_dtype() + test_array_nested() + test_array_subclass() + test_method_decorator() + test_slots() + test_origbases() + test_metaclass() + test_enummeta() + test_inherit() diff --git a/lib/python3.12/site-packages/dill/tests/test_extendpickle.py b/lib/python3.12/site-packages/dill/tests/test_extendpickle.py new file mode 100644 index 0000000000000000000000000000000000000000..3b274d4bf2714c25ecddbbe13904b72759211d66 --- /dev/null +++ b/lib/python3.12/site-packages/dill/tests/test_extendpickle.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python +# +# Author: Mike McKerns (mmckerns @caltech and @uqfoundation) +# Copyright (c) 2008-2016 California Institute of Technology. +# Copyright (c) 2016-2024 The Uncertainty Quantification Foundation. +# License: 3-clause BSD. The full license text is available at: +# - https://github.com/uqfoundation/dill/blob/master/LICENSE + +import dill as pickle +from io import BytesIO as StringIO + + +def my_fn(x): + return x * 17 + + +def test_extend(): + obj = lambda : my_fn(34) + assert obj() == 578 + + obj_io = StringIO() + pickler = pickle.Pickler(obj_io) + pickler.dump(obj) + + obj_str = obj_io.getvalue() + + obj2_io = StringIO(obj_str) + unpickler = pickle.Unpickler(obj2_io) + obj2 = unpickler.load() + + assert obj2() == 578 + + +def test_isdill(): + obj_io = StringIO() + pickler = pickle.Pickler(obj_io) + assert pickle._dill.is_dill(pickler) is True + + pickler = pickle._dill.StockPickler(obj_io) + assert pickle._dill.is_dill(pickler) is False + + try: + import multiprocess as mp + pickler = mp.reduction.ForkingPickler(obj_io) + assert pickle._dill.is_dill(pickler, child=True) is True + assert pickle._dill.is_dill(pickler, child=False) is False + except Exception: + pass + + +if __name__ == '__main__': + test_extend() + test_isdill() diff --git a/lib/python3.12/site-packages/dill/tests/test_logger.py b/lib/python3.12/site-packages/dill/tests/test_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..b878a1054b36bfc4e39e3a575ef407c094ebd3a8 --- /dev/null +++ b/lib/python3.12/site-packages/dill/tests/test_logger.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python + +# Author: Leonardo Gama (@leogama) +# Copyright (c) 2022-2024 The Uncertainty Quantification Foundation. +# License: 3-clause BSD. The full license text is available at: +# - https://github.com/uqfoundation/dill/blob/master/LICENSE + +import logging +import re +import tempfile + +import dill +from dill import detect +from dill.logger import stderr_handler, adapter as logger + +try: + from StringIO import StringIO +except ImportError: + from io import StringIO + +test_obj = {'a': (1, 2), 'b': object(), 'f': lambda x: x**2, 'big': list(range(10))} + +def test_logging(should_trace): + buffer = StringIO() + handler = logging.StreamHandler(buffer) + logger.addHandler(handler) + try: + dill.dumps(test_obj) + if should_trace: + regex = re.compile(r'(\S*┬ \w.*[^)]' # begin pickling object + r'|│*└ # \w.* \[\d+ (\wi)?B])' # object written (with size) + ) + for line in buffer.getvalue().splitlines(): + assert regex.fullmatch(line) + return buffer.getvalue() + else: + assert buffer.getvalue() == "" + finally: + logger.removeHandler(handler) + buffer.close() + +def test_trace_to_file(stream_trace): + file = tempfile.NamedTemporaryFile(mode='r') + with detect.trace(file.name, mode='w'): + dill.dumps(test_obj) + file_trace = file.read() + file.close() + # Apparently, objects can change location in memory... + reghex = re.compile(r'0x[0-9A-Za-z]+') + file_trace, stream_trace = reghex.sub('0x', file_trace), reghex.sub('0x', stream_trace) + # PyPy prints dictionary contents with repr(dict)... + regdict = re.compile(r'(dict\.__repr__ of ).*') + file_trace, stream_trace = regdict.sub(r'\1{}>', file_trace), regdict.sub(r'\1{}>', stream_trace) + assert file_trace == stream_trace + +if __name__ == '__main__': + logger.removeHandler(stderr_handler) + test_logging(should_trace=False) + detect.trace(True) + test_logging(should_trace=True) + detect.trace(False) + test_logging(should_trace=False) + + loglevel = logging.ERROR + logger.setLevel(loglevel) + with detect.trace(): + stream_trace = test_logging(should_trace=True) + test_logging(should_trace=False) + assert logger.getEffectiveLevel() == loglevel + test_trace_to_file(stream_trace) diff --git a/lib/python3.12/site-packages/dill/tests/test_nested.py b/lib/python3.12/site-packages/dill/tests/test_nested.py new file mode 100644 index 0000000000000000000000000000000000000000..9109f55c86b0539632a7bde27e8172c8ba714acd --- /dev/null +++ b/lib/python3.12/site-packages/dill/tests/test_nested.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python +# +# Author: Mike McKerns (mmckerns @caltech and @uqfoundation) +# Copyright (c) 2008-2016 California Institute of Technology. +# Copyright (c) 2016-2024 The Uncertainty Quantification Foundation. +# License: 3-clause BSD. The full license text is available at: +# - https://github.com/uqfoundation/dill/blob/master/LICENSE +""" +test dill's ability to handle nested functions +""" + +import os +import math + +import dill as pickle +pickle.settings['recurse'] = True + + +# the nested function: pickle should fail here, but dill is ok. +def adder(augend): + zero = [0] + + def inner(addend): + return addend + augend + zero[0] + return inner + + +# rewrite the nested function using a class: standard pickle should work here. +class cadder(object): + def __init__(self, augend): + self.augend = augend + self.zero = [0] + + def __call__(self, addend): + return addend + self.augend + self.zero[0] + + +# rewrite again, but as an old-style class +class c2adder: + def __init__(self, augend): + self.augend = augend + self.zero = [0] + + def __call__(self, addend): + return addend + self.augend + self.zero[0] + + +# some basic class stuff +class basic(object): + pass + + +class basic2: + pass + + +x = 5 +y = 1 + + +def test_basic(): + a = [0, 1, 2] + pa = pickle.dumps(a) + pmath = pickle.dumps(math) #XXX: FAILS in pickle + pmap = pickle.dumps(map) + # ... + la = pickle.loads(pa) + lmath = pickle.loads(pmath) + lmap = pickle.loads(pmap) + assert list(map(math.sin, a)) == list(lmap(lmath.sin, la)) + + +def test_basic_class(): + pbasic2 = pickle.dumps(basic2) + _pbasic2 = pickle.loads(pbasic2)() + pbasic = pickle.dumps(basic) + _pbasic = pickle.loads(pbasic)() + + +def test_c2adder(): + pc2adder = pickle.dumps(c2adder) + pc2add5 = pickle.loads(pc2adder)(x) + assert pc2add5(y) == x+y + + +def test_pickled_cadder(): + pcadder = pickle.dumps(cadder) + pcadd5 = pickle.loads(pcadder)(x) + assert pcadd5(y) == x+y + + +def test_raw_adder_and_inner(): + add5 = adder(x) + assert add5(y) == x+y + + +def test_pickled_adder(): + padder = pickle.dumps(adder) + padd5 = pickle.loads(padder)(x) + assert padd5(y) == x+y + + +def test_pickled_inner(): + add5 = adder(x) + pinner = pickle.dumps(add5) #XXX: FAILS in pickle + p5add = pickle.loads(pinner) + assert p5add(y) == x+y + + +def test_moduledict_where_not_main(): + try: + from . import test_moduledict + except ImportError: + import test_moduledict + name = 'test_moduledict.py' + if os.path.exists(name) and os.path.exists(name+'c'): + os.remove(name+'c') + + if os.path.exists(name) and hasattr(test_moduledict, "__cached__") \ + and os.path.exists(test_moduledict.__cached__): + os.remove(getattr(test_moduledict, "__cached__")) + + if os.path.exists("__pycache__") and not os.listdir("__pycache__"): + os.removedirs("__pycache__") + + +if __name__ == '__main__': + test_basic() + test_basic_class() + test_c2adder() + test_pickled_cadder() + test_raw_adder_and_inner() + test_pickled_adder() + test_pickled_inner() + test_moduledict_where_not_main() diff --git a/lib/python3.12/site-packages/dill/tests/test_restricted.py b/lib/python3.12/site-packages/dill/tests/test_restricted.py new file mode 100644 index 0000000000000000000000000000000000000000..cdb773e36df6304e44c5a883f9f4e5beaee0d674 --- /dev/null +++ b/lib/python3.12/site-packages/dill/tests/test_restricted.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python +# +# Author: Kirill Makhonin (@kirillmakhonin) +# Copyright (c) 2008-2016 California Institute of Technology. +# Copyright (c) 2016-2024 The Uncertainty Quantification Foundation. +# License: 3-clause BSD. The full license text is available at: +# - https://github.com/uqfoundation/dill/blob/master/LICENSE + +import dill + +class RestrictedType: + def __bool__(*args, **kwargs): + raise Exception('Restricted function') + + __eq__ = __lt__ = __le__ = __ne__ = __gt__ = __ge__ = __hash__ = __bool__ + +glob_obj = RestrictedType() + +def restricted_func(): + a = glob_obj + +def test_function_with_restricted_object(): + deserialized = dill.loads(dill.dumps(restricted_func, recurse=True)) + + +if __name__ == '__main__': + test_function_with_restricted_object() diff --git a/lib/python3.12/site-packages/dill/tests/test_source.py b/lib/python3.12/site-packages/dill/tests/test_source.py new file mode 100644 index 0000000000000000000000000000000000000000..8771a542a27a15140139b3de6645d9bdd0615eed --- /dev/null +++ b/lib/python3.12/site-packages/dill/tests/test_source.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python +# +# Author: Mike McKerns (mmckerns @caltech and @uqfoundation) +# Copyright (c) 2008-2016 California Institute of Technology. +# Copyright (c) 2016-2024 The Uncertainty Quantification Foundation. +# License: 3-clause BSD. The full license text is available at: +# - https://github.com/uqfoundation/dill/blob/master/LICENSE + +from dill.source import getsource, getname, _wrap, likely_import +from dill.source import getimportable +from dill._dill import IS_PYPY + +import sys +PY310b = 0x30a00b1 + +f = lambda x: x**2 +def g(x): return f(x) - x + +def h(x): + def g(x): return x + return g(x) - x + +class Foo(object): + def bar(self, x): + return x*x+x +_foo = Foo() + +def add(x,y): + return x+y + +# yes, same as 'f', but things are tricky when it comes to pointers +squared = lambda x:x**2 + +class Bar: + pass +_bar = Bar() + + # inspect.getsourcelines # dill.source.getblocks +def test_getsource(): + assert getsource(f) == 'f = lambda x: x**2\n' + assert getsource(g) == 'def g(x): return f(x) - x\n' + assert getsource(h) == 'def h(x):\n def g(x): return x\n return g(x) - x\n' + assert getname(f) == 'f' + assert getname(g) == 'g' + assert getname(h) == 'h' + assert _wrap(f)(4) == 16 + assert _wrap(g)(4) == 12 + assert _wrap(h)(4) == 0 + + assert getname(Foo) == 'Foo' + assert getname(Bar) == 'Bar' + assert getsource(Bar) == 'class Bar:\n pass\n' + assert getsource(Foo) == 'class Foo(object):\n def bar(self, x):\n return x*x+x\n' + #XXX: add getsource for _foo, _bar + +# test itself +def test_itself(): + assert likely_import(likely_import)=='from dill.source import likely_import\n' + +# builtin functions and objects +def test_builtin(): + assert likely_import(pow) == 'pow\n' + assert likely_import(100) == '100\n' + assert likely_import(True) == 'True\n' + assert likely_import(pow, explicit=True) == 'from builtins import pow\n' + assert likely_import(100, explicit=True) == '100\n' + assert likely_import(True, explicit=True) == 'True\n' + # this is kinda BS... you can't import a None + assert likely_import(None) == 'None\n' + assert likely_import(None, explicit=True) == 'None\n' + + +# other imported functions +def test_imported(): + from math import sin + assert likely_import(sin) == 'from math import sin\n' + +# interactively defined functions +def test_dynamic(): + assert likely_import(add) == 'from %s import add\n' % __name__ + # interactive lambdas + assert likely_import(squared) == 'from %s import squared\n' % __name__ + +# classes and class instances +def test_classes(): + from io import BytesIO as StringIO + y = "from _io import BytesIO\n" + x = y if (IS_PYPY or sys.hexversion >= PY310b) else "from io import BytesIO\n" + s = StringIO() + + assert likely_import(StringIO) == x + assert likely_import(s) == y + # interactively defined classes and class instances + assert likely_import(Foo) == 'from %s import Foo\n' % __name__ + assert likely_import(_foo) == 'from %s import Foo\n' % __name__ + + +# test getimportable +def test_importable(): + assert getimportable(add) == 'from %s import add\n' % __name__ + assert getimportable(squared) == 'from %s import squared\n' % __name__ + assert getimportable(Foo) == 'from %s import Foo\n' % __name__ + assert getimportable(Foo.bar) == 'from %s import bar\n' % __name__ + assert getimportable(_foo.bar) == 'from %s import bar\n' % __name__ + assert getimportable(None) == 'None\n' + assert getimportable(100) == '100\n' + + assert getimportable(add, byname=False) == 'def add(x,y):\n return x+y\n' + assert getimportable(squared, byname=False) == 'squared = lambda x:x**2\n' + assert getimportable(None, byname=False) == 'None\n' + assert getimportable(Bar, byname=False) == 'class Bar:\n pass\n' + assert getimportable(Foo, byname=False) == 'class Foo(object):\n def bar(self, x):\n return x*x+x\n' + assert getimportable(Foo.bar, byname=False) == 'def bar(self, x):\n return x*x+x\n' + assert getimportable(Foo.bar, byname=True) == 'from %s import bar\n' % __name__ + assert getimportable(Foo.bar, alias='memo', byname=True) == 'from %s import bar as memo\n' % __name__ + assert getimportable(Foo, alias='memo', byname=True) == 'from %s import Foo as memo\n' % __name__ + assert getimportable(squared, alias='memo', byname=True) == 'from %s import squared as memo\n' % __name__ + assert getimportable(squared, alias='memo', byname=False) == 'memo = squared = lambda x:x**2\n' + assert getimportable(add, alias='memo', byname=False) == 'def add(x,y):\n return x+y\n\nmemo = add\n' + assert getimportable(None, alias='memo', byname=False) == 'memo = None\n' + assert getimportable(100, alias='memo', byname=False) == 'memo = 100\n' + assert getimportable(add, explicit=True) == 'from %s import add\n' % __name__ + assert getimportable(squared, explicit=True) == 'from %s import squared\n' % __name__ + assert getimportable(Foo, explicit=True) == 'from %s import Foo\n' % __name__ + assert getimportable(Foo.bar, explicit=True) == 'from %s import bar\n' % __name__ + assert getimportable(_foo.bar, explicit=True) == 'from %s import bar\n' % __name__ + assert getimportable(None, explicit=True) == 'None\n' + assert getimportable(100, explicit=True) == '100\n' + + +def test_numpy(): + try: + from numpy import array + x = array([1,2,3]) + assert getimportable(x) == 'from numpy import array\narray([1, 2, 3])\n' + assert getimportable(array) == 'from %s import array\n' % array.__module__ + assert getimportable(x, byname=False) == 'from numpy import array\narray([1, 2, 3])\n' + assert getimportable(array, byname=False) == 'from %s import array\n' % array.__module__ + except ImportError: pass + +#NOTE: if before likely_import(pow), will cause pow to throw AssertionError +def test_foo(): + assert getimportable(_foo, byname=False).startswith("import dill\nclass Foo(object):\n def bar(self, x):\n return x*x+x\ndill.loads(") + +if __name__ == '__main__': + test_getsource() + test_itself() + test_builtin() + test_imported() + test_dynamic() + test_classes() + test_importable() + test_numpy() + test_foo() diff --git a/lib/python3.12/site-packages/dill/tests/test_weakref.py b/lib/python3.12/site-packages/dill/tests/test_weakref.py new file mode 100644 index 0000000000000000000000000000000000000000..df5cbce9309e45b56a84b1d605f69586b68876cf --- /dev/null +++ b/lib/python3.12/site-packages/dill/tests/test_weakref.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python +# +# Author: Mike McKerns (mmckerns @caltech and @uqfoundation) +# Copyright (c) 2008-2016 California Institute of Technology. +# Copyright (c) 2016-2024 The Uncertainty Quantification Foundation. +# License: 3-clause BSD. The full license text is available at: +# - https://github.com/uqfoundation/dill/blob/master/LICENSE + +import dill +dill.settings['recurse'] = True +import weakref + +class _class: + def _method(self): + pass + +class _callable_class: + def __call__(self): + pass + +def _function(): + pass + + +def test_weakref(): + o = _class() + oc = _callable_class() + f = _function + x = _class + + # ReferenceType + r = weakref.ref(o) + d_r = weakref.ref(_class()) + fr = weakref.ref(f) + xr = weakref.ref(x) + + # ProxyType + p = weakref.proxy(o) + d_p = weakref.proxy(_class()) + + # CallableProxyType + cp = weakref.proxy(oc) + d_cp = weakref.proxy(_callable_class()) + fp = weakref.proxy(f) + xp = weakref.proxy(x) + + objlist = [r,d_r,fr,xr, p,d_p, cp,d_cp,fp,xp] + #dill.detect.trace(True) + + for obj in objlist: + res = dill.detect.errors(obj) + if res: + print ("%r:\n %s" % (obj, res)) + # else: + # print ("PASS: %s" % obj) + assert not res + +def test_dictproxy(): + from dill._dill import DictProxyType + try: + m = DictProxyType({"foo": "bar"}) + except Exception: + m = type.__dict__ + mp = dill.copy(m) + assert mp.items() == m.items() + + +if __name__ == '__main__': + test_weakref() + from dill._dill import IS_PYPY + if not IS_PYPY: + test_dictproxy() diff --git a/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/INSTALLER b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/License.txt b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/License.txt new file mode 100644 index 0000000000000000000000000000000000000000..b491c70e0aef319022ded661e111ddbd45b8a17f --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/License.txt @@ -0,0 +1,1568 @@ +End User License Agreement +-------------------------- + + +Preface +------- + +The Software License Agreement in Chapter 1 and the Supplement +in Chapter 2 contain license terms and conditions that govern +the use of NVIDIA software. By accepting this agreement, you +agree to comply with all the terms and conditions applicable +to the product(s) included herein. + + +NVIDIA Driver + + +Description + +This package contains the operating system driver and +fundamental system software components for NVIDIA GPUs. + + +NVIDIA CUDA Toolkit + + +Description + +The NVIDIA CUDA Toolkit provides command-line and graphical +tools for building, debugging and optimizing the performance +of applications accelerated by NVIDIA GPUs, runtime and math +libraries, and documentation including programming guides, +user manuals, and API references. + + +Default Install Location of CUDA Toolkit + +Windows platform: + +%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v#.# + +Linux platform: + +/usr/local/cuda-#.# + +Mac platform: + +/Developer/NVIDIA/CUDA-#.# + + +NVIDIA CUDA Samples + + +Description + +This package includes over 100+ CUDA examples that demonstrate +various CUDA programming principles, and efficient CUDA +implementation of algorithms in specific application domains. + + +Default Install Location of CUDA Samples + +Windows platform: + +%ProgramData%\NVIDIA Corporation\CUDA Samples\v#.# + +Linux platform: + +/usr/local/cuda-#.#/samples + +and + +$HOME/NVIDIA_CUDA-#.#_Samples + +Mac platform: + +/Developer/NVIDIA/CUDA-#.#/samples + + +NVIDIA Nsight Visual Studio Edition (Windows only) + + +Description + +NVIDIA Nsight Development Platform, Visual Studio Edition is a +development environment integrated into Microsoft Visual +Studio that provides tools for debugging, profiling, analyzing +and optimizing your GPU computing and graphics applications. + + +Default Install Location of Nsight Visual Studio Edition + +Windows platform: + +%ProgramFiles(x86)%\NVIDIA Corporation\Nsight Visual Studio Edition #.# + + +1. License Agreement for NVIDIA Software Development Kits +--------------------------------------------------------- + + +Release Date: July 26, 2018 +--------------------------- + + +Important NoticeRead before downloading, installing, +copying or using the licensed software: +------------------------------------------------------- + +This license agreement, including exhibits attached +("Agreement”) is a legal agreement between you and NVIDIA +Corporation ("NVIDIA") and governs your use of a NVIDIA +software development kit (“SDK”). + +Each SDK has its own set of software and materials, but here +is a description of the types of items that may be included in +a SDK: source code, header files, APIs, data sets and assets +(examples include images, textures, models, scenes, videos, +native API input/output files), binary software, sample code, +libraries, utility programs, programming code and +documentation. + +This Agreement can be accepted only by an adult of legal age +of majority in the country in which the SDK is used. + +If you are entering into this Agreement on behalf of a company +or other legal entity, you represent that you have the legal +authority to bind the entity to this Agreement, in which case +“you” will mean the entity you represent. + +If you don’t have the required age or authority to accept +this Agreement, or if you don’t accept all the terms and +conditions of this Agreement, do not download, install or use +the SDK. + +You agree to use the SDK only for purposes that are permitted +by (a) this Agreement, and (b) any applicable law, regulation +or generally accepted practices or guidelines in the relevant +jurisdictions. + + +1.1. License + + +1.1.1. License Grant + +Subject to the terms of this Agreement, NVIDIA hereby grants +you a non-exclusive, non-transferable license, without the +right to sublicense (except as expressly provided in this +Agreement) to: + + 1. Install and use the SDK, + + 2. Modify and create derivative works of sample source code + delivered in the SDK, and + + 3. Distribute those portions of the SDK that are identified + in this Agreement as distributable, as incorporated in + object code format into a software application that meets + the distribution requirements indicated in this Agreement. + + +1.1.2. Distribution Requirements + +These are the distribution requirements for you to exercise +the distribution grant: + + 1. Your application must have material additional + functionality, beyond the included portions of the SDK. + + 2. The distributable portions of the SDK shall only be + accessed by your application. + + 3. The following notice shall be included in modifications + and derivative works of sample source code distributed: + “This software contains source code provided by NVIDIA + Corporation.” + + 4. Unless a developer tool is identified in this Agreement + as distributable, it is delivered for your internal use + only. + + 5. The terms under which you distribute your application + must be consistent with the terms of this Agreement, + including (without limitation) terms relating to the + license grant and license restrictions and protection of + NVIDIA’s intellectual property rights. Additionally, you + agree that you will protect the privacy, security and + legal rights of your application users. + + 6. You agree to notify NVIDIA in writing of any known or + suspected distribution or use of the SDK not in compliance + with the requirements of this Agreement, and to enforce + the terms of your agreements with respect to distributed + SDK. + + +1.1.3. Authorized Users + +You may allow employees and contractors of your entity or of +your subsidiary(ies) to access and use the SDK from your +secure network to perform work on your behalf. + +If you are an academic institution you may allow users +enrolled or employed by the academic institution to access and +use the SDK from your secure network. + +You are responsible for the compliance with the terms of this +Agreement by your authorized users. If you become aware that +your authorized users didn’t follow the terms of this +Agreement, you agree to take reasonable steps to resolve the +non-compliance and prevent new occurrences. + + +1.1.4. Pre-Release SDK + +The SDK versions identified as alpha, beta, preview or +otherwise as pre-release, may not be fully functional, may +contain errors or design flaws, and may have reduced or +different security, privacy, accessibility, availability, and +reliability standards relative to commercial versions of +NVIDIA software and materials. Use of a pre-release SDK may +result in unexpected results, loss of data, project delays or +other unpredictable damage or loss. + +You may use a pre-release SDK at your own risk, understanding +that pre-release SDKs are not intended for use in production +or business-critical systems. + +NVIDIA may choose not to make available a commercial version +of any pre-release SDK. NVIDIA may also choose to abandon +development and terminate the availability of a pre-release +SDK at any time without liability. + + +1.1.5. Updates + +NVIDIA may, at its option, make available patches, workarounds +or other updates to this SDK. Unless the updates are provided +with their separate governing terms, they are deemed part of +the SDK licensed to you as provided in this Agreement. You +agree that the form and content of the SDK that NVIDIA +provides may change without prior notice to you. While NVIDIA +generally maintains compatibility between versions, NVIDIA may +in some cases make changes that introduce incompatibilities in +future versions of the SDK. + + +1.1.6. Third Party Licenses + +The SDK may come bundled with, or otherwise include or be +distributed with, third party software licensed by a NVIDIA +supplier and/or open source software provided under an open +source license. Use of third party software is subject to the +third-party license terms, or in the absence of third party +terms, the terms of this Agreement. Copyright to third party +software is held by the copyright holders indicated in the +third-party software or license. + + +1.1.7. Reservation of Rights + +NVIDIA reserves all rights, title, and interest in and to the +SDK, not expressly granted to you under this Agreement. + + +1.2. Limitations + +The following license limitations apply to your use of the +SDK: + + 1. You may not reverse engineer, decompile or disassemble, + or remove copyright or other proprietary notices from any + portion of the SDK or copies of the SDK. + + 2. Except as expressly provided in this Agreement, you may + not copy, sell, rent, sublicense, transfer, distribute, + modify, or create derivative works of any portion of the + SDK. For clarity, you may not distribute or sublicense the + SDK as a stand-alone product. + + 3. Unless you have an agreement with NVIDIA for this + purpose, you may not indicate that an application created + with the SDK is sponsored or endorsed by NVIDIA. + + 4. You may not bypass, disable, or circumvent any + encryption, security, digital rights management or + authentication mechanism in the SDK. + + 5. You may not use the SDK in any manner that would cause it + to become subject to an open source software license. As + examples, licenses that require as a condition of use, + modification, and/or distribution that the SDK be: + + a. Disclosed or distributed in source code form; + + b. Licensed for the purpose of making derivative works; + or + + c. Redistributable at no charge. + + 6. Unless you have an agreement with NVIDIA for this + purpose, you may not use the SDK with any system or + application where the use or failure of the system or + application can reasonably be expected to threaten or + result in personal injury, death, or catastrophic loss. + Examples include use in avionics, navigation, military, + medical, life support or other life critical applications. + NVIDIA does not design, test or manufacture the SDK for + these critical uses and NVIDIA shall not be liable to you + or any third party, in whole or in part, for any claims or + damages arising from such uses. + + 7. You agree to defend, indemnify and hold harmless NVIDIA + and its affiliates, and their respective employees, + contractors, agents, officers and directors, from and + against any and all claims, damages, obligations, losses, + liabilities, costs or debt, fines, restitutions and + expenses (including but not limited to attorney’s fees + and costs incident to establishing the right of + indemnification) arising out of or related to your use of + the SDK outside of the scope of this Agreement, or not in + compliance with its terms. + + +1.3. Ownership + + 1. NVIDIA or its licensors hold all rights, title and + interest in and to the SDK and its modifications and + derivative works, including their respective intellectual + property rights, subject to your rights described in this + section. This SDK may include software and materials from + NVIDIA’s licensors, and these licensors are intended + third party beneficiaries that may enforce this Agreement + with respect to their intellectual property rights. + + 2. You hold all rights, title and interest in and to your + applications and your derivative works of the sample + source code delivered in the SDK, including their + respective intellectual property rights, subject to + NVIDIA’s rights described in this section. + + 3. You may, but don’t have to, provide to NVIDIA + suggestions, feature requests or other feedback regarding + the SDK, including possible enhancements or modifications + to the SDK. For any feedback that you voluntarily provide, + you hereby grant NVIDIA and its affiliates a perpetual, + non-exclusive, worldwide, irrevocable license to use, + reproduce, modify, license, sublicense (through multiple + tiers of sublicensees), and distribute (through multiple + tiers of distributors) it without the payment of any + royalties or fees to you. NVIDIA will use feedback at its + choice. NVIDIA is constantly looking for ways to improve + its products, so you may send feedback to NVIDIA through + the developer portal at https://developer.nvidia.com. + + +1.4. No Warranties + +THE SDK IS PROVIDED BY NVIDIA “AS IS” AND “WITH ALL +FAULTS.” TO THE MAXIMUM EXTENT PERMITTED BY LAW, NVIDIA AND +ITS AFFILIATES EXPRESSLY DISCLAIM ALL WARRANTIES OF ANY KIND +OR NATURE, WHETHER EXPRESS, IMPLIED OR STATUTORY, INCLUDING, +BUT NOT LIMITED TO, ANY WARRANTIES OF MERCHANTABILITY, FITNESS +FOR A PARTICULAR PURPOSE, TITLE, NON-INFRINGEMENT, OR THE +ABSENCE OF ANY DEFECTS THEREIN, WHETHER LATENT OR PATENT. NO +WARRANTY IS MADE ON THE BASIS OF TRADE USAGE, COURSE OF +DEALING OR COURSE OF TRADE. + + +1.5. Limitation of Liability + +TO THE MAXIMUM EXTENT PERMITTED BY LAW, NVIDIA AND ITS +AFFILIATES SHALL NOT BE LIABLE FOR ANY SPECIAL, INCIDENTAL, +PUNITIVE OR CONSEQUENTIAL DAMAGES, OR ANY LOST PROFITS, LOSS +OF USE, LOSS OF DATA OR LOSS OF GOODWILL, OR THE COSTS OF +PROCURING SUBSTITUTE PRODUCTS, ARISING OUT OF OR IN CONNECTION +WITH THIS AGREEMENT OR THE USE OR PERFORMANCE OF THE SDK, +WHETHER SUCH LIABILITY ARISES FROM ANY CLAIM BASED UPON BREACH +OF CONTRACT, BREACH OF WARRANTY, TORT (INCLUDING NEGLIGENCE), +PRODUCT LIABILITY OR ANY OTHER CAUSE OF ACTION OR THEORY OF +LIABILITY. IN NO EVENT WILL NVIDIA’S AND ITS AFFILIATES +TOTAL CUMULATIVE LIABILITY UNDER OR ARISING OUT OF THIS +AGREEMENT EXCEED US$10.00. THE NATURE OF THE LIABILITY OR THE +NUMBER OF CLAIMS OR SUITS SHALL NOT ENLARGE OR EXTEND THIS +LIMIT. + +These exclusions and limitations of liability shall apply +regardless if NVIDIA or its affiliates have been advised of +the possibility of such damages, and regardless of whether a +remedy fails its essential purpose. These exclusions and +limitations of liability form an essential basis of the +bargain between the parties, and, absent any of these +exclusions or limitations of liability, the provisions of this +Agreement, including, without limitation, the economic terms, +would be substantially different. + + +1.6. Termination + + 1. This Agreement will continue to apply until terminated by + either you or NVIDIA as described below. + + 2. If you want to terminate this Agreement, you may do so by + stopping to use the SDK. + + 3. NVIDIA may, at any time, terminate this Agreement if: + + a. (i) you fail to comply with any term of this + Agreement and the non-compliance is not fixed within + thirty (30) days following notice from NVIDIA (or + immediately if you violate NVIDIA’s intellectual + property rights); + + b. (ii) you commence or participate in any legal + proceeding against NVIDIA with respect to the SDK; or + + c. (iii) NVIDIA decides to no longer provide the SDK in + a country or, in NVIDIA’s sole discretion, the + continued use of it is no longer commercially viable. + + 4. Upon any termination of this Agreement, you agree to + promptly discontinue use of the SDK and destroy all copies + in your possession or control. Your prior distributions in + accordance with this Agreement are not affected by the + termination of this Agreement. Upon written request, you + will certify in writing that you have complied with your + commitments under this section. Upon any termination of + this Agreement all provisions survive except for the + license grant provisions. + + +1.7. General + +If you wish to assign this Agreement or your rights and +obligations, including by merger, consolidation, dissolution +or operation of law, contact NVIDIA to ask for permission. Any +attempted assignment not approved by NVIDIA in writing shall +be void and of no effect. NVIDIA may assign, delegate or +transfer this Agreement and its rights and obligations, and if +to a non-affiliate you will be notified. + +You agree to cooperate with NVIDIA and provide reasonably +requested information to verify your compliance with this +Agreement. + +This Agreement will be governed in all respects by the laws of +the United States and of the State of Delaware as those laws +are applied to contracts entered into and performed entirely +within Delaware by Delaware residents, without regard to the +conflicts of laws principles. The United Nations Convention on +Contracts for the International Sale of Goods is specifically +disclaimed. You agree to all terms of this Agreement in the +English language. + +The state or federal courts residing in Santa Clara County, +California shall have exclusive jurisdiction over any dispute +or claim arising out of this Agreement. Notwithstanding this, +you agree that NVIDIA shall still be allowed to apply for +injunctive remedies or an equivalent type of urgent legal +relief in any jurisdiction. + +If any court of competent jurisdiction determines that any +provision of this Agreement is illegal, invalid or +unenforceable, such provision will be construed as limited to +the extent necessary to be consistent with and fully +enforceable under the law and the remaining provisions will +remain in full force and effect. Unless otherwise specified, +remedies are cumulative. + +Each party acknowledges and agrees that the other is an +independent contractor in the performance of this Agreement. + +The SDK has been developed entirely at private expense and is +“commercial items” consisting of “commercial computer +software” and “commercial computer software +documentation” provided with RESTRICTED RIGHTS. Use, +duplication or disclosure by the U.S. Government or a U.S. +Government subcontractor is subject to the restrictions in +this Agreement pursuant to DFARS 227.7202-3(a) or as set forth +in subparagraphs (c)(1) and (2) of the Commercial Computer +Software - Restricted Rights clause at FAR 52.227-19, as +applicable. Contractor/manufacturer is NVIDIA, 2788 San Tomas +Expressway, Santa Clara, CA 95051. + +The SDK is subject to United States export laws and +regulations. You agree that you will not ship, transfer or +export the SDK into any country, or use the SDK in any manner, +prohibited by the United States Bureau of Industry and +Security or economic sanctions regulations administered by the +U.S. Department of Treasury’s Office of Foreign Assets +Control (OFAC), or any applicable export laws, restrictions or +regulations. These laws include restrictions on destinations, +end users and end use. By accepting this Agreement, you +confirm that you are not a resident or citizen of any country +currently embargoed by the U.S. and that you are not otherwise +prohibited from receiving the SDK. + +Any notice delivered by NVIDIA to you under this Agreement +will be delivered via mail, email or fax. You agree that any +notices that NVIDIA sends you electronically will satisfy any +legal communication requirements. Please direct your legal +notices or other correspondence to NVIDIA Corporation, 2788 +San Tomas Expressway, Santa Clara, California 95051, United +States of America, Attention: Legal Department. + +This Agreement and any exhibits incorporated into this +Agreement constitute the entire agreement of the parties with +respect to the subject matter of this Agreement and supersede +all prior negotiations or documentation exchanged between the +parties relating to this SDK license. Any additional and/or +conflicting terms on documents issued by you are null, void, +and invalid. Any amendment or waiver under this Agreement +shall be in writing and signed by representatives of both +parties. + + +2. CUDA Toolkit Supplement to Software License Agreement for +NVIDIA Software Development Kits +------------------------------------------------------------ + + +Release date: August 16, 2018 +----------------------------- + +The terms in this supplement govern your use of the NVIDIA +CUDA Toolkit SDK under the terms of your license agreement +(“Agreement”) as modified by this supplement. Capitalized +terms used but not defined below have the meaning assigned to +them in the Agreement. + +This supplement is an exhibit to the Agreement and is +incorporated as an integral part of the Agreement. In the +event of conflict between the terms in this supplement and the +terms in the Agreement, the terms in this supplement govern. + + +2.1. License Scope + +The SDK is licensed for you to develop applications only for +use in systems with NVIDIA GPUs. + + +2.2. Distribution + +The portions of the SDK that are distributable under the +Agreement are listed in Attachment A. + + +2.3. Operating Systems + +Those portions of the SDK designed exclusively for use on the +Linux or FreeBSD operating systems, or other operating systems +derived from the source code to these operating systems, may +be copied and redistributed for use in accordance with this +Agreement, provided that the object code files are not +modified in any way (except for unzipping of compressed +files). + + +2.4. Audio and Video Encoders and Decoders + +You acknowledge and agree that it is your sole responsibility +to obtain any additional third-party licenses required to +make, have made, use, have used, sell, import, and offer for +sale your products or services that include or incorporate any +third-party software and content relating to audio and/or +video encoders and decoders from, including but not limited +to, Microsoft, Thomson, Fraunhofer IIS, Sisvel S.p.A., +MPEG-LA, and Coding Technologies. NVIDIA does not grant to you +under this Agreement any necessary patent or other rights with +respect to any audio and/or video encoders and decoders. + + +2.5. Licensing + +If the distribution terms in this Agreement are not suitable +for your organization, or for any questions regarding this +Agreement, please contact NVIDIA at +nvidia-compute-license-questions@nvidia.com. + + +2.6. Attachment A + +The following portions of the SDK are distributable under the +Agreement: + +Component + +CUDA Runtime + +Windows + +cudart.dll, cudart_static.lib, cudadevrt.lib + +Mac OSX + +libcudart.dylib, libcudart_static.a, libcudadevrt.a + +Linux + +libcudart.so, libcudart_static.a, libcudadevrt.a + +Android + +libcudart.so, libcudart_static.a, libcudadevrt.a + +Component + +CUDA FFT Library + +Windows + +cufft.dll, cufftw.dll, cufft.lib, cufftw.lib + +Mac OSX + +libcufft.dylib, libcufft_static.a, libcufftw.dylib, +libcufftw_static.a + +Linux + +libcufft.so, libcufft_static.a, libcufftw.so, +libcufftw_static.a + +Android + +libcufft.so, libcufft_static.a, libcufftw.so, +libcufftw_static.a + +Component + +CUDA BLAS Library + +Windows + +cublas.dll, cublasLt.dll + +Mac OSX + +libcublas.dylib, libcublasLt.dylib, libcublas_static.a, +libcublasLt_static.a + +Linux + +libcublas.so, libcublasLt.so, libcublas_static.a, +libcublasLt_static.a + +Android + +libcublas.so, libcublasLt.so, libcublas_static.a, +libcublasLt_static.a + +Component + +NVIDIA "Drop-in" BLAS Library + +Windows + +nvblas.dll + +Mac OSX + +libnvblas.dylib + +Linux + +libnvblas.so + +Component + +CUDA Sparse Matrix Library + +Windows + +cusparse.dll, cusparse.lib + +Mac OSX + +libcusparse.dylib, libcusparse_static.a + +Linux + +libcusparse.so, libcusparse_static.a + +Android + +libcusparse.so, libcusparse_static.a + +Component + +CUDA Linear Solver Library + +Windows + +cusolver.dll, cusolver.lib + +Mac OSX + +libcusolver.dylib, libcusolver_static.a + +Linux + +libcusolver.so, libcusolver_static.a + +Android + +libcusolver.so, libcusolver_static.a + +Component + +CUDA Random Number Generation Library + +Windows + +curand.dll, curand.lib + +Mac OSX + +libcurand.dylib, libcurand_static.a + +Linux + +libcurand.so, libcurand_static.a + +Android + +libcurand.so, libcurand_static.a + +Component + +CUDA Accelerated Graph Library + +Component + +NVIDIA Performance Primitives Library + +Windows + +nppc.dll, nppc.lib, nppial.dll, nppial.lib, nppicc.dll, +nppicc.lib, nppicom.dll, nppicom.lib, nppidei.dll, +nppidei.lib, nppif.dll, nppif.lib, nppig.dll, nppig.lib, +nppim.dll, nppim.lib, nppist.dll, nppist.lib, nppisu.dll, +nppisu.lib, nppitc.dll, nppitc.lib, npps.dll, npps.lib + +Mac OSX + +libnppc.dylib, libnppc_static.a, libnppial.dylib, +libnppial_static.a, libnppicc.dylib, libnppicc_static.a, +libnppicom.dylib, libnppicom_static.a, libnppidei.dylib, +libnppidei_static.a, libnppif.dylib, libnppif_static.a, +libnppig.dylib, libnppig_static.a, libnppim.dylib, +libnppisu_static.a, libnppitc.dylib, libnppitc_static.a, +libnpps.dylib, libnpps_static.a + +Linux + +libnppc.so, libnppc_static.a, libnppial.so, +libnppial_static.a, libnppicc.so, libnppicc_static.a, +libnppicom.so, libnppicom_static.a, libnppidei.so, +libnppidei_static.a, libnppif.so, libnppif_static.a +libnppig.so, libnppig_static.a, libnppim.so, +libnppim_static.a, libnppist.so, libnppist_static.a, +libnppisu.so, libnppisu_static.a, libnppitc.so +libnppitc_static.a, libnpps.so, libnpps_static.a + +Android + +libnppc.so, libnppc_static.a, libnppial.so, +libnppial_static.a, libnppicc.so, libnppicc_static.a, +libnppicom.so, libnppicom_static.a, libnppidei.so, +libnppidei_static.a, libnppif.so, libnppif_static.a +libnppig.so, libnppig_static.a, libnppim.so, +libnppim_static.a, libnppist.so, libnppist_static.a, +libnppisu.so, libnppisu_static.a, libnppitc.so +libnppitc_static.a, libnpps.so, libnpps_static.a + +Component + +NVIDIA JPEG Library + +Linux + +libnvjpeg.so, libnvjpeg_static.a + +Component + +Internal common library required for statically linking to +cuBLAS, cuSPARSE, cuFFT, cuRAND, nvJPEG and NPP + +Mac OSX + +libculibos.a + +Linux + +libculibos.a + +Component + +NVIDIA Runtime Compilation Library and Header + +All + +nvrtc.h + +Windows + +nvrtc.dll, nvrtc-builtins.dll + +Mac OSX + +libnvrtc.dylib, libnvrtc-builtins.dylib + +Linux + +libnvrtc.so, libnvrtc-builtins.so + +Component + +NVIDIA Optimizing Compiler Library + +Windows + +nvvm.dll + +Mac OSX + +libnvvm.dylib + +Linux + +libnvvm.so + +Component + +NVIDIA Common Device Math Functions Library + +Windows + +libdevice.10.bc + +Mac OSX + +libdevice.10.bc + +Linux + +libdevice.10.bc + +Component + +CUDA Occupancy Calculation Header Library + +All + +cuda_occupancy.h + +Component + +CUDA Half Precision Headers + +All + +cuda_fp16.h, cuda_fp16.hpp + +Component + +CUDA Profiling Tools Interface (CUPTI) Library + +Windows + +cupti.dll + +Mac OSX + +libcupti.dylib + +Linux + +libcupti.so + +Component + +NVIDIA Tools Extension Library + +Windows + +nvToolsExt.dll, nvToolsExt.lib + +Mac OSX + +libnvToolsExt.dylib + +Linux + +libnvToolsExt.so + +Component + +NVIDIA CUDA Driver Libraries + +Linux + +libcuda.so, libnvidia-fatbinaryloader.so, +libnvidia-ptxjitcompiler.so + +The NVIDIA CUDA Driver Libraries are only distributable in +applications that meet this criteria: + + 1. The application was developed starting from a NVIDIA CUDA + container obtained from Docker Hub or the NVIDIA GPU + Cloud, and + + 2. The resulting application is packaged as a Docker + container and distributed to users on Docker Hub or the + NVIDIA GPU Cloud only. + + +2.7. Attachment B + + +Additional Licensing Obligations + +The following third party components included in the SOFTWARE +are licensed to Licensee pursuant to the following terms and +conditions: + + 1. Licensee's use of the GDB third party component is + subject to the terms and conditions of GNU GPL v3: + + This product includes copyrighted third-party software licensed + under the terms of the GNU General Public License v3 ("GPL v3"). + All third-party software packages are copyright by their respective + authors. GPL v3 terms and conditions are hereby incorporated into + the Agreement by this reference: http://www.gnu.org/licenses/gpl.txt + + Consistent with these licensing requirements, the software + listed below is provided under the terms of the specified + open source software licenses. To obtain source code for + software provided under licenses that require + redistribution of source code, including the GNU General + Public License (GPL) and GNU Lesser General Public License + (LGPL), contact oss-requests@nvidia.com. This offer is + valid for a period of three (3) years from the date of the + distribution of this product by NVIDIA CORPORATION. + + Component License + CUDA-GDB GPL v3 + + 2. Licensee represents and warrants that any and all third + party licensing and/or royalty payment obligations in + connection with Licensee's use of the H.264 video codecs + are solely the responsibility of Licensee. + + 3. Licensee's use of the Thrust library is subject to the + terms and conditions of the Apache License Version 2.0. + All third-party software packages are copyright by their + respective authors. Apache License Version 2.0 terms and + conditions are hereby incorporated into the Agreement by + this reference. + http://www.apache.org/licenses/LICENSE-2.0.html + + In addition, Licensee acknowledges the following notice: + Thrust includes source code from the Boost Iterator, + Tuple, System, and Random Number libraries. + + Boost Software License - Version 1.0 - August 17th, 2003 + . . . . + + Permission is hereby granted, free of charge, to any person or + organization obtaining a copy of the software and accompanying + documentation covered by this license (the "Software") to use, + reproduce, display, distribute, execute, and transmit the Software, + and to prepare derivative works of the Software, and to permit + third-parties to whom the Software is furnished to do so, all + subject to the following: + + The copyright notices in the Software and this entire statement, + including the above license grant, this restriction and the following + disclaimer, must be included in all copies of the Software, in whole + or in part, and all derivative works of the Software, unless such + copies or derivative works are solely in the form of machine-executable + object code generated by a source language processor. + + 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, TITLE AND + NON-INFRINGEMENT. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR + ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE FOR ANY DAMAGES OR + OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, ARISING + FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR + OTHER DEALINGS IN THE SOFTWARE. + + 4. Licensee's use of the LLVM third party component is + subject to the following terms and conditions: + + ====================================================== + LLVM Release License + ====================================================== + University of Illinois/NCSA + Open Source License + + Copyright (c) 2003-2010 University of Illinois at Urbana-Champaign. + All rights reserved. + + Developed by: + + LLVM Team + + University of Illinois at Urbana-Champaign + + http://llvm.org + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to + deal with 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: + + * Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimers. + + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimers in the + documentation and/or other materials provided with the distribution. + + * Neither the names of the LLVM Team, University of Illinois at Urbana- + Champaign, nor the names of its contributors may be used to endorse or + promote products derived from this Software without specific prior + written permission. + + 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 CONTRIBUTORS 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 WITH THE SOFTWARE. + + 5. Licensee's use (e.g. nvprof) of the PCRE third party + component is subject to the following terms and + conditions: + + ------------ + PCRE LICENCE + ------------ + PCRE is a library of functions to support regular expressions whose syntax + and semantics are as close as possible to those of the Perl 5 language. + Release 8 of PCRE is distributed under the terms of the "BSD" licence, as + specified below. The documentation for PCRE, supplied in the "doc" + directory, is distributed under the same terms as the software itself. The + basic library functions are written in C and are freestanding. Also + included in the distribution is a set of C++ wrapper functions, and a just- + in-time compiler that can be used to optimize pattern matching. These are + both optional features that can be omitted when the library is built. + + THE BASIC LIBRARY FUNCTIONS + --------------------------- + Written by: Philip Hazel + Email local part: ph10 + Email domain: cam.ac.uk + University of Cambridge Computing Service, + Cambridge, England. + Copyright (c) 1997-2012 University of Cambridge + All rights reserved. + + PCRE JUST-IN-TIME COMPILATION SUPPORT + ------------------------------------- + Written by: Zoltan Herczeg + Email local part: hzmester + Emain domain: freemail.hu + Copyright(c) 2010-2012 Zoltan Herczeg + All rights reserved. + + STACK-LESS JUST-IN-TIME COMPILER + -------------------------------- + Written by: Zoltan Herczeg + Email local part: hzmester + Emain domain: freemail.hu + Copyright(c) 2009-2012 Zoltan Herczeg + All rights reserved. + + THE C++ WRAPPER FUNCTIONS + ------------------------- + Contributed by: Google Inc. + Copyright (c) 2007-2012, Google Inc. + All rights reserved. + + THE "BSD" LICENCE + ----------------- + 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 University of Cambridge nor the name of Google + Inc. nor the names of their 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. + + 6. Some of the cuBLAS library routines were written by or + derived from code written by Vasily Volkov and are subject + to the Modified Berkeley Software Distribution License as + follows: + + Copyright (c) 2007-2009, Regents of the University of California + + 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 University of California, Berkeley 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 AUTHOR "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 AUTHOR 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. + + 7. Some of the cuBLAS library routines were written by or + derived from code written by Davide Barbieri and are + subject to the Modified Berkeley Software Distribution + License as follows: + + Copyright (c) 2008-2009 Davide Barbieri @ University of Rome Tor Vergata. + + 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. + * The name of the author may not be used to endorse or promote + products derived from this software without specific prior + written permission. + + THIS SOFTWARE IS PROVIDED BY THE AUTHOR "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 AUTHOR 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. + + 8. Some of the cuBLAS library routines were derived from + code developed by the University of Tennessee and are + subject to the Modified Berkeley Software Distribution + License as follows: + + Copyright (c) 2010 The University of Tennessee. + + 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 listed in this license in the documentation and/or + other materials provided with the distribution. + * Neither the name of the copyright holders 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. + + 9. Some of the cuBLAS library routines were written by or + derived from code written by Jonathan Hogg and are subject + to the Modified Berkeley Software Distribution License as + follows: + + Copyright (c) 2012, The Science and Technology Facilities Council (STFC). + + 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 STFC 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 STFC 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. + + 10. Some of the cuBLAS library routines were written by or + derived from code written by Ahmad M. Abdelfattah, David + Keyes, and Hatem Ltaief, and are subject to the Apache + License, Version 2.0, as follows: + + -- (C) Copyright 2013 King Abdullah University of Science and Technology + Authors: + Ahmad Abdelfattah (ahmad.ahmad@kaust.edu.sa) + David Keyes (david.keyes@kaust.edu.sa) + Hatem Ltaief (hatem.ltaief@kaust.edu.sa) + + 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 King Abdullah University of Science and + Technology 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 + HOLDERS 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 + + 11. Some of the cuSPARSE library routines were written by or + derived from code written by Li-Wen Chang and are subject + to the NCSA Open Source License as follows: + + Copyright (c) 2012, University of Illinois. + + All rights reserved. + + Developed by: IMPACT Group, University of Illinois, http://impact.crhc.illinois.edu + + Permission is hereby granted, free of charge, to any person obtaining + a copy of this software and associated documentation files (the + "Software"), to deal with 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: + * 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 + disclaimers in the documentation and/or other materials provided + with the distribution. + * Neither the names of IMPACT Group, University of Illinois, nor + the names of its contributors may be used to endorse or promote + products derived from this Software without specific prior + written permission. + + 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 CONTRIBUTORS 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 WITH THE + SOFTWARE. + + 12. Some of the cuRAND library routines were written by or + derived from code written by Mutsuo Saito and Makoto + Matsumoto and are subject to the following license: + + Copyright (c) 2009, 2010 Mutsuo Saito, Makoto Matsumoto and Hiroshima + University. All rights reserved. + + Copyright (c) 2011 Mutsuo Saito, Makoto Matsumoto, Hiroshima + University and University of Tokyo. 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 Hiroshima University 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. + + 13. Some of the cuRAND library routines were derived from + code developed by D. E. Shaw Research and are subject to + the following license: + + Copyright 2010-2011, D. E. Shaw Research. + + 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 D. E. Shaw Research 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. + + 14. Some of the Math library routines were written by or + derived from code developed by Norbert Juffa and are + subject to the following license: + + Copyright (c) 2015-2017, Norbert Juffa + All rights reserved. + + 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. + + 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. + + 15. Licensee's use of the lz4 third party component is + subject to the following terms and conditions: + + Copyright (C) 2011-2013, Yann Collet. + BSD 2-Clause License (http://www.opensource.org/licenses/bsd-license.php) + + 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. + + 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. + + 16. The NPP library uses code from the Boost Math Toolkit, + and is subject to the following license: + + Boost Software License - Version 1.0 - August 17th, 2003 + . . . . + + Permission is hereby granted, free of charge, to any person or + organization obtaining a copy of the software and accompanying + documentation covered by this license (the "Software") to use, + reproduce, display, distribute, execute, and transmit the Software, + and to prepare derivative works of the Software, and to permit + third-parties to whom the Software is furnished to do so, all + subject to the following: + + The copyright notices in the Software and this entire statement, + including the above license grant, this restriction and the following + disclaimer, must be included in all copies of the Software, in whole + or in part, and all derivative works of the Software, unless such + copies or derivative works are solely in the form of machine-executable + object code generated by a source language processor. + + 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, TITLE AND + NON-INFRINGEMENT. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR + ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE FOR ANY DAMAGES OR + OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, ARISING + FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR + OTHER DEALINGS IN THE SOFTWARE. + + 17. Portions of the Nsight Eclipse Edition is subject to the + following license: + + The Eclipse Foundation makes available all content in this plug-in + ("Content"). Unless otherwise indicated below, the Content is provided + to you under the terms and conditions of the Eclipse Public License + Version 1.0 ("EPL"). A copy of the EPL is available at http:// + www.eclipse.org/legal/epl-v10.html. For purposes of the EPL, "Program" + will mean the Content. + + If you did not receive this Content directly from the Eclipse + Foundation, the Content is being redistributed by another party + ("Redistributor") and different terms and conditions may apply to your + use of any object code in the Content. Check the Redistributor's + license that was provided with the Content. If no such license exists, + contact the Redistributor. Unless otherwise indicated below, the terms + and conditions of the EPL still apply to any source code in the + Content and such source code may be obtained at http://www.eclipse.org. + + 18. Some of the cuBLAS library routines uses code from + OpenAI, which is subject to the following license: + + License URL + https://github.com/openai/openai-gemm/blob/master/LICENSE + + License Text + The MIT License + + Copyright (c) 2016 OpenAI (http://openai.com), 2016 Google Inc. + + 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. + + 19. Licensee's use of the Visual Studio Setup Configuration + Samples is subject to the following license: + + The MIT License (MIT) + Copyright (C) Microsoft Corporation. All rights reserved. + + 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. + + 20. Licensee's use of linmath.h header for CPU functions for + GL vector/matrix operations from lunarG is subject to the + Apache License Version 2.0. + + 21. The DX12-CUDA sample uses the d3dx12.h header, which is + subject to the MIT license . + +----------------- diff --git a/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/METADATA b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..0b8cf29e6f82235cc8fa638b4594844111370546 --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/METADATA @@ -0,0 +1,35 @@ +Metadata-Version: 2.1 +Name: nvidia-cublas-cu12 +Version: 12.4.5.8 +Summary: CUBLAS native runtime libraries +Home-page: https://developer.nvidia.com/cuda-zone +Author: Nvidia CUDA Installer Team +Author-email: cuda_installer@nvidia.com +License: NVIDIA Proprietary Software +Keywords: cuda,nvidia,runtime,machine learning,deep learning +Classifier: Development Status :: 4 - Beta +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Education +Classifier: Intended Audience :: Science/Research +Classifier: License :: Other/Proprietary License +Classifier: Natural Language :: English +Classifier: Programming Language :: Python :: 3 +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 :: Only +Classifier: Topic :: Scientific/Engineering +Classifier: Topic :: Scientific/Engineering :: Mathematics +Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence +Classifier: Topic :: Software Development +Classifier: Topic :: Software Development :: Libraries +Classifier: Operating System :: Microsoft :: Windows +Classifier: Operating System :: POSIX :: Linux +Requires-Python: >=3 +License-File: License.txt + +CUBLAS native runtime libraries diff --git a/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/RECORD b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..b7926eab61036989e0642119b74d5fda406ce404 --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/RECORD @@ -0,0 +1,24 @@ +nvidia/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia/__pycache__/__init__.cpython-312.pyc,, +nvidia/cublas/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia/cublas/__pycache__/__init__.cpython-312.pyc,, +nvidia/cublas/include/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia/cublas/include/__pycache__/__init__.cpython-312.pyc,, +nvidia/cublas/include/cublas.h,sha256=a0lLqy-k47NuwyDjuueC3W0Mpc908MTU7o5sMJqE-1w,41246 +nvidia/cublas/include/cublasLt.h,sha256=51KyHQc7T9rxmVfNimP9O6vka8JqBdebjZKCWKZakt4,77626 +nvidia/cublas/include/cublasXt.h,sha256=CW9dyXYGSUW1wEXrVVyhU6OxBK1PUvMoYdVGlQT7L9A,37380 +nvidia/cublas/include/cublas_api.h,sha256=XRArlgDy_4hWuEt8XafRsE9KRJ5XVo06Nh113cgg-7o,370663 +nvidia/cublas/include/cublas_v2.h,sha256=qxMdB5jb97luEfw61LEAB-Wlr8A9DLBvO4rRypDCNKw,15460 +nvidia/cublas/include/nvblas.h,sha256=dXCLR-2oUiJFzLsDtIAK09m42ct4G0HWdYzBUuDPXpc,23341 +nvidia/cublas/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia/cublas/lib/__pycache__/__init__.cpython-312.pyc,, +nvidia/cublas/lib/libcublas.so.12,sha256=TMRVJkSaldOYU4l4XphbeDJGDqPbF4vJ6ZY74hEfvug,109604768 +nvidia/cublas/lib/libcublasLt.so.12,sha256=RKgTqi2giDD5CD-B0Otz8a5AUqTZsLDeSAqPbNnrMHg,441938896 +nvidia/cublas/lib/libnvblas.so.12,sha256=fCpY3FQVQgg5IwHQ_j1ToSDkwevquegM6R_plIuurck,757496 +nvidia_cublas_cu12-12.4.5.8.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +nvidia_cublas_cu12-12.4.5.8.dist-info/License.txt,sha256=rW9YU_ugyg0VnQ9Y1JrkmDDC-Mk_epJki5zpCttMbM0,59262 +nvidia_cublas_cu12-12.4.5.8.dist-info/METADATA,sha256=FtdQvmVmrqzO9Vp7VbNtbQWUxXF45arMsnGnwYdlZuc,1505 +nvidia_cublas_cu12-12.4.5.8.dist-info/RECORD,, +nvidia_cublas_cu12-12.4.5.8.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia_cublas_cu12-12.4.5.8.dist-info/WHEEL,sha256=XDTs3wIbcE-BcRO08VJlZpA6z9OaC1mOKPCGGGwuM2g,109 +nvidia_cublas_cu12-12.4.5.8.dist-info/top_level.txt,sha256=fTkAtiFuL16nUrB9ytDDtpytz2t0B4NvYTnRzwAhO14,7 diff --git a/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/REQUESTED b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/WHEEL b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..e6c30e957cfb045017a9fef3430bb8ee87c4a074 --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: bdist_wheel (0.42.0) +Root-Is-Purelib: true +Tag: py3-none-manylinux2014_x86_64 + diff --git a/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/top_level.txt b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..862f7abf232cdfbb928609856247292e81c9decb --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cublas_cu12-12.4.5.8.dist-info/top_level.txt @@ -0,0 +1 @@ +nvidia diff --git a/lib/python3.12/site-packages/pydantic/__init__.py b/lib/python3.12/site-packages/pydantic/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..716ca403ae0894726d21408ca83913219fc92500 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/__init__.py @@ -0,0 +1,445 @@ +import typing +from importlib import import_module +from warnings import warn + +from ._migration import getattr_migration +from .version import VERSION + +if typing.TYPE_CHECKING: + # import of virtually everything is supported via `__getattr__` below, + # but we need them here for type checking and IDE support + import pydantic_core + from pydantic_core.core_schema import ( + FieldSerializationInfo, + SerializationInfo, + SerializerFunctionWrapHandler, + ValidationInfo, + ValidatorFunctionWrapHandler, + ) + + from . import dataclasses + from .aliases import AliasChoices, AliasGenerator, AliasPath + from .annotated_handlers import GetCoreSchemaHandler, GetJsonSchemaHandler + from .config import ConfigDict, with_config + from .errors import * + from .fields import Field, PrivateAttr, computed_field + from .functional_serializers import ( + PlainSerializer, + SerializeAsAny, + WrapSerializer, + field_serializer, + model_serializer, + ) + from .functional_validators import ( + AfterValidator, + BeforeValidator, + InstanceOf, + ModelWrapValidatorHandler, + PlainValidator, + SkipValidation, + WrapValidator, + field_validator, + model_validator, + ) + from .json_schema import WithJsonSchema + from .main import * + from .networks import * + from .type_adapter import TypeAdapter + from .types import * + from .validate_call_decorator import validate_call + from .warnings import ( + PydanticDeprecatedSince20, + PydanticDeprecatedSince26, + PydanticDeprecatedSince29, + PydanticDeprecatedSince210, + PydanticDeprecatedSince211, + PydanticDeprecationWarning, + PydanticExperimentalWarning, + ) + + # this encourages pycharm to import `ValidationError` from here, not pydantic_core + ValidationError = pydantic_core.ValidationError + from .deprecated.class_validators import root_validator, validator + from .deprecated.config import BaseConfig, Extra + from .deprecated.tools import * + from .root_model import RootModel + +__version__ = VERSION +__all__ = ( + # dataclasses + 'dataclasses', + # functional validators + 'field_validator', + 'model_validator', + 'AfterValidator', + 'BeforeValidator', + 'PlainValidator', + 'WrapValidator', + 'SkipValidation', + 'InstanceOf', + 'ModelWrapValidatorHandler', + # JSON Schema + 'WithJsonSchema', + # deprecated V1 functional validators, these are imported via `__getattr__` below + 'root_validator', + 'validator', + # functional serializers + 'field_serializer', + 'model_serializer', + 'PlainSerializer', + 'SerializeAsAny', + 'WrapSerializer', + # config + 'ConfigDict', + 'with_config', + # deprecated V1 config, these are imported via `__getattr__` below + 'BaseConfig', + 'Extra', + # validate_call + 'validate_call', + # errors + 'PydanticErrorCodes', + 'PydanticUserError', + 'PydanticSchemaGenerationError', + 'PydanticImportError', + 'PydanticUndefinedAnnotation', + 'PydanticInvalidForJsonSchema', + 'PydanticForbiddenQualifier', + # fields + 'Field', + 'computed_field', + 'PrivateAttr', + # alias + 'AliasChoices', + 'AliasGenerator', + 'AliasPath', + # main + 'BaseModel', + 'create_model', + # network + 'AnyUrl', + 'AnyHttpUrl', + 'FileUrl', + 'HttpUrl', + 'FtpUrl', + 'WebsocketUrl', + 'AnyWebsocketUrl', + 'UrlConstraints', + 'EmailStr', + 'NameEmail', + 'IPvAnyAddress', + 'IPvAnyInterface', + 'IPvAnyNetwork', + 'PostgresDsn', + 'CockroachDsn', + 'AmqpDsn', + 'RedisDsn', + 'MongoDsn', + 'KafkaDsn', + 'NatsDsn', + 'MySQLDsn', + 'MariaDBDsn', + 'ClickHouseDsn', + 'SnowflakeDsn', + 'validate_email', + # root_model + 'RootModel', + # deprecated tools, these are imported via `__getattr__` below + 'parse_obj_as', + 'schema_of', + 'schema_json_of', + # types + 'Strict', + 'StrictStr', + 'conbytes', + 'conlist', + 'conset', + 'confrozenset', + 'constr', + 'StringConstraints', + 'ImportString', + 'conint', + 'PositiveInt', + 'NegativeInt', + 'NonNegativeInt', + 'NonPositiveInt', + 'confloat', + 'PositiveFloat', + 'NegativeFloat', + 'NonNegativeFloat', + 'NonPositiveFloat', + 'FiniteFloat', + 'condecimal', + 'condate', + 'UUID1', + 'UUID3', + 'UUID4', + 'UUID5', + 'UUID6', + 'UUID7', + 'UUID8', + 'FilePath', + 'DirectoryPath', + 'NewPath', + 'Json', + 'Secret', + 'SecretStr', + 'SecretBytes', + 'SocketPath', + 'StrictBool', + 'StrictBytes', + 'StrictInt', + 'StrictFloat', + 'PaymentCardNumber', + 'ByteSize', + 'PastDate', + 'FutureDate', + 'PastDatetime', + 'FutureDatetime', + 'AwareDatetime', + 'NaiveDatetime', + 'AllowInfNan', + 'EncoderProtocol', + 'EncodedBytes', + 'EncodedStr', + 'Base64Encoder', + 'Base64Bytes', + 'Base64Str', + 'Base64UrlBytes', + 'Base64UrlStr', + 'GetPydanticSchema', + 'Tag', + 'Discriminator', + 'JsonValue', + 'FailFast', + # type_adapter + 'TypeAdapter', + # version + '__version__', + 'VERSION', + # warnings + 'PydanticDeprecatedSince20', + 'PydanticDeprecatedSince26', + 'PydanticDeprecatedSince29', + 'PydanticDeprecatedSince210', + 'PydanticDeprecatedSince211', + 'PydanticDeprecationWarning', + 'PydanticExperimentalWarning', + # annotated handlers + 'GetCoreSchemaHandler', + 'GetJsonSchemaHandler', + # pydantic_core + 'ValidationError', + 'ValidationInfo', + 'SerializationInfo', + 'ValidatorFunctionWrapHandler', + 'FieldSerializationInfo', + 'SerializerFunctionWrapHandler', + 'OnErrorOmit', +) + +# A mapping of {: (package, )} defining dynamic imports +_dynamic_imports: 'dict[str, tuple[str, str]]' = { + 'dataclasses': (__spec__.parent, '__module__'), + # functional validators + 'field_validator': (__spec__.parent, '.functional_validators'), + 'model_validator': (__spec__.parent, '.functional_validators'), + 'AfterValidator': (__spec__.parent, '.functional_validators'), + 'BeforeValidator': (__spec__.parent, '.functional_validators'), + 'PlainValidator': (__spec__.parent, '.functional_validators'), + 'WrapValidator': (__spec__.parent, '.functional_validators'), + 'SkipValidation': (__spec__.parent, '.functional_validators'), + 'InstanceOf': (__spec__.parent, '.functional_validators'), + 'ModelWrapValidatorHandler': (__spec__.parent, '.functional_validators'), + # JSON Schema + 'WithJsonSchema': (__spec__.parent, '.json_schema'), + # functional serializers + 'field_serializer': (__spec__.parent, '.functional_serializers'), + 'model_serializer': (__spec__.parent, '.functional_serializers'), + 'PlainSerializer': (__spec__.parent, '.functional_serializers'), + 'SerializeAsAny': (__spec__.parent, '.functional_serializers'), + 'WrapSerializer': (__spec__.parent, '.functional_serializers'), + # config + 'ConfigDict': (__spec__.parent, '.config'), + 'with_config': (__spec__.parent, '.config'), + # validate call + 'validate_call': (__spec__.parent, '.validate_call_decorator'), + # errors + 'PydanticErrorCodes': (__spec__.parent, '.errors'), + 'PydanticUserError': (__spec__.parent, '.errors'), + 'PydanticSchemaGenerationError': (__spec__.parent, '.errors'), + 'PydanticImportError': (__spec__.parent, '.errors'), + 'PydanticUndefinedAnnotation': (__spec__.parent, '.errors'), + 'PydanticInvalidForJsonSchema': (__spec__.parent, '.errors'), + 'PydanticForbiddenQualifier': (__spec__.parent, '.errors'), + # fields + 'Field': (__spec__.parent, '.fields'), + 'computed_field': (__spec__.parent, '.fields'), + 'PrivateAttr': (__spec__.parent, '.fields'), + # alias + 'AliasChoices': (__spec__.parent, '.aliases'), + 'AliasGenerator': (__spec__.parent, '.aliases'), + 'AliasPath': (__spec__.parent, '.aliases'), + # main + 'BaseModel': (__spec__.parent, '.main'), + 'create_model': (__spec__.parent, '.main'), + # network + 'AnyUrl': (__spec__.parent, '.networks'), + 'AnyHttpUrl': (__spec__.parent, '.networks'), + 'FileUrl': (__spec__.parent, '.networks'), + 'HttpUrl': (__spec__.parent, '.networks'), + 'FtpUrl': (__spec__.parent, '.networks'), + 'WebsocketUrl': (__spec__.parent, '.networks'), + 'AnyWebsocketUrl': (__spec__.parent, '.networks'), + 'UrlConstraints': (__spec__.parent, '.networks'), + 'EmailStr': (__spec__.parent, '.networks'), + 'NameEmail': (__spec__.parent, '.networks'), + 'IPvAnyAddress': (__spec__.parent, '.networks'), + 'IPvAnyInterface': (__spec__.parent, '.networks'), + 'IPvAnyNetwork': (__spec__.parent, '.networks'), + 'PostgresDsn': (__spec__.parent, '.networks'), + 'CockroachDsn': (__spec__.parent, '.networks'), + 'AmqpDsn': (__spec__.parent, '.networks'), + 'RedisDsn': (__spec__.parent, '.networks'), + 'MongoDsn': (__spec__.parent, '.networks'), + 'KafkaDsn': (__spec__.parent, '.networks'), + 'NatsDsn': (__spec__.parent, '.networks'), + 'MySQLDsn': (__spec__.parent, '.networks'), + 'MariaDBDsn': (__spec__.parent, '.networks'), + 'ClickHouseDsn': (__spec__.parent, '.networks'), + 'SnowflakeDsn': (__spec__.parent, '.networks'), + 'validate_email': (__spec__.parent, '.networks'), + # root_model + 'RootModel': (__spec__.parent, '.root_model'), + # types + 'Strict': (__spec__.parent, '.types'), + 'StrictStr': (__spec__.parent, '.types'), + 'conbytes': (__spec__.parent, '.types'), + 'conlist': (__spec__.parent, '.types'), + 'conset': (__spec__.parent, '.types'), + 'confrozenset': (__spec__.parent, '.types'), + 'constr': (__spec__.parent, '.types'), + 'StringConstraints': (__spec__.parent, '.types'), + 'ImportString': (__spec__.parent, '.types'), + 'conint': (__spec__.parent, '.types'), + 'PositiveInt': (__spec__.parent, '.types'), + 'NegativeInt': (__spec__.parent, '.types'), + 'NonNegativeInt': (__spec__.parent, '.types'), + 'NonPositiveInt': (__spec__.parent, '.types'), + 'confloat': (__spec__.parent, '.types'), + 'PositiveFloat': (__spec__.parent, '.types'), + 'NegativeFloat': (__spec__.parent, '.types'), + 'NonNegativeFloat': (__spec__.parent, '.types'), + 'NonPositiveFloat': (__spec__.parent, '.types'), + 'FiniteFloat': (__spec__.parent, '.types'), + 'condecimal': (__spec__.parent, '.types'), + 'condate': (__spec__.parent, '.types'), + 'UUID1': (__spec__.parent, '.types'), + 'UUID3': (__spec__.parent, '.types'), + 'UUID4': (__spec__.parent, '.types'), + 'UUID5': (__spec__.parent, '.types'), + 'UUID6': (__spec__.parent, '.types'), + 'UUID7': (__spec__.parent, '.types'), + 'UUID8': (__spec__.parent, '.types'), + 'FilePath': (__spec__.parent, '.types'), + 'DirectoryPath': (__spec__.parent, '.types'), + 'NewPath': (__spec__.parent, '.types'), + 'Json': (__spec__.parent, '.types'), + 'Secret': (__spec__.parent, '.types'), + 'SecretStr': (__spec__.parent, '.types'), + 'SecretBytes': (__spec__.parent, '.types'), + 'StrictBool': (__spec__.parent, '.types'), + 'StrictBytes': (__spec__.parent, '.types'), + 'StrictInt': (__spec__.parent, '.types'), + 'StrictFloat': (__spec__.parent, '.types'), + 'PaymentCardNumber': (__spec__.parent, '.types'), + 'ByteSize': (__spec__.parent, '.types'), + 'PastDate': (__spec__.parent, '.types'), + 'SocketPath': (__spec__.parent, '.types'), + 'FutureDate': (__spec__.parent, '.types'), + 'PastDatetime': (__spec__.parent, '.types'), + 'FutureDatetime': (__spec__.parent, '.types'), + 'AwareDatetime': (__spec__.parent, '.types'), + 'NaiveDatetime': (__spec__.parent, '.types'), + 'AllowInfNan': (__spec__.parent, '.types'), + 'EncoderProtocol': (__spec__.parent, '.types'), + 'EncodedBytes': (__spec__.parent, '.types'), + 'EncodedStr': (__spec__.parent, '.types'), + 'Base64Encoder': (__spec__.parent, '.types'), + 'Base64Bytes': (__spec__.parent, '.types'), + 'Base64Str': (__spec__.parent, '.types'), + 'Base64UrlBytes': (__spec__.parent, '.types'), + 'Base64UrlStr': (__spec__.parent, '.types'), + 'GetPydanticSchema': (__spec__.parent, '.types'), + 'Tag': (__spec__.parent, '.types'), + 'Discriminator': (__spec__.parent, '.types'), + 'JsonValue': (__spec__.parent, '.types'), + 'OnErrorOmit': (__spec__.parent, '.types'), + 'FailFast': (__spec__.parent, '.types'), + # type_adapter + 'TypeAdapter': (__spec__.parent, '.type_adapter'), + # warnings + 'PydanticDeprecatedSince20': (__spec__.parent, '.warnings'), + 'PydanticDeprecatedSince26': (__spec__.parent, '.warnings'), + 'PydanticDeprecatedSince29': (__spec__.parent, '.warnings'), + 'PydanticDeprecatedSince210': (__spec__.parent, '.warnings'), + 'PydanticDeprecatedSince211': (__spec__.parent, '.warnings'), + 'PydanticDeprecationWarning': (__spec__.parent, '.warnings'), + 'PydanticExperimentalWarning': (__spec__.parent, '.warnings'), + # annotated handlers + 'GetCoreSchemaHandler': (__spec__.parent, '.annotated_handlers'), + 'GetJsonSchemaHandler': (__spec__.parent, '.annotated_handlers'), + # pydantic_core stuff + 'ValidationError': ('pydantic_core', '.'), + 'ValidationInfo': ('pydantic_core', '.core_schema'), + 'SerializationInfo': ('pydantic_core', '.core_schema'), + 'ValidatorFunctionWrapHandler': ('pydantic_core', '.core_schema'), + 'FieldSerializationInfo': ('pydantic_core', '.core_schema'), + 'SerializerFunctionWrapHandler': ('pydantic_core', '.core_schema'), + # deprecated, mostly not included in __all__ + 'root_validator': (__spec__.parent, '.deprecated.class_validators'), + 'validator': (__spec__.parent, '.deprecated.class_validators'), + 'BaseConfig': (__spec__.parent, '.deprecated.config'), + 'Extra': (__spec__.parent, '.deprecated.config'), + 'parse_obj_as': (__spec__.parent, '.deprecated.tools'), + 'schema_of': (__spec__.parent, '.deprecated.tools'), + 'schema_json_of': (__spec__.parent, '.deprecated.tools'), + # deprecated dynamic imports + 'FieldValidationInfo': ('pydantic_core', '.core_schema'), + 'GenerateSchema': (__spec__.parent, '._internal._generate_schema'), +} +_deprecated_dynamic_imports = {'FieldValidationInfo', 'GenerateSchema'} + +_getattr_migration = getattr_migration(__name__) + + +def __getattr__(attr_name: str) -> object: + if attr_name in _deprecated_dynamic_imports: + warn( + f'Importing {attr_name} from `pydantic` is deprecated. This feature is either no longer supported, or is not public.', + DeprecationWarning, + stacklevel=2, + ) + + dynamic_attr = _dynamic_imports.get(attr_name) + if dynamic_attr is None: + return _getattr_migration(attr_name) + + package, module_name = dynamic_attr + + if module_name == '__module__': + result = import_module(f'.{attr_name}', package=package) + globals()[attr_name] = result + return result + else: + module = import_module(module_name, package=package) + result = getattr(module, attr_name) + g = globals() + for k, (_, v_module_name) in _dynamic_imports.items(): + if v_module_name == module_name and k not in _deprecated_dynamic_imports: + g[k] = getattr(module, k) + return result + + +def __dir__() -> 'list[str]': + return list(__all__) diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..88727fb3b7266d427de4e9f6577b6c7c2a2eb135 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/_migration.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/_migration.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84a556c8edb20037ad77dd18200ed8eae1fdf06d Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/_migration.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/alias_generators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/alias_generators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f564b7f74277aafcf2cae1251f6236fb0bc00bca Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/alias_generators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/aliases.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/aliases.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3356555d6af1a16152ec5e44acf1c60fcd88b9e7 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/aliases.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/annotated_handlers.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/annotated_handlers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0f0737cd2db2758d8b615090deca08baba4290d Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/annotated_handlers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/class_validators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/class_validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d5ec9316b5c94991057eb88d3490f77b6a0c6f5 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/class_validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/color.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/color.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..881a780d0e52aabb5c3c011627dbe119b4dbcc23 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/color.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8825d9d7d26378253744b46616fc5a13bd549b7 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/dataclasses.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/dataclasses.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b6e628c6ad49ae44316367d1da253ed4db6ee16e Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/dataclasses.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/datetime_parse.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/datetime_parse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d6b6fa1b659c35f2100e37bd1daf3286983b05c3 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/datetime_parse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/decorator.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/decorator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83f45c4751ed4b7d3d49569b4125cb922e46af4b Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/decorator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/env_settings.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/env_settings.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..740989ff60dbdbf555d9b75a8657f8a674c7ac74 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/env_settings.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/error_wrappers.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/error_wrappers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08720e9d559571df0491597d7e97e4ba1db98810 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/error_wrappers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/errors.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/errors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a480f513ad11c3ec8535da6859aea7737fe86157 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/errors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/fields.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/fields.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51c3f3cb28e89845880ed17a780718738dccb34d Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/fields.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/functional_serializers.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/functional_serializers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19af54b7a29c47d5411a26f5ad257ec6b611882c Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/functional_serializers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/functional_validators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/functional_validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dbd81d31a042d2e5bd69670a5f004cd9145dc563 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/functional_validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/generics.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/generics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42421d6226be631edf68fea1bca73b814665f420 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/generics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/json.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/json.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51e9480202cfec413aba0686507d7a5c13060b5d Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/json.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/main.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/main.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..034804f1b63031e6279669ab755abefca1c69b57 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/main.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/mypy.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/mypy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49592ae7865ba9871bef532af5d91a7402da8905 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/mypy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/networks.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/networks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..19f9340e6682c17a94f807b5b8b82a9777ec087e Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/networks.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/parse.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/parse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..064ae94c900f022f0107c5d116df91c376a84198 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/parse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/root_model.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/root_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea123f6515980113edcaf13a56bc7476d0427360 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/root_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/schema.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/schema.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..32e0e31df3ec0f0dd3eb229026485052c38c38c9 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/schema.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/tools.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/tools.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b190a35bc48343f05ceaaf63b35110d2504c64db Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/tools.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/type_adapter.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/type_adapter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e7e8e113ec2a84e33588edb7c94ae6043e98d9d Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/type_adapter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c899bf73d19ef0e9a97558c20727170004897fba Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/typing.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/typing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a8c71978579527015caadc62867dce76868f113 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/typing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b979384b683bee7acd38a3192b1ac5adf17a54b Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/validate_call_decorator.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/validate_call_decorator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..272e4b117a9f32ee550fdcaaf94aa53d243afa5b Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/validate_call_decorator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/validators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d83b396ca9019e8d91ffffa480697cac5c901c0 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/version.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55a2fcb0ae23120fd3227f5a228aa433b85d1482 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/version.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/__pycache__/warnings.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/__pycache__/warnings.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7c71a0dac4f6f62113a9b48e636b967c5b391fd4 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/__pycache__/warnings.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__init__.py b/lib/python3.12/site-packages/pydantic/_internal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07c0fc8b4d09e68fa1c712d519e25d731afdef70 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_config.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..446f2570221e75601eca68b0cd8c4345bc321880 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_core_metadata.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_core_metadata.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9a2b1e268b4d22a8675596cc42cbff6b289bcb74 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_core_metadata.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_core_utils.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_core_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ade43427d8a376105c0c21a28a686320f5edd199 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_core_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_dataclasses.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_dataclasses.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92d3557677c83ee2037937735e423d40db573e45 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_dataclasses.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_decorators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_decorators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..03cb80e9557544877c7e56aee5125c65d66d9204 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_decorators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_decorators_v1.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_decorators_v1.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4c6cef2c8dd8807dd718dc6791e2acfed1abda3 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_decorators_v1.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_discriminated_union.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_discriminated_union.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0da86e41d0d38577269428a4af61327a2404eb90 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_discriminated_union.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_docs_extraction.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_docs_extraction.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..95c1cb62981c43385363223873953dff0a73fe31 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_docs_extraction.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_fields.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_fields.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8d323d5482114d076a52da6656f082a2514faec Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_fields.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_forward_ref.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_forward_ref.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a762b9d80d246c1f6af79a5b27957b7f0144237 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_forward_ref.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_generics.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_generics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c84bedf661147f4b7a85f7513a91449df31bfb79 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_generics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_git.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_git.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff420353c5744010300b0ef31b62421fc4e850c4 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_git.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_import_utils.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_import_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..606c3a55a8ff004f03093a2b98953bf8693a2bfa Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_import_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_internal_dataclass.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_internal_dataclass.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..329d59e5175e1174e32045cb1acad8b8fdbabf12 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_internal_dataclass.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_known_annotated_metadata.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_known_annotated_metadata.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7da288374c227407c80d35735a8622ddfdb257af Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_known_annotated_metadata.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_mock_val_ser.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_mock_val_ser.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d356cd428b6ff92e8e3aa8c8b6ff119937b8ba5d Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_mock_val_ser.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_model_construction.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_model_construction.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e00a94ae74aba4b02a58961b6d2be5b451269db Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_model_construction.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_namespace_utils.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_namespace_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c32e09d3bc7b8b8f55cd73cc23e301d06685a02f Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_namespace_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_repr.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_repr.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9740d42aac1b2ee53947af5174177df3ddbe83ae Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_repr.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_schema_gather.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_schema_gather.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0db8fdeeb6c7ad7f637f800f12b74a02138d2641 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_schema_gather.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_schema_generation_shared.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_schema_generation_shared.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dfd158ea6825b8fddd5933ae5c8902785962600c Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_schema_generation_shared.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_serializers.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_serializers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5734c27e581014ab43010b4e39e124c427a066f9 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_serializers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_signature.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_signature.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8bf7f50d08a1bb6efadbf5a1902afec81b4bcebc Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_signature.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_typing_extra.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_typing_extra.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ec85cf6c033caf1d5d082ca255ad8d420c1ed797 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_typing_extra.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f86b78a4accbecfe449a0a5d801ee26fb6b5db3b Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_validate_call.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_validate_call.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b514abb19eb1810da7a6e28e337f50deb47ce04f Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_validate_call.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_validators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5324f94c363060d7a84fe2463d48e7d28d2817ff Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/_internal/__pycache__/_validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/_internal/_config.py b/lib/python3.12/site-packages/pydantic/_internal/_config.py new file mode 100644 index 0000000000000000000000000000000000000000..fe7126463379f6ceec7b860496e16fbfddbab6c9 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_config.py @@ -0,0 +1,373 @@ +from __future__ import annotations as _annotations + +import warnings +from contextlib import contextmanager +from re import Pattern +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Literal, + cast, +) + +from pydantic_core import core_schema +from typing_extensions import Self + +from ..aliases import AliasGenerator +from ..config import ConfigDict, ExtraValues, JsonDict, JsonEncoder, JsonSchemaExtraCallable +from ..errors import PydanticUserError +from ..warnings import PydanticDeprecatedSince20, PydanticDeprecatedSince210 + +if not TYPE_CHECKING: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + +if TYPE_CHECKING: + from .._internal._schema_generation_shared import GenerateSchema + from ..fields import ComputedFieldInfo, FieldInfo + +DEPRECATION_MESSAGE = 'Support for class-based `config` is deprecated, use ConfigDict instead.' + + +class ConfigWrapper: + """Internal wrapper for Config which exposes ConfigDict items as attributes.""" + + __slots__ = ('config_dict',) + + config_dict: ConfigDict + + # all annotations are copied directly from ConfigDict, and should be kept up to date, a test will fail if they + # stop matching + title: str | None + str_to_lower: bool + str_to_upper: bool + str_strip_whitespace: bool + str_min_length: int + str_max_length: int | None + extra: ExtraValues | None + frozen: bool + populate_by_name: bool + use_enum_values: bool + validate_assignment: bool + arbitrary_types_allowed: bool + from_attributes: bool + # whether to use the actual key provided in the data (e.g. alias or first alias for "field required" errors) instead of field_names + # to construct error `loc`s, default `True` + loc_by_alias: bool + alias_generator: Callable[[str], str] | AliasGenerator | None + model_title_generator: Callable[[type], str] | None + field_title_generator: Callable[[str, FieldInfo | ComputedFieldInfo], str] | None + ignored_types: tuple[type, ...] + allow_inf_nan: bool + json_schema_extra: JsonDict | JsonSchemaExtraCallable | None + json_encoders: dict[type[object], JsonEncoder] | None + + # new in V2 + strict: bool + # whether instances of models and dataclasses (including subclass instances) should re-validate, default 'never' + revalidate_instances: Literal['always', 'never', 'subclass-instances'] + ser_json_timedelta: Literal['iso8601', 'float'] + ser_json_bytes: Literal['utf8', 'base64', 'hex'] + val_json_bytes: Literal['utf8', 'base64', 'hex'] + ser_json_inf_nan: Literal['null', 'constants', 'strings'] + # whether to validate default values during validation, default False + validate_default: bool + validate_return: bool + protected_namespaces: tuple[str | Pattern[str], ...] + hide_input_in_errors: bool + defer_build: bool + plugin_settings: dict[str, object] | None + schema_generator: type[GenerateSchema] | None + json_schema_serialization_defaults_required: bool + json_schema_mode_override: Literal['validation', 'serialization', None] + coerce_numbers_to_str: bool + regex_engine: Literal['rust-regex', 'python-re'] + validation_error_cause: bool + use_attribute_docstrings: bool + cache_strings: bool | Literal['all', 'keys', 'none'] + validate_by_alias: bool + validate_by_name: bool + serialize_by_alias: bool + + def __init__(self, config: ConfigDict | dict[str, Any] | type[Any] | None, *, check: bool = True): + if check: + self.config_dict = prepare_config(config) + else: + self.config_dict = cast(ConfigDict, config) + + @classmethod + def for_model(cls, bases: tuple[type[Any], ...], namespace: dict[str, Any], kwargs: dict[str, Any]) -> Self: + """Build a new `ConfigWrapper` instance for a `BaseModel`. + + The config wrapper built based on (in descending order of priority): + - options from `kwargs` + - options from the `namespace` + - options from the base classes (`bases`) + + Args: + bases: A tuple of base classes. + namespace: The namespace of the class being created. + kwargs: The kwargs passed to the class being created. + + Returns: + A `ConfigWrapper` instance for `BaseModel`. + """ + config_new = ConfigDict() + for base in bases: + config = getattr(base, 'model_config', None) + if config: + config_new.update(config.copy()) + + config_class_from_namespace = namespace.get('Config') + config_dict_from_namespace = namespace.get('model_config') + + raw_annotations = namespace.get('__annotations__', {}) + if raw_annotations.get('model_config') and config_dict_from_namespace is None: + raise PydanticUserError( + '`model_config` cannot be used as a model field name. Use `model_config` for model configuration.', + code='model-config-invalid-field-name', + ) + + if config_class_from_namespace and config_dict_from_namespace: + raise PydanticUserError('"Config" and "model_config" cannot be used together', code='config-both') + + config_from_namespace = config_dict_from_namespace or prepare_config(config_class_from_namespace) + + config_new.update(config_from_namespace) + + for k in list(kwargs.keys()): + if k in config_keys: + config_new[k] = kwargs.pop(k) + + return cls(config_new) + + # we don't show `__getattr__` to type checkers so missing attributes cause errors + if not TYPE_CHECKING: # pragma: no branch + + def __getattr__(self, name: str) -> Any: + try: + return self.config_dict[name] + except KeyError: + try: + return config_defaults[name] + except KeyError: + raise AttributeError(f'Config has no attribute {name!r}') from None + + def core_config(self, title: str | None) -> core_schema.CoreConfig: + """Create a pydantic-core config. + + We don't use getattr here since we don't want to populate with defaults. + + Args: + title: The title to use if not set in config. + + Returns: + A `CoreConfig` object created from config. + """ + config = self.config_dict + + if config.get('schema_generator') is not None: + warnings.warn( + 'The `schema_generator` setting has been deprecated since v2.10. This setting no longer has any effect.', + PydanticDeprecatedSince210, + stacklevel=2, + ) + + if (populate_by_name := config.get('populate_by_name')) is not None: + # We include this patch for backwards compatibility purposes, but this config setting will be deprecated in v3.0, and likely removed in v4.0. + # Thus, the above warning and this patch can be removed then as well. + if config.get('validate_by_name') is None: + config['validate_by_alias'] = True + config['validate_by_name'] = populate_by_name + + # We dynamically patch validate_by_name to be True if validate_by_alias is set to False + # and validate_by_name is not explicitly set. + if config.get('validate_by_alias') is False and config.get('validate_by_name') is None: + config['validate_by_name'] = True + + if (not config.get('validate_by_alias', True)) and (not config.get('validate_by_name', False)): + raise PydanticUserError( + 'At least one of `validate_by_alias` or `validate_by_name` must be set to True.', + code='validate-by-alias-and-name-false', + ) + + return core_schema.CoreConfig( + **{ # pyright: ignore[reportArgumentType] + k: v + for k, v in ( + ('title', config.get('title') or title or None), + ('extra_fields_behavior', config.get('extra')), + ('allow_inf_nan', config.get('allow_inf_nan')), + ('str_strip_whitespace', config.get('str_strip_whitespace')), + ('str_to_lower', config.get('str_to_lower')), + ('str_to_upper', config.get('str_to_upper')), + ('strict', config.get('strict')), + ('ser_json_timedelta', config.get('ser_json_timedelta')), + ('ser_json_bytes', config.get('ser_json_bytes')), + ('val_json_bytes', config.get('val_json_bytes')), + ('ser_json_inf_nan', config.get('ser_json_inf_nan')), + ('from_attributes', config.get('from_attributes')), + ('loc_by_alias', config.get('loc_by_alias')), + ('revalidate_instances', config.get('revalidate_instances')), + ('validate_default', config.get('validate_default')), + ('str_max_length', config.get('str_max_length')), + ('str_min_length', config.get('str_min_length')), + ('hide_input_in_errors', config.get('hide_input_in_errors')), + ('coerce_numbers_to_str', config.get('coerce_numbers_to_str')), + ('regex_engine', config.get('regex_engine')), + ('validation_error_cause', config.get('validation_error_cause')), + ('cache_strings', config.get('cache_strings')), + ('validate_by_alias', config.get('validate_by_alias')), + ('validate_by_name', config.get('validate_by_name')), + ('serialize_by_alias', config.get('serialize_by_alias')), + ) + if v is not None + } + ) + + def __repr__(self): + c = ', '.join(f'{k}={v!r}' for k, v in self.config_dict.items()) + return f'ConfigWrapper({c})' + + +class ConfigWrapperStack: + """A stack of `ConfigWrapper` instances.""" + + def __init__(self, config_wrapper: ConfigWrapper): + self._config_wrapper_stack: list[ConfigWrapper] = [config_wrapper] + + @property + def tail(self) -> ConfigWrapper: + return self._config_wrapper_stack[-1] + + @contextmanager + def push(self, config_wrapper: ConfigWrapper | ConfigDict | None): + if config_wrapper is None: + yield + return + + if not isinstance(config_wrapper, ConfigWrapper): + config_wrapper = ConfigWrapper(config_wrapper, check=False) + + self._config_wrapper_stack.append(config_wrapper) + try: + yield + finally: + self._config_wrapper_stack.pop() + + +config_defaults = ConfigDict( + title=None, + str_to_lower=False, + str_to_upper=False, + str_strip_whitespace=False, + str_min_length=0, + str_max_length=None, + # let the model / dataclass decide how to handle it + extra=None, + frozen=False, + populate_by_name=False, + use_enum_values=False, + validate_assignment=False, + arbitrary_types_allowed=False, + from_attributes=False, + loc_by_alias=True, + alias_generator=None, + model_title_generator=None, + field_title_generator=None, + ignored_types=(), + allow_inf_nan=True, + json_schema_extra=None, + strict=False, + revalidate_instances='never', + ser_json_timedelta='iso8601', + ser_json_bytes='utf8', + val_json_bytes='utf8', + ser_json_inf_nan='null', + validate_default=False, + validate_return=False, + protected_namespaces=('model_validate', 'model_dump'), + hide_input_in_errors=False, + json_encoders=None, + defer_build=False, + schema_generator=None, + plugin_settings=None, + json_schema_serialization_defaults_required=False, + json_schema_mode_override=None, + coerce_numbers_to_str=False, + regex_engine='rust-regex', + validation_error_cause=False, + use_attribute_docstrings=False, + cache_strings=True, + validate_by_alias=True, + validate_by_name=False, + serialize_by_alias=False, +) + + +def prepare_config(config: ConfigDict | dict[str, Any] | type[Any] | None) -> ConfigDict: + """Create a `ConfigDict` instance from an existing dict, a class (e.g. old class-based config) or None. + + Args: + config: The input config. + + Returns: + A ConfigDict object created from config. + """ + if config is None: + return ConfigDict() + + if not isinstance(config, dict): + warnings.warn(DEPRECATION_MESSAGE, DeprecationWarning) + config = {k: getattr(config, k) for k in dir(config) if not k.startswith('__')} + + config_dict = cast(ConfigDict, config) + check_deprecated(config_dict) + return config_dict + + +config_keys = set(ConfigDict.__annotations__.keys()) + + +V2_REMOVED_KEYS = { + 'allow_mutation', + 'error_msg_templates', + 'fields', + 'getter_dict', + 'smart_union', + 'underscore_attrs_are_private', + 'json_loads', + 'json_dumps', + 'copy_on_model_validation', + 'post_init_call', +} +V2_RENAMED_KEYS = { + 'allow_population_by_field_name': 'validate_by_name', + 'anystr_lower': 'str_to_lower', + 'anystr_strip_whitespace': 'str_strip_whitespace', + 'anystr_upper': 'str_to_upper', + 'keep_untouched': 'ignored_types', + 'max_anystr_length': 'str_max_length', + 'min_anystr_length': 'str_min_length', + 'orm_mode': 'from_attributes', + 'schema_extra': 'json_schema_extra', + 'validate_all': 'validate_default', +} + + +def check_deprecated(config_dict: ConfigDict) -> None: + """Check for deprecated config keys and warn the user. + + Args: + config_dict: The input config. + """ + deprecated_removed_keys = V2_REMOVED_KEYS & config_dict.keys() + deprecated_renamed_keys = V2_RENAMED_KEYS.keys() & config_dict.keys() + if deprecated_removed_keys or deprecated_renamed_keys: + renamings = {k: V2_RENAMED_KEYS[k] for k in sorted(deprecated_renamed_keys)} + renamed_bullets = [f'* {k!r} has been renamed to {v!r}' for k, v in renamings.items()] + removed_bullets = [f'* {k!r} has been removed' for k in sorted(deprecated_removed_keys)] + message = '\n'.join(['Valid config keys have changed in V2:'] + renamed_bullets + removed_bullets) + warnings.warn(message, UserWarning) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_core_metadata.py b/lib/python3.12/site-packages/pydantic/_internal/_core_metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..9f2510c03429b6de24063bec5937e4334077722d --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_core_metadata.py @@ -0,0 +1,97 @@ +from __future__ import annotations as _annotations + +from typing import TYPE_CHECKING, Any, TypedDict, cast +from warnings import warn + +if TYPE_CHECKING: + from ..config import JsonDict, JsonSchemaExtraCallable + from ._schema_generation_shared import ( + GetJsonSchemaFunction, + ) + + +class CoreMetadata(TypedDict, total=False): + """A `TypedDict` for holding the metadata dict of the schema. + + Attributes: + pydantic_js_functions: List of JSON schema functions that resolve refs during application. + pydantic_js_annotation_functions: List of JSON schema functions that don't resolve refs during application. + pydantic_js_prefer_positional_arguments: Whether JSON schema generator will + prefer positional over keyword arguments for an 'arguments' schema. + custom validation function. Only applies to before, plain, and wrap validators. + pydantic_js_updates: key / value pair updates to apply to the JSON schema for a type. + pydantic_js_extra: WIP, either key/value pair updates to apply to the JSON schema, or a custom callable. + pydantic_internal_union_tag_key: Used internally by the `Tag` metadata to specify the tag used for a discriminated union. + pydantic_internal_union_discriminator: Used internally to specify the discriminator value for a discriminated union + when the discriminator was applied to a `'definition-ref'` schema, and that reference was missing at the time + of the annotation application. + + TODO: Perhaps we should move this structure to pydantic-core. At the moment, though, + it's easier to iterate on if we leave it in pydantic until we feel there is a semi-stable API. + + TODO: It's unfortunate how functionally oriented JSON schema generation is, especially that which occurs during + the core schema generation process. It's inevitable that we need to store some json schema related information + on core schemas, given that we generate JSON schemas directly from core schemas. That being said, debugging related + issues is quite difficult when JSON schema information is disguised via dynamically defined functions. + """ + + pydantic_js_functions: list[GetJsonSchemaFunction] + pydantic_js_annotation_functions: list[GetJsonSchemaFunction] + pydantic_js_prefer_positional_arguments: bool + pydantic_js_updates: JsonDict + pydantic_js_extra: JsonDict | JsonSchemaExtraCallable + pydantic_internal_union_tag_key: str + pydantic_internal_union_discriminator: str + + +def update_core_metadata( + core_metadata: Any, + /, + *, + pydantic_js_functions: list[GetJsonSchemaFunction] | None = None, + pydantic_js_annotation_functions: list[GetJsonSchemaFunction] | None = None, + pydantic_js_updates: JsonDict | None = None, + pydantic_js_extra: JsonDict | JsonSchemaExtraCallable | None = None, +) -> None: + from ..json_schema import PydanticJsonSchemaWarning + + """Update CoreMetadata instance in place. When we make modifications in this function, they + take effect on the `core_metadata` reference passed in as the first (and only) positional argument. + + First, cast to `CoreMetadata`, then finish with a cast to `dict[str, Any]` for core schema compatibility. + We do this here, instead of before / after each call to this function so that this typing hack + can be easily removed if/when we move `CoreMetadata` to `pydantic-core`. + + For parameter descriptions, see `CoreMetadata` above. + """ + core_metadata = cast(CoreMetadata, core_metadata) + + if pydantic_js_functions: + core_metadata.setdefault('pydantic_js_functions', []).extend(pydantic_js_functions) + + if pydantic_js_annotation_functions: + core_metadata.setdefault('pydantic_js_annotation_functions', []).extend(pydantic_js_annotation_functions) + + if pydantic_js_updates: + if (existing_updates := core_metadata.get('pydantic_js_updates')) is not None: + core_metadata['pydantic_js_updates'] = {**existing_updates, **pydantic_js_updates} + else: + core_metadata['pydantic_js_updates'] = pydantic_js_updates + + if pydantic_js_extra is not None: + existing_pydantic_js_extra = core_metadata.get('pydantic_js_extra') + if existing_pydantic_js_extra is None: + core_metadata['pydantic_js_extra'] = pydantic_js_extra + if isinstance(existing_pydantic_js_extra, dict): + if isinstance(pydantic_js_extra, dict): + core_metadata['pydantic_js_extra'] = {**existing_pydantic_js_extra, **pydantic_js_extra} + if callable(pydantic_js_extra): + warn( + 'Composing `dict` and `callable` type `json_schema_extra` is not supported.' + 'The `callable` type is being ignored.' + "If you'd like support for this behavior, please open an issue on pydantic.", + PydanticJsonSchemaWarning, + ) + if callable(existing_pydantic_js_extra): + # if ever there's a case of a callable, we'll just keep the last json schema extra spec + core_metadata['pydantic_js_extra'] = pydantic_js_extra diff --git a/lib/python3.12/site-packages/pydantic/_internal/_core_utils.py b/lib/python3.12/site-packages/pydantic/_internal/_core_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..cf8cf7c5f078536490a3d802cf421daae962971b --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_core_utils.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import inspect +import os +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Union + +from pydantic_core import CoreSchema, core_schema +from pydantic_core import validate_core_schema as _validate_core_schema +from typing_extensions import TypeGuard, get_args, get_origin +from typing_inspection import typing_objects + +from . import _repr +from ._typing_extra import is_generic_alias + +if TYPE_CHECKING: + from rich.console import Console + +AnyFunctionSchema = Union[ + core_schema.AfterValidatorFunctionSchema, + core_schema.BeforeValidatorFunctionSchema, + core_schema.WrapValidatorFunctionSchema, + core_schema.PlainValidatorFunctionSchema, +] + + +FunctionSchemaWithInnerSchema = Union[ + core_schema.AfterValidatorFunctionSchema, + core_schema.BeforeValidatorFunctionSchema, + core_schema.WrapValidatorFunctionSchema, +] + +CoreSchemaField = Union[ + core_schema.ModelField, core_schema.DataclassField, core_schema.TypedDictField, core_schema.ComputedField +] +CoreSchemaOrField = Union[core_schema.CoreSchema, CoreSchemaField] + +_CORE_SCHEMA_FIELD_TYPES = {'typed-dict-field', 'dataclass-field', 'model-field', 'computed-field'} +_FUNCTION_WITH_INNER_SCHEMA_TYPES = {'function-before', 'function-after', 'function-wrap'} +_LIST_LIKE_SCHEMA_WITH_ITEMS_TYPES = {'list', 'set', 'frozenset'} + + +def is_core_schema( + schema: CoreSchemaOrField, +) -> TypeGuard[CoreSchema]: + return schema['type'] not in _CORE_SCHEMA_FIELD_TYPES + + +def is_core_schema_field( + schema: CoreSchemaOrField, +) -> TypeGuard[CoreSchemaField]: + return schema['type'] in _CORE_SCHEMA_FIELD_TYPES + + +def is_function_with_inner_schema( + schema: CoreSchemaOrField, +) -> TypeGuard[FunctionSchemaWithInnerSchema]: + return schema['type'] in _FUNCTION_WITH_INNER_SCHEMA_TYPES + + +def is_list_like_schema_with_items_schema( + schema: CoreSchema, +) -> TypeGuard[core_schema.ListSchema | core_schema.SetSchema | core_schema.FrozenSetSchema]: + return schema['type'] in _LIST_LIKE_SCHEMA_WITH_ITEMS_TYPES + + +def get_type_ref(type_: Any, args_override: tuple[type[Any], ...] | None = None) -> str: + """Produces the ref to be used for this type by pydantic_core's core schemas. + + This `args_override` argument was added for the purpose of creating valid recursive references + when creating generic models without needing to create a concrete class. + """ + origin = get_origin(type_) or type_ + + args = get_args(type_) if is_generic_alias(type_) else (args_override or ()) + generic_metadata = getattr(type_, '__pydantic_generic_metadata__', None) + if generic_metadata: + origin = generic_metadata['origin'] or origin + args = generic_metadata['args'] or args + + module_name = getattr(origin, '__module__', '') + if typing_objects.is_typealiastype(origin): + type_ref = f'{module_name}.{origin.__name__}:{id(origin)}' + else: + try: + qualname = getattr(origin, '__qualname__', f'') + except Exception: + qualname = getattr(origin, '__qualname__', '') + type_ref = f'{module_name}.{qualname}:{id(origin)}' + + arg_refs: list[str] = [] + for arg in args: + if isinstance(arg, str): + # Handle string literals as a special case; we may be able to remove this special handling if we + # wrap them in a ForwardRef at some point. + arg_ref = f'{arg}:str-{id(arg)}' + else: + arg_ref = f'{_repr.display_as_type(arg)}:{id(arg)}' + arg_refs.append(arg_ref) + if arg_refs: + type_ref = f'{type_ref}[{",".join(arg_refs)}]' + return type_ref + + +def get_ref(s: core_schema.CoreSchema) -> None | str: + """Get the ref from the schema if it has one. + This exists just for type checking to work correctly. + """ + return s.get('ref', None) + + +def validate_core_schema(schema: CoreSchema) -> CoreSchema: + if os.getenv('PYDANTIC_VALIDATE_CORE_SCHEMAS'): + return _validate_core_schema(schema) + return schema + + +def _clean_schema_for_pretty_print(obj: Any, strip_metadata: bool = True) -> Any: # pragma: no cover + """A utility function to remove irrelevant information from a core schema.""" + if isinstance(obj, Mapping): + new_dct = {} + for k, v in obj.items(): + if k == 'metadata' and strip_metadata: + new_metadata = {} + + for meta_k, meta_v in v.items(): + if meta_k in ('pydantic_js_functions', 'pydantic_js_annotation_functions'): + new_metadata['js_metadata'] = '' + else: + new_metadata[meta_k] = _clean_schema_for_pretty_print(meta_v, strip_metadata=strip_metadata) + + if list(new_metadata.keys()) == ['js_metadata']: + new_metadata = {''} + + new_dct[k] = new_metadata + # Remove some defaults: + elif k in ('custom_init', 'root_model') and not v: + continue + else: + new_dct[k] = _clean_schema_for_pretty_print(v, strip_metadata=strip_metadata) + + return new_dct + elif isinstance(obj, Sequence) and not isinstance(obj, str): + return [_clean_schema_for_pretty_print(v, strip_metadata=strip_metadata) for v in obj] + else: + return obj + + +def pretty_print_core_schema( + val: Any, + *, + console: Console | None = None, + max_depth: int | None = None, + strip_metadata: bool = True, +) -> None: # pragma: no cover + """Pretty-print a core schema using the `rich` library. + + Args: + val: The core schema to print, or a Pydantic model/dataclass/type adapter + (in which case the cached core schema is fetched and printed). + console: A rich console to use when printing. Defaults to the global rich console instance. + max_depth: The number of nesting levels which may be printed. + strip_metadata: Whether to strip metadata in the output. If `True` any known core metadata + attributes will be stripped (but custom attributes are kept). Defaults to `True`. + """ + # lazy import: + from rich.pretty import pprint + + # circ. imports: + from pydantic import BaseModel, TypeAdapter + from pydantic.dataclasses import is_pydantic_dataclass + + if (inspect.isclass(val) and issubclass(val, BaseModel)) or is_pydantic_dataclass(val): + val = val.__pydantic_core_schema__ + if isinstance(val, TypeAdapter): + val = val.core_schema + cleaned_schema = _clean_schema_for_pretty_print(val, strip_metadata=strip_metadata) + + pprint(cleaned_schema, console=console, max_depth=max_depth) + + +pps = pretty_print_core_schema diff --git a/lib/python3.12/site-packages/pydantic/_internal/_dataclasses.py b/lib/python3.12/site-packages/pydantic/_internal/_dataclasses.py new file mode 100644 index 0000000000000000000000000000000000000000..03f156fc79f47f9416e02653bd44008b53f1cc54 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_dataclasses.py @@ -0,0 +1,238 @@ +"""Private logic for creating pydantic dataclasses.""" + +from __future__ import annotations as _annotations + +import dataclasses +import typing +import warnings +from functools import partial, wraps +from typing import Any, ClassVar + +from pydantic_core import ( + ArgsKwargs, + SchemaSerializer, + SchemaValidator, + core_schema, +) +from typing_extensions import TypeGuard + +from ..errors import PydanticUndefinedAnnotation +from ..plugin._schema_validator import PluggableSchemaValidator, create_schema_validator +from ..warnings import PydanticDeprecatedSince20 +from . import _config, _decorators +from ._fields import collect_dataclass_fields +from ._generate_schema import GenerateSchema, InvalidSchemaError +from ._generics import get_standard_typevars_map +from ._mock_val_ser import set_dataclass_mocks +from ._namespace_utils import NsResolver +from ._signature import generate_pydantic_signature +from ._utils import LazyClassAttribute + +if typing.TYPE_CHECKING: + from _typeshed import DataclassInstance as StandardDataclass + + from ..config import ConfigDict + from ..fields import FieldInfo + + class PydanticDataclass(StandardDataclass, typing.Protocol): + """A protocol containing attributes only available once a class has been decorated as a Pydantic dataclass. + + Attributes: + __pydantic_config__: Pydantic-specific configuration settings for the dataclass. + __pydantic_complete__: Whether dataclass building is completed, or if there are still undefined fields. + __pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer. + __pydantic_decorators__: Metadata containing the decorators defined on the dataclass. + __pydantic_fields__: Metadata about the fields defined on the dataclass. + __pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the dataclass. + __pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the dataclass. + """ + + __pydantic_config__: ClassVar[ConfigDict] + __pydantic_complete__: ClassVar[bool] + __pydantic_core_schema__: ClassVar[core_schema.CoreSchema] + __pydantic_decorators__: ClassVar[_decorators.DecoratorInfos] + __pydantic_fields__: ClassVar[dict[str, FieldInfo]] + __pydantic_serializer__: ClassVar[SchemaSerializer] + __pydantic_validator__: ClassVar[SchemaValidator | PluggableSchemaValidator] + + @classmethod + def __pydantic_fields_complete__(cls) -> bool: ... + +else: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + + +def set_dataclass_fields( + cls: type[StandardDataclass], + ns_resolver: NsResolver | None = None, + config_wrapper: _config.ConfigWrapper | None = None, +) -> None: + """Collect and set `cls.__pydantic_fields__`. + + Args: + cls: The class. + ns_resolver: Namespace resolver to use when getting dataclass annotations. + config_wrapper: The config wrapper instance, defaults to `None`. + """ + typevars_map = get_standard_typevars_map(cls) + fields = collect_dataclass_fields( + cls, ns_resolver=ns_resolver, typevars_map=typevars_map, config_wrapper=config_wrapper + ) + + cls.__pydantic_fields__ = fields # type: ignore + + +def complete_dataclass( + cls: type[Any], + config_wrapper: _config.ConfigWrapper, + *, + raise_errors: bool = True, + ns_resolver: NsResolver | None = None, + _force_build: bool = False, +) -> bool: + """Finish building a pydantic dataclass. + + This logic is called on a class which has already been wrapped in `dataclasses.dataclass()`. + + This is somewhat analogous to `pydantic._internal._model_construction.complete_model_class`. + + Args: + cls: The class. + config_wrapper: The config wrapper instance. + raise_errors: Whether to raise errors, defaults to `True`. + ns_resolver: The namespace resolver instance to use when collecting dataclass fields + and during schema building. + _force_build: Whether to force building the dataclass, no matter if + [`defer_build`][pydantic.config.ConfigDict.defer_build] is set. + + Returns: + `True` if building a pydantic dataclass is successfully completed, `False` otherwise. + + Raises: + PydanticUndefinedAnnotation: If `raise_error` is `True` and there is an undefined annotations. + """ + original_init = cls.__init__ + + # dataclass.__init__ must be defined here so its `__qualname__` can be changed since functions can't be copied, + # and so that the mock validator is used if building was deferred: + def __init__(__dataclass_self__: PydanticDataclass, *args: Any, **kwargs: Any) -> None: + __tracebackhide__ = True + s = __dataclass_self__ + s.__pydantic_validator__.validate_python(ArgsKwargs(args, kwargs), self_instance=s) + + __init__.__qualname__ = f'{cls.__qualname__}.__init__' + + cls.__init__ = __init__ # type: ignore + cls.__pydantic_config__ = config_wrapper.config_dict # type: ignore + + set_dataclass_fields(cls, ns_resolver, config_wrapper=config_wrapper) + + if not _force_build and config_wrapper.defer_build: + set_dataclass_mocks(cls) + return False + + if hasattr(cls, '__post_init_post_parse__'): + warnings.warn( + 'Support for `__post_init_post_parse__` has been dropped, the method will not be called', DeprecationWarning + ) + + typevars_map = get_standard_typevars_map(cls) + gen_schema = GenerateSchema( + config_wrapper, + ns_resolver=ns_resolver, + typevars_map=typevars_map, + ) + + # set __signature__ attr only for the class, but not for its instances + # (because instances can define `__call__`, and `inspect.signature` shouldn't + # use the `__signature__` attribute and instead generate from `__call__`). + cls.__signature__ = LazyClassAttribute( + '__signature__', + partial( + generate_pydantic_signature, + # It's important that we reference the `original_init` here, + # as it is the one synthesized by the stdlib `dataclass` module: + init=original_init, + fields=cls.__pydantic_fields__, # type: ignore + validate_by_name=config_wrapper.validate_by_name, + extra=config_wrapper.extra, + is_dataclass=True, + ), + ) + + try: + schema = gen_schema.generate_schema(cls) + except PydanticUndefinedAnnotation as e: + if raise_errors: + raise + set_dataclass_mocks(cls, f'`{e.name}`') + return False + + core_config = config_wrapper.core_config(title=cls.__name__) + + try: + schema = gen_schema.clean_schema(schema) + except InvalidSchemaError: + set_dataclass_mocks(cls) + return False + + # We are about to set all the remaining required properties expected for this cast; + # __pydantic_decorators__ and __pydantic_fields__ should already be set + cls = typing.cast('type[PydanticDataclass]', cls) + # debug(schema) + + cls.__pydantic_core_schema__ = schema + cls.__pydantic_validator__ = validator = create_schema_validator( + schema, cls, cls.__module__, cls.__qualname__, 'dataclass', core_config, config_wrapper.plugin_settings + ) + cls.__pydantic_serializer__ = SchemaSerializer(schema, core_config) + + if config_wrapper.validate_assignment: + + @wraps(cls.__setattr__) + def validated_setattr(instance: Any, field: str, value: str, /) -> None: + validator.validate_assignment(instance, field, value) + + cls.__setattr__ = validated_setattr.__get__(None, cls) # type: ignore + + cls.__pydantic_complete__ = True + return True + + +def is_builtin_dataclass(_cls: type[Any]) -> TypeGuard[type[StandardDataclass]]: + """Returns True if a class is a stdlib dataclass and *not* a pydantic dataclass. + + We check that + - `_cls` is a dataclass + - `_cls` does not inherit from a processed pydantic dataclass (and thus have a `__pydantic_validator__`) + - `_cls` does not have any annotations that are not dataclass fields + e.g. + ```python + import dataclasses + + import pydantic.dataclasses + + @dataclasses.dataclass + class A: + x: int + + @pydantic.dataclasses.dataclass + class B(A): + y: int + ``` + In this case, when we first check `B`, we make an extra check and look at the annotations ('y'), + which won't be a superset of all the dataclass fields (only the stdlib fields i.e. 'x') + + Args: + cls: The class. + + Returns: + `True` if the class is a stdlib dataclass, `False` otherwise. + """ + return ( + dataclasses.is_dataclass(_cls) + and not hasattr(_cls, '__pydantic_validator__') + and set(_cls.__dataclass_fields__).issuperset(set(getattr(_cls, '__annotations__', {}))) + ) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_decorators.py b/lib/python3.12/site-packages/pydantic/_internal/_decorators.py new file mode 100644 index 0000000000000000000000000000000000000000..92880a49fcced8ae1fb7c8eff6bc9e7e5977cb17 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_decorators.py @@ -0,0 +1,838 @@ +"""Logic related to validators applied to models etc. via the `@field_validator` and `@model_validator` decorators.""" + +from __future__ import annotations as _annotations + +import types +from collections import deque +from collections.abc import Iterable +from dataclasses import dataclass, field +from functools import cached_property, partial, partialmethod +from inspect import Parameter, Signature, isdatadescriptor, ismethoddescriptor, signature +from itertools import islice +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, Literal, TypeVar, Union + +from pydantic_core import PydanticUndefined, PydanticUndefinedType, core_schema +from typing_extensions import TypeAlias, is_typeddict + +from ..errors import PydanticUserError +from ._core_utils import get_type_ref +from ._internal_dataclass import slots_true +from ._namespace_utils import GlobalsNamespace, MappingNamespace +from ._typing_extra import get_function_type_hints +from ._utils import can_be_positional + +if TYPE_CHECKING: + from ..fields import ComputedFieldInfo + from ..functional_validators import FieldValidatorModes + + +@dataclass(**slots_true) +class ValidatorDecoratorInfo: + """A container for data from `@validator` so that we can access it + while building the pydantic-core schema. + + Attributes: + decorator_repr: A class variable representing the decorator string, '@validator'. + fields: A tuple of field names the validator should be called on. + mode: The proposed validator mode. + each_item: For complex objects (sets, lists etc.) whether to validate individual + elements rather than the whole object. + always: Whether this method and other validators should be called even if the value is missing. + check_fields: Whether to check that the fields actually exist on the model. + """ + + decorator_repr: ClassVar[str] = '@validator' + + fields: tuple[str, ...] + mode: Literal['before', 'after'] + each_item: bool + always: bool + check_fields: bool | None + + +@dataclass(**slots_true) +class FieldValidatorDecoratorInfo: + """A container for data from `@field_validator` so that we can access it + while building the pydantic-core schema. + + Attributes: + decorator_repr: A class variable representing the decorator string, '@field_validator'. + fields: A tuple of field names the validator should be called on. + mode: The proposed validator mode. + check_fields: Whether to check that the fields actually exist on the model. + json_schema_input_type: The input type of the function. This is only used to generate + the appropriate JSON Schema (in validation mode) and can only specified + when `mode` is either `'before'`, `'plain'` or `'wrap'`. + """ + + decorator_repr: ClassVar[str] = '@field_validator' + + fields: tuple[str, ...] + mode: FieldValidatorModes + check_fields: bool | None + json_schema_input_type: Any + + +@dataclass(**slots_true) +class RootValidatorDecoratorInfo: + """A container for data from `@root_validator` so that we can access it + while building the pydantic-core schema. + + Attributes: + decorator_repr: A class variable representing the decorator string, '@root_validator'. + mode: The proposed validator mode. + """ + + decorator_repr: ClassVar[str] = '@root_validator' + mode: Literal['before', 'after'] + + +@dataclass(**slots_true) +class FieldSerializerDecoratorInfo: + """A container for data from `@field_serializer` so that we can access it + while building the pydantic-core schema. + + Attributes: + decorator_repr: A class variable representing the decorator string, '@field_serializer'. + fields: A tuple of field names the serializer should be called on. + mode: The proposed serializer mode. + return_type: The type of the serializer's return value. + when_used: The serialization condition. Accepts a string with values `'always'`, `'unless-none'`, `'json'`, + and `'json-unless-none'`. + check_fields: Whether to check that the fields actually exist on the model. + """ + + decorator_repr: ClassVar[str] = '@field_serializer' + fields: tuple[str, ...] + mode: Literal['plain', 'wrap'] + return_type: Any + when_used: core_schema.WhenUsed + check_fields: bool | None + + +@dataclass(**slots_true) +class ModelSerializerDecoratorInfo: + """A container for data from `@model_serializer` so that we can access it + while building the pydantic-core schema. + + Attributes: + decorator_repr: A class variable representing the decorator string, '@model_serializer'. + mode: The proposed serializer mode. + return_type: The type of the serializer's return value. + when_used: The serialization condition. Accepts a string with values `'always'`, `'unless-none'`, `'json'`, + and `'json-unless-none'`. + """ + + decorator_repr: ClassVar[str] = '@model_serializer' + mode: Literal['plain', 'wrap'] + return_type: Any + when_used: core_schema.WhenUsed + + +@dataclass(**slots_true) +class ModelValidatorDecoratorInfo: + """A container for data from `@model_validator` so that we can access it + while building the pydantic-core schema. + + Attributes: + decorator_repr: A class variable representing the decorator string, '@model_validator'. + mode: The proposed serializer mode. + """ + + decorator_repr: ClassVar[str] = '@model_validator' + mode: Literal['wrap', 'before', 'after'] + + +DecoratorInfo: TypeAlias = """Union[ + ValidatorDecoratorInfo, + FieldValidatorDecoratorInfo, + RootValidatorDecoratorInfo, + FieldSerializerDecoratorInfo, + ModelSerializerDecoratorInfo, + ModelValidatorDecoratorInfo, + ComputedFieldInfo, +]""" + +ReturnType = TypeVar('ReturnType') +DecoratedType: TypeAlias = ( + 'Union[classmethod[Any, Any, ReturnType], staticmethod[Any, ReturnType], Callable[..., ReturnType], property]' +) + + +@dataclass # can't use slots here since we set attributes on `__post_init__` +class PydanticDescriptorProxy(Generic[ReturnType]): + """Wrap a classmethod, staticmethod, property or unbound function + and act as a descriptor that allows us to detect decorated items + from the class' attributes. + + This class' __get__ returns the wrapped item's __get__ result, + which makes it transparent for classmethods and staticmethods. + + Attributes: + wrapped: The decorator that has to be wrapped. + decorator_info: The decorator info. + shim: A wrapper function to wrap V1 style function. + """ + + wrapped: DecoratedType[ReturnType] + decorator_info: DecoratorInfo + shim: Callable[[Callable[..., Any]], Callable[..., Any]] | None = None + + def __post_init__(self): + for attr in 'setter', 'deleter': + if hasattr(self.wrapped, attr): + f = partial(self._call_wrapped_attr, name=attr) + setattr(self, attr, f) + + def _call_wrapped_attr(self, func: Callable[[Any], None], *, name: str) -> PydanticDescriptorProxy[ReturnType]: + self.wrapped = getattr(self.wrapped, name)(func) + if isinstance(self.wrapped, property): + # update ComputedFieldInfo.wrapped_property + from ..fields import ComputedFieldInfo + + if isinstance(self.decorator_info, ComputedFieldInfo): + self.decorator_info.wrapped_property = self.wrapped + return self + + def __get__(self, obj: object | None, obj_type: type[object] | None = None) -> PydanticDescriptorProxy[ReturnType]: + try: + return self.wrapped.__get__(obj, obj_type) + except AttributeError: + # not a descriptor, e.g. a partial object + return self.wrapped # type: ignore[return-value] + + def __set_name__(self, instance: Any, name: str) -> None: + if hasattr(self.wrapped, '__set_name__'): + self.wrapped.__set_name__(instance, name) # pyright: ignore[reportFunctionMemberAccess] + + def __getattr__(self, name: str, /) -> Any: + """Forward checks for __isabstractmethod__ and such.""" + return getattr(self.wrapped, name) + + +DecoratorInfoType = TypeVar('DecoratorInfoType', bound=DecoratorInfo) + + +@dataclass(**slots_true) +class Decorator(Generic[DecoratorInfoType]): + """A generic container class to join together the decorator metadata + (metadata from decorator itself, which we have when the + decorator is called but not when we are building the core-schema) + and the bound function (which we have after the class itself is created). + + Attributes: + cls_ref: The class ref. + cls_var_name: The decorated function name. + func: The decorated function. + shim: A wrapper function to wrap V1 style function. + info: The decorator info. + """ + + cls_ref: str + cls_var_name: str + func: Callable[..., Any] + shim: Callable[[Any], Any] | None + info: DecoratorInfoType + + @staticmethod + def build( + cls_: Any, + *, + cls_var_name: str, + shim: Callable[[Any], Any] | None, + info: DecoratorInfoType, + ) -> Decorator[DecoratorInfoType]: + """Build a new decorator. + + Args: + cls_: The class. + cls_var_name: The decorated function name. + shim: A wrapper function to wrap V1 style function. + info: The decorator info. + + Returns: + The new decorator instance. + """ + func = get_attribute_from_bases(cls_, cls_var_name) + if shim is not None: + func = shim(func) + func = unwrap_wrapped_function(func, unwrap_partial=False) + if not callable(func): + # This branch will get hit for classmethod properties + attribute = get_attribute_from_base_dicts(cls_, cls_var_name) # prevents the binding call to `__get__` + if isinstance(attribute, PydanticDescriptorProxy): + func = unwrap_wrapped_function(attribute.wrapped) + return Decorator( + cls_ref=get_type_ref(cls_), + cls_var_name=cls_var_name, + func=func, + shim=shim, + info=info, + ) + + def bind_to_cls(self, cls: Any) -> Decorator[DecoratorInfoType]: + """Bind the decorator to a class. + + Args: + cls: the class. + + Returns: + The new decorator instance. + """ + return self.build( + cls, + cls_var_name=self.cls_var_name, + shim=self.shim, + info=self.info, + ) + + +def get_bases(tp: type[Any]) -> tuple[type[Any], ...]: + """Get the base classes of a class or typeddict. + + Args: + tp: The type or class to get the bases. + + Returns: + The base classes. + """ + if is_typeddict(tp): + return tp.__orig_bases__ # type: ignore + try: + return tp.__bases__ + except AttributeError: + return () + + +def mro(tp: type[Any]) -> tuple[type[Any], ...]: + """Calculate the Method Resolution Order of bases using the C3 algorithm. + + See https://www.python.org/download/releases/2.3/mro/ + """ + # try to use the existing mro, for performance mainly + # but also because it helps verify the implementation below + if not is_typeddict(tp): + try: + return tp.__mro__ + except AttributeError: + # GenericAlias and some other cases + pass + + bases = get_bases(tp) + return (tp,) + mro_for_bases(bases) + + +def mro_for_bases(bases: tuple[type[Any], ...]) -> tuple[type[Any], ...]: + def merge_seqs(seqs: list[deque[type[Any]]]) -> Iterable[type[Any]]: + while True: + non_empty = [seq for seq in seqs if seq] + if not non_empty: + # Nothing left to process, we're done. + return + candidate: type[Any] | None = None + for seq in non_empty: # Find merge candidates among seq heads. + candidate = seq[0] + not_head = [s for s in non_empty if candidate in islice(s, 1, None)] + if not_head: + # Reject the candidate. + candidate = None + else: + break + if not candidate: + raise TypeError('Inconsistent hierarchy, no C3 MRO is possible') + yield candidate + for seq in non_empty: + # Remove candidate. + if seq[0] == candidate: + seq.popleft() + + seqs = [deque(mro(base)) for base in bases] + [deque(bases)] + return tuple(merge_seqs(seqs)) + + +_sentinel = object() + + +def get_attribute_from_bases(tp: type[Any] | tuple[type[Any], ...], name: str) -> Any: + """Get the attribute from the next class in the MRO that has it, + aiming to simulate calling the method on the actual class. + + The reason for iterating over the mro instead of just getting + the attribute (which would do that for us) is to support TypedDict, + which lacks a real __mro__, but can have a virtual one constructed + from its bases (as done here). + + Args: + tp: The type or class to search for the attribute. If a tuple, this is treated as a set of base classes. + name: The name of the attribute to retrieve. + + Returns: + Any: The attribute value, if found. + + Raises: + AttributeError: If the attribute is not found in any class in the MRO. + """ + if isinstance(tp, tuple): + for base in mro_for_bases(tp): + attribute = base.__dict__.get(name, _sentinel) + if attribute is not _sentinel: + attribute_get = getattr(attribute, '__get__', None) + if attribute_get is not None: + return attribute_get(None, tp) + return attribute + raise AttributeError(f'{name} not found in {tp}') + else: + try: + return getattr(tp, name) + except AttributeError: + return get_attribute_from_bases(mro(tp), name) + + +def get_attribute_from_base_dicts(tp: type[Any], name: str) -> Any: + """Get an attribute out of the `__dict__` following the MRO. + This prevents the call to `__get__` on the descriptor, and allows + us to get the original function for classmethod properties. + + Args: + tp: The type or class to search for the attribute. + name: The name of the attribute to retrieve. + + Returns: + Any: The attribute value, if found. + + Raises: + KeyError: If the attribute is not found in any class's `__dict__` in the MRO. + """ + for base in reversed(mro(tp)): + if name in base.__dict__: + return base.__dict__[name] + return tp.__dict__[name] # raise the error + + +@dataclass(**slots_true) +class DecoratorInfos: + """Mapping of name in the class namespace to decorator info. + + note that the name in the class namespace is the function or attribute name + not the field name! + """ + + validators: dict[str, Decorator[ValidatorDecoratorInfo]] = field(default_factory=dict) + field_validators: dict[str, Decorator[FieldValidatorDecoratorInfo]] = field(default_factory=dict) + root_validators: dict[str, Decorator[RootValidatorDecoratorInfo]] = field(default_factory=dict) + field_serializers: dict[str, Decorator[FieldSerializerDecoratorInfo]] = field(default_factory=dict) + model_serializers: dict[str, Decorator[ModelSerializerDecoratorInfo]] = field(default_factory=dict) + model_validators: dict[str, Decorator[ModelValidatorDecoratorInfo]] = field(default_factory=dict) + computed_fields: dict[str, Decorator[ComputedFieldInfo]] = field(default_factory=dict) + + @staticmethod + def build(model_dc: type[Any]) -> DecoratorInfos: # noqa: C901 (ignore complexity) + """We want to collect all DecFunc instances that exist as + attributes in the namespace of the class (a BaseModel or dataclass) + that called us + But we want to collect these in the order of the bases + So instead of getting them all from the leaf class (the class that called us), + we traverse the bases from root (the oldest ancestor class) to leaf + and collect all of the instances as we go, taking care to replace + any duplicate ones with the last one we see to mimic how function overriding + works with inheritance. + If we do replace any functions we put the replacement into the position + the replaced function was in; that is, we maintain the order. + """ + # reminder: dicts are ordered and replacement does not alter the order + res = DecoratorInfos() + for base in reversed(mro(model_dc)[1:]): + existing: DecoratorInfos | None = base.__dict__.get('__pydantic_decorators__') + if existing is None: + existing = DecoratorInfos.build(base) + res.validators.update({k: v.bind_to_cls(model_dc) for k, v in existing.validators.items()}) + res.field_validators.update({k: v.bind_to_cls(model_dc) for k, v in existing.field_validators.items()}) + res.root_validators.update({k: v.bind_to_cls(model_dc) for k, v in existing.root_validators.items()}) + res.field_serializers.update({k: v.bind_to_cls(model_dc) for k, v in existing.field_serializers.items()}) + res.model_serializers.update({k: v.bind_to_cls(model_dc) for k, v in existing.model_serializers.items()}) + res.model_validators.update({k: v.bind_to_cls(model_dc) for k, v in existing.model_validators.items()}) + res.computed_fields.update({k: v.bind_to_cls(model_dc) for k, v in existing.computed_fields.items()}) + + to_replace: list[tuple[str, Any]] = [] + + for var_name, var_value in vars(model_dc).items(): + if isinstance(var_value, PydanticDescriptorProxy): + info = var_value.decorator_info + if isinstance(info, ValidatorDecoratorInfo): + res.validators[var_name] = Decorator.build( + model_dc, cls_var_name=var_name, shim=var_value.shim, info=info + ) + elif isinstance(info, FieldValidatorDecoratorInfo): + res.field_validators[var_name] = Decorator.build( + model_dc, cls_var_name=var_name, shim=var_value.shim, info=info + ) + elif isinstance(info, RootValidatorDecoratorInfo): + res.root_validators[var_name] = Decorator.build( + model_dc, cls_var_name=var_name, shim=var_value.shim, info=info + ) + elif isinstance(info, FieldSerializerDecoratorInfo): + # check whether a serializer function is already registered for fields + for field_serializer_decorator in res.field_serializers.values(): + # check that each field has at most one serializer function. + # serializer functions for the same field in subclasses are allowed, + # and are treated as overrides + if field_serializer_decorator.cls_var_name == var_name: + continue + for f in info.fields: + if f in field_serializer_decorator.info.fields: + raise PydanticUserError( + 'Multiple field serializer functions were defined ' + f'for field {f!r}, this is not allowed.', + code='multiple-field-serializers', + ) + res.field_serializers[var_name] = Decorator.build( + model_dc, cls_var_name=var_name, shim=var_value.shim, info=info + ) + elif isinstance(info, ModelValidatorDecoratorInfo): + res.model_validators[var_name] = Decorator.build( + model_dc, cls_var_name=var_name, shim=var_value.shim, info=info + ) + elif isinstance(info, ModelSerializerDecoratorInfo): + res.model_serializers[var_name] = Decorator.build( + model_dc, cls_var_name=var_name, shim=var_value.shim, info=info + ) + else: + from ..fields import ComputedFieldInfo + + isinstance(var_value, ComputedFieldInfo) + res.computed_fields[var_name] = Decorator.build( + model_dc, cls_var_name=var_name, shim=None, info=info + ) + to_replace.append((var_name, var_value.wrapped)) + if to_replace: + # If we can save `__pydantic_decorators__` on the class we'll be able to check for it above + # so then we don't need to re-process the type, which means we can discard our descriptor wrappers + # and replace them with the thing they are wrapping (see the other setattr call below) + # which allows validator class methods to also function as regular class methods + model_dc.__pydantic_decorators__ = res + for name, value in to_replace: + setattr(model_dc, name, value) + return res + + +def inspect_validator(validator: Callable[..., Any], mode: FieldValidatorModes) -> bool: + """Look at a field or model validator function and determine whether it takes an info argument. + + An error is raised if the function has an invalid signature. + + Args: + validator: The validator function to inspect. + mode: The proposed validator mode. + + Returns: + Whether the validator takes an info argument. + """ + try: + sig = signature(validator) + except (ValueError, TypeError): + # `inspect.signature` might not be able to infer a signature, e.g. with C objects. + # In this case, we assume no info argument is present: + return False + n_positional = count_positional_required_params(sig) + if mode == 'wrap': + if n_positional == 3: + return True + elif n_positional == 2: + return False + else: + assert mode in {'before', 'after', 'plain'}, f"invalid mode: {mode!r}, expected 'before', 'after' or 'plain" + if n_positional == 2: + return True + elif n_positional == 1: + return False + + raise PydanticUserError( + f'Unrecognized field_validator function signature for {validator} with `mode={mode}`:{sig}', + code='validator-signature', + ) + + +def inspect_field_serializer(serializer: Callable[..., Any], mode: Literal['plain', 'wrap']) -> tuple[bool, bool]: + """Look at a field serializer function and determine if it is a field serializer, + and whether it takes an info argument. + + An error is raised if the function has an invalid signature. + + Args: + serializer: The serializer function to inspect. + mode: The serializer mode, either 'plain' or 'wrap'. + + Returns: + Tuple of (is_field_serializer, info_arg). + """ + try: + sig = signature(serializer) + except (ValueError, TypeError): + # `inspect.signature` might not be able to infer a signature, e.g. with C objects. + # In this case, we assume no info argument is present and this is not a method: + return (False, False) + + first = next(iter(sig.parameters.values()), None) + is_field_serializer = first is not None and first.name == 'self' + + n_positional = count_positional_required_params(sig) + if is_field_serializer: + # -1 to correct for self parameter + info_arg = _serializer_info_arg(mode, n_positional - 1) + else: + info_arg = _serializer_info_arg(mode, n_positional) + + if info_arg is None: + raise PydanticUserError( + f'Unrecognized field_serializer function signature for {serializer} with `mode={mode}`:{sig}', + code='field-serializer-signature', + ) + + return is_field_serializer, info_arg + + +def inspect_annotated_serializer(serializer: Callable[..., Any], mode: Literal['plain', 'wrap']) -> bool: + """Look at a serializer function used via `Annotated` and determine whether it takes an info argument. + + An error is raised if the function has an invalid signature. + + Args: + serializer: The serializer function to check. + mode: The serializer mode, either 'plain' or 'wrap'. + + Returns: + info_arg + """ + try: + sig = signature(serializer) + except (ValueError, TypeError): + # `inspect.signature` might not be able to infer a signature, e.g. with C objects. + # In this case, we assume no info argument is present: + return False + info_arg = _serializer_info_arg(mode, count_positional_required_params(sig)) + if info_arg is None: + raise PydanticUserError( + f'Unrecognized field_serializer function signature for {serializer} with `mode={mode}`:{sig}', + code='field-serializer-signature', + ) + else: + return info_arg + + +def inspect_model_serializer(serializer: Callable[..., Any], mode: Literal['plain', 'wrap']) -> bool: + """Look at a model serializer function and determine whether it takes an info argument. + + An error is raised if the function has an invalid signature. + + Args: + serializer: The serializer function to check. + mode: The serializer mode, either 'plain' or 'wrap'. + + Returns: + `info_arg` - whether the function expects an info argument. + """ + if isinstance(serializer, (staticmethod, classmethod)) or not is_instance_method_from_sig(serializer): + raise PydanticUserError( + '`@model_serializer` must be applied to instance methods', code='model-serializer-instance-method' + ) + + sig = signature(serializer) + info_arg = _serializer_info_arg(mode, count_positional_required_params(sig)) + if info_arg is None: + raise PydanticUserError( + f'Unrecognized model_serializer function signature for {serializer} with `mode={mode}`:{sig}', + code='model-serializer-signature', + ) + else: + return info_arg + + +def _serializer_info_arg(mode: Literal['plain', 'wrap'], n_positional: int) -> bool | None: + if mode == 'plain': + if n_positional == 1: + # (input_value: Any, /) -> Any + return False + elif n_positional == 2: + # (model: Any, input_value: Any, /) -> Any + return True + else: + assert mode == 'wrap', f"invalid mode: {mode!r}, expected 'plain' or 'wrap'" + if n_positional == 2: + # (input_value: Any, serializer: SerializerFunctionWrapHandler, /) -> Any + return False + elif n_positional == 3: + # (input_value: Any, serializer: SerializerFunctionWrapHandler, info: SerializationInfo, /) -> Any + return True + + return None + + +AnyDecoratorCallable: TypeAlias = ( + 'Union[classmethod[Any, Any, Any], staticmethod[Any, Any], partialmethod[Any], Callable[..., Any]]' +) + + +def is_instance_method_from_sig(function: AnyDecoratorCallable) -> bool: + """Whether the function is an instance method. + + It will consider a function as instance method if the first parameter of + function is `self`. + + Args: + function: The function to check. + + Returns: + `True` if the function is an instance method, `False` otherwise. + """ + sig = signature(unwrap_wrapped_function(function)) + first = next(iter(sig.parameters.values()), None) + if first and first.name == 'self': + return True + return False + + +def ensure_classmethod_based_on_signature(function: AnyDecoratorCallable) -> Any: + """Apply the `@classmethod` decorator on the function. + + Args: + function: The function to apply the decorator on. + + Return: + The `@classmethod` decorator applied function. + """ + if not isinstance( + unwrap_wrapped_function(function, unwrap_class_static_method=False), classmethod + ) and _is_classmethod_from_sig(function): + return classmethod(function) # type: ignore[arg-type] + return function + + +def _is_classmethod_from_sig(function: AnyDecoratorCallable) -> bool: + sig = signature(unwrap_wrapped_function(function)) + first = next(iter(sig.parameters.values()), None) + if first and first.name == 'cls': + return True + return False + + +def unwrap_wrapped_function( + func: Any, + *, + unwrap_partial: bool = True, + unwrap_class_static_method: bool = True, +) -> Any: + """Recursively unwraps a wrapped function until the underlying function is reached. + This handles property, functools.partial, functools.partialmethod, staticmethod, and classmethod. + + Args: + func: The function to unwrap. + unwrap_partial: If True (default), unwrap partial and partialmethod decorators. + unwrap_class_static_method: If True (default), also unwrap classmethod and staticmethod + decorators. If False, only unwrap partial and partialmethod decorators. + + Returns: + The underlying function of the wrapped function. + """ + # Define the types we want to check against as a single tuple. + unwrap_types = ( + (property, cached_property) + + ((partial, partialmethod) if unwrap_partial else ()) + + ((staticmethod, classmethod) if unwrap_class_static_method else ()) + ) + + while isinstance(func, unwrap_types): + if unwrap_class_static_method and isinstance(func, (classmethod, staticmethod)): + func = func.__func__ + elif isinstance(func, (partial, partialmethod)): + func = func.func + elif isinstance(func, property): + func = func.fget # arbitrary choice, convenient for computed fields + else: + # Make coverage happy as it can only get here in the last possible case + assert isinstance(func, cached_property) + func = func.func # type: ignore + + return func + + +_function_like = ( + partial, + partialmethod, + types.FunctionType, + types.BuiltinFunctionType, + types.MethodType, + types.WrapperDescriptorType, + types.MethodWrapperType, + types.MemberDescriptorType, +) + + +def get_callable_return_type( + callable_obj: Any, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, +) -> Any | PydanticUndefinedType: + """Get the callable return type. + + Args: + callable_obj: The callable to analyze. + globalns: The globals namespace to use during type annotation evaluation. + localns: The locals namespace to use during type annotation evaluation. + + Returns: + The function return type. + """ + if isinstance(callable_obj, type): + # types are callables, and we assume the return type + # is the type itself (e.g. `int()` results in an instance of `int`). + return callable_obj + + if not isinstance(callable_obj, _function_like): + call_func = getattr(type(callable_obj), '__call__', None) # noqa: B004 + if call_func is not None: + callable_obj = call_func + + hints = get_function_type_hints( + unwrap_wrapped_function(callable_obj), + include_keys={'return'}, + globalns=globalns, + localns=localns, + ) + return hints.get('return', PydanticUndefined) + + +def count_positional_required_params(sig: Signature) -> int: + """Get the number of positional (required) arguments of a signature. + + This function should only be used to inspect signatures of validation and serialization functions. + The first argument (the value being serialized or validated) is counted as a required argument + even if a default value exists. + + Returns: + The number of positional arguments of a signature. + """ + parameters = list(sig.parameters.values()) + return sum( + 1 + for param in parameters + if can_be_positional(param) + # First argument is the value being validated/serialized, and can have a default value + # (e.g. `float`, which has signature `(x=0, /)`). We assume other parameters (the info arg + # for instance) should be required, and thus without any default value. + and (param.default is Parameter.empty or param is parameters[0]) + ) + + +def ensure_property(f: Any) -> Any: + """Ensure that a function is a `property` or `cached_property`, or is a valid descriptor. + + Args: + f: The function to check. + + Returns: + The function, or a `property` or `cached_property` instance wrapping the function. + """ + if ismethoddescriptor(f) or isdatadescriptor(f): + return f + else: + return property(f) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_decorators_v1.py b/lib/python3.12/site-packages/pydantic/_internal/_decorators_v1.py new file mode 100644 index 0000000000000000000000000000000000000000..34273779d77666f5cd775f6a2684641092db67ee --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_decorators_v1.py @@ -0,0 +1,174 @@ +"""Logic for V1 validators, e.g. `@validator` and `@root_validator`.""" + +from __future__ import annotations as _annotations + +from inspect import Parameter, signature +from typing import Any, Union, cast + +from pydantic_core import core_schema +from typing_extensions import Protocol + +from ..errors import PydanticUserError +from ._utils import can_be_positional + + +class V1OnlyValueValidator(Protocol): + """A simple validator, supported for V1 validators and V2 validators.""" + + def __call__(self, __value: Any) -> Any: ... + + +class V1ValidatorWithValues(Protocol): + """A validator with `values` argument, supported for V1 validators and V2 validators.""" + + def __call__(self, __value: Any, values: dict[str, Any]) -> Any: ... + + +class V1ValidatorWithValuesKwOnly(Protocol): + """A validator with keyword only `values` argument, supported for V1 validators and V2 validators.""" + + def __call__(self, __value: Any, *, values: dict[str, Any]) -> Any: ... + + +class V1ValidatorWithKwargs(Protocol): + """A validator with `kwargs` argument, supported for V1 validators and V2 validators.""" + + def __call__(self, __value: Any, **kwargs: Any) -> Any: ... + + +class V1ValidatorWithValuesAndKwargs(Protocol): + """A validator with `values` and `kwargs` arguments, supported for V1 validators and V2 validators.""" + + def __call__(self, __value: Any, values: dict[str, Any], **kwargs: Any) -> Any: ... + + +V1Validator = Union[ + V1ValidatorWithValues, V1ValidatorWithValuesKwOnly, V1ValidatorWithKwargs, V1ValidatorWithValuesAndKwargs +] + + +def can_be_keyword(param: Parameter) -> bool: + return param.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) + + +def make_generic_v1_field_validator(validator: V1Validator) -> core_schema.WithInfoValidatorFunction: + """Wrap a V1 style field validator for V2 compatibility. + + Args: + validator: The V1 style field validator. + + Returns: + A wrapped V2 style field validator. + + Raises: + PydanticUserError: If the signature is not supported or the parameters are + not available in Pydantic V2. + """ + sig = signature(validator) + + needs_values_kw = False + + for param_num, (param_name, parameter) in enumerate(sig.parameters.items()): + if can_be_keyword(parameter) and param_name in ('field', 'config'): + raise PydanticUserError( + 'The `field` and `config` parameters are not available in Pydantic V2, ' + 'please use the `info` parameter instead.', + code='validator-field-config-info', + ) + if parameter.kind is Parameter.VAR_KEYWORD: + needs_values_kw = True + elif can_be_keyword(parameter) and param_name == 'values': + needs_values_kw = True + elif can_be_positional(parameter) and param_num == 0: + # value + continue + elif parameter.default is Parameter.empty: # ignore params with defaults e.g. bound by functools.partial + raise PydanticUserError( + f'Unsupported signature for V1 style validator {validator}: {sig} is not supported.', + code='validator-v1-signature', + ) + + if needs_values_kw: + # (v, **kwargs), (v, values, **kwargs), (v, *, values, **kwargs) or (v, *, values) + val1 = cast(V1ValidatorWithValues, validator) + + def wrapper1(value: Any, info: core_schema.ValidationInfo) -> Any: + return val1(value, values=info.data) + + return wrapper1 + else: + val2 = cast(V1OnlyValueValidator, validator) + + def wrapper2(value: Any, _: core_schema.ValidationInfo) -> Any: + return val2(value) + + return wrapper2 + + +RootValidatorValues = dict[str, Any] +# technically tuple[model_dict, model_extra, fields_set] | tuple[dataclass_dict, init_vars] +RootValidatorFieldsTuple = tuple[Any, ...] + + +class V1RootValidatorFunction(Protocol): + """A simple root validator, supported for V1 validators and V2 validators.""" + + def __call__(self, __values: RootValidatorValues) -> RootValidatorValues: ... + + +class V2CoreBeforeRootValidator(Protocol): + """V2 validator with mode='before'.""" + + def __call__(self, __values: RootValidatorValues, __info: core_schema.ValidationInfo) -> RootValidatorValues: ... + + +class V2CoreAfterRootValidator(Protocol): + """V2 validator with mode='after'.""" + + def __call__( + self, __fields_tuple: RootValidatorFieldsTuple, __info: core_schema.ValidationInfo + ) -> RootValidatorFieldsTuple: ... + + +def make_v1_generic_root_validator( + validator: V1RootValidatorFunction, pre: bool +) -> V2CoreBeforeRootValidator | V2CoreAfterRootValidator: + """Wrap a V1 style root validator for V2 compatibility. + + Args: + validator: The V1 style field validator. + pre: Whether the validator is a pre validator. + + Returns: + A wrapped V2 style validator. + """ + if pre is True: + # mode='before' for pydantic-core + def _wrapper1(values: RootValidatorValues, _: core_schema.ValidationInfo) -> RootValidatorValues: + return validator(values) + + return _wrapper1 + + # mode='after' for pydantic-core + def _wrapper2(fields_tuple: RootValidatorFieldsTuple, _: core_schema.ValidationInfo) -> RootValidatorFieldsTuple: + if len(fields_tuple) == 2: + # dataclass, this is easy + values, init_vars = fields_tuple + values = validator(values) + return values, init_vars + else: + # ugly hack: to match v1 behaviour, we merge values and model_extra, then split them up based on fields + # afterwards + model_dict, model_extra, fields_set = fields_tuple + if model_extra: + fields = set(model_dict.keys()) + model_dict.update(model_extra) + model_dict_new = validator(model_dict) + for k in list(model_dict_new.keys()): + if k not in fields: + model_extra[k] = model_dict_new.pop(k) + else: + model_dict_new = validator(model_dict) + return model_dict_new, model_extra, fields_set + + return _wrapper2 diff --git a/lib/python3.12/site-packages/pydantic/_internal/_discriminated_union.py b/lib/python3.12/site-packages/pydantic/_internal/_discriminated_union.py new file mode 100644 index 0000000000000000000000000000000000000000..5dd6fdaf4182be93ba9891dec8be54ef0e201bcb --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_discriminated_union.py @@ -0,0 +1,479 @@ +from __future__ import annotations as _annotations + +from collections.abc import Hashable, Sequence +from typing import TYPE_CHECKING, Any, cast + +from pydantic_core import CoreSchema, core_schema + +from ..errors import PydanticUserError +from . import _core_utils +from ._core_utils import ( + CoreSchemaField, +) + +if TYPE_CHECKING: + from ..types import Discriminator + from ._core_metadata import CoreMetadata + + +class MissingDefinitionForUnionRef(Exception): + """Raised when applying a discriminated union discriminator to a schema + requires a definition that is not yet defined + """ + + def __init__(self, ref: str) -> None: + self.ref = ref + super().__init__(f'Missing definition for ref {self.ref!r}') + + +def set_discriminator_in_metadata(schema: CoreSchema, discriminator: Any) -> None: + metadata = cast('CoreMetadata', schema.setdefault('metadata', {})) + metadata['pydantic_internal_union_discriminator'] = discriminator + + +def apply_discriminator( + schema: core_schema.CoreSchema, + discriminator: str | Discriminator, + definitions: dict[str, core_schema.CoreSchema] | None = None, +) -> core_schema.CoreSchema: + """Applies the discriminator and returns a new core schema. + + Args: + schema: The input schema. + discriminator: The name of the field which will serve as the discriminator. + definitions: A mapping of schema ref to schema. + + Returns: + The new core schema. + + Raises: + TypeError: + - If `discriminator` is used with invalid union variant. + - If `discriminator` is used with `Union` type with one variant. + - If `discriminator` value mapped to multiple choices. + MissingDefinitionForUnionRef: + If the definition for ref is missing. + PydanticUserError: + - If a model in union doesn't have a discriminator field. + - If discriminator field has a non-string alias. + - If discriminator fields have different aliases. + - If discriminator field not of type `Literal`. + """ + from ..types import Discriminator + + if isinstance(discriminator, Discriminator): + if isinstance(discriminator.discriminator, str): + discriminator = discriminator.discriminator + else: + return discriminator._convert_schema(schema) + + return _ApplyInferredDiscriminator(discriminator, definitions or {}).apply(schema) + + +class _ApplyInferredDiscriminator: + """This class is used to convert an input schema containing a union schema into one where that union is + replaced with a tagged-union, with all the associated debugging and performance benefits. + + This is done by: + * Validating that the input schema is compatible with the provided discriminator + * Introspecting the schema to determine which discriminator values should map to which union choices + * Handling various edge cases such as 'definitions', 'default', 'nullable' schemas, and more + + I have chosen to implement the conversion algorithm in this class, rather than a function, + to make it easier to maintain state while recursively walking the provided CoreSchema. + """ + + def __init__(self, discriminator: str, definitions: dict[str, core_schema.CoreSchema]): + # `discriminator` should be the name of the field which will serve as the discriminator. + # It must be the python name of the field, and *not* the field's alias. Note that as of now, + # all members of a discriminated union _must_ use a field with the same name as the discriminator. + # This may change if/when we expose a way to manually specify the TaggedUnionSchema's choices. + self.discriminator = discriminator + + # `definitions` should contain a mapping of schema ref to schema for all schemas which might + # be referenced by some choice + self.definitions = definitions + + # `_discriminator_alias` will hold the value, if present, of the alias for the discriminator + # + # Note: following the v1 implementation, we currently disallow the use of different aliases + # for different choices. This is not a limitation of pydantic_core, but if we try to handle + # this, the inference logic gets complicated very quickly, and could result in confusing + # debugging challenges for users making subtle mistakes. + # + # Rather than trying to do the most powerful inference possible, I think we should eventually + # expose a way to more-manually control the way the TaggedUnionSchema is constructed through + # the use of a new type which would be placed as an Annotation on the Union type. This would + # provide the full flexibility/power of pydantic_core's TaggedUnionSchema where necessary for + # more complex cases, without over-complicating the inference logic for the common cases. + self._discriminator_alias: str | None = None + + # `_should_be_nullable` indicates whether the converted union has `None` as an allowed value. + # If `None` is an acceptable value of the (possibly-wrapped) union, we ignore it while + # constructing the TaggedUnionSchema, but set the `_should_be_nullable` attribute to True. + # Once we have constructed the TaggedUnionSchema, if `_should_be_nullable` is True, we ensure + # that the final schema gets wrapped as a NullableSchema. This has the same semantics on the + # python side, but resolves the issue that `None` cannot correspond to any discriminator values. + self._should_be_nullable = False + + # `_is_nullable` is used to track if the final produced schema will definitely be nullable; + # we set it to True if the input schema is wrapped in a nullable schema that we know will be preserved + # as an indication that, even if None is discovered as one of the union choices, we will not need to wrap + # the final value in another nullable schema. + # + # This is more complicated than just checking for the final outermost schema having type 'nullable' thanks + # to the possible presence of other wrapper schemas such as DefinitionsSchema, WithDefaultSchema, etc. + self._is_nullable = False + + # `_choices_to_handle` serves as a stack of choices to add to the tagged union. Initially, choices + # from the union in the wrapped schema will be appended to this list, and the recursive choice-handling + # algorithm may add more choices to this stack as (nested) unions are encountered. + self._choices_to_handle: list[core_schema.CoreSchema] = [] + + # `_tagged_union_choices` is built during the call to `apply`, and will hold the choices to be included + # in the output TaggedUnionSchema that will replace the union from the input schema + self._tagged_union_choices: dict[Hashable, core_schema.CoreSchema] = {} + + # `_used` is changed to True after applying the discriminator to prevent accidental reuse + self._used = False + + def apply(self, schema: core_schema.CoreSchema) -> core_schema.CoreSchema: + """Return a new CoreSchema based on `schema` that uses a tagged-union with the discriminator provided + to this class. + + Args: + schema: The input schema. + + Returns: + The new core schema. + + Raises: + TypeError: + - If `discriminator` is used with invalid union variant. + - If `discriminator` is used with `Union` type with one variant. + - If `discriminator` value mapped to multiple choices. + ValueError: + If the definition for ref is missing. + PydanticUserError: + - If a model in union doesn't have a discriminator field. + - If discriminator field has a non-string alias. + - If discriminator fields have different aliases. + - If discriminator field not of type `Literal`. + """ + assert not self._used + schema = self._apply_to_root(schema) + if self._should_be_nullable and not self._is_nullable: + schema = core_schema.nullable_schema(schema) + self._used = True + return schema + + def _apply_to_root(self, schema: core_schema.CoreSchema) -> core_schema.CoreSchema: + """This method handles the outer-most stage of recursion over the input schema: + unwrapping nullable or definitions schemas, and calling the `_handle_choice` + method iteratively on the choices extracted (recursively) from the possibly-wrapped union. + """ + if schema['type'] == 'nullable': + self._is_nullable = True + wrapped = self._apply_to_root(schema['schema']) + nullable_wrapper = schema.copy() + nullable_wrapper['schema'] = wrapped + return nullable_wrapper + + if schema['type'] == 'definitions': + wrapped = self._apply_to_root(schema['schema']) + definitions_wrapper = schema.copy() + definitions_wrapper['schema'] = wrapped + return definitions_wrapper + + if schema['type'] != 'union': + # If the schema is not a union, it probably means it just had a single member and + # was flattened by pydantic_core. + # However, it still may make sense to apply the discriminator to this schema, + # as a way to get discriminated-union-style error messages, so we allow this here. + schema = core_schema.union_schema([schema]) + + # Reverse the choices list before extending the stack so that they get handled in the order they occur + choices_schemas = [v[0] if isinstance(v, tuple) else v for v in schema['choices'][::-1]] + self._choices_to_handle.extend(choices_schemas) + while self._choices_to_handle: + choice = self._choices_to_handle.pop() + self._handle_choice(choice) + + if self._discriminator_alias is not None and self._discriminator_alias != self.discriminator: + # * We need to annotate `discriminator` as a union here to handle both branches of this conditional + # * We need to annotate `discriminator` as list[list[str | int]] and not list[list[str]] due to the + # invariance of list, and because list[list[str | int]] is the type of the discriminator argument + # to tagged_union_schema below + # * See the docstring of pydantic_core.core_schema.tagged_union_schema for more details about how to + # interpret the value of the discriminator argument to tagged_union_schema. (The list[list[str]] here + # is the appropriate way to provide a list of fallback attributes to check for a discriminator value.) + discriminator: str | list[list[str | int]] = [[self.discriminator], [self._discriminator_alias]] + else: + discriminator = self.discriminator + return core_schema.tagged_union_schema( + choices=self._tagged_union_choices, + discriminator=discriminator, + custom_error_type=schema.get('custom_error_type'), + custom_error_message=schema.get('custom_error_message'), + custom_error_context=schema.get('custom_error_context'), + strict=False, + from_attributes=True, + ref=schema.get('ref'), + metadata=schema.get('metadata'), + serialization=schema.get('serialization'), + ) + + def _handle_choice(self, choice: core_schema.CoreSchema) -> None: + """This method handles the "middle" stage of recursion over the input schema. + Specifically, it is responsible for handling each choice of the outermost union + (and any "coalesced" choices obtained from inner unions). + + Here, "handling" entails: + * Coalescing nested unions and compatible tagged-unions + * Tracking the presence of 'none' and 'nullable' schemas occurring as choices + * Validating that each allowed discriminator value maps to a unique choice + * Updating the _tagged_union_choices mapping that will ultimately be used to build the TaggedUnionSchema. + """ + if choice['type'] == 'definition-ref': + if choice['schema_ref'] not in self.definitions: + raise MissingDefinitionForUnionRef(choice['schema_ref']) + + if choice['type'] == 'none': + self._should_be_nullable = True + elif choice['type'] == 'definitions': + self._handle_choice(choice['schema']) + elif choice['type'] == 'nullable': + self._should_be_nullable = True + self._handle_choice(choice['schema']) # unwrap the nullable schema + elif choice['type'] == 'union': + # Reverse the choices list before extending the stack so that they get handled in the order they occur + choices_schemas = [v[0] if isinstance(v, tuple) else v for v in choice['choices'][::-1]] + self._choices_to_handle.extend(choices_schemas) + elif choice['type'] not in { + 'model', + 'typed-dict', + 'tagged-union', + 'lax-or-strict', + 'dataclass', + 'dataclass-args', + 'definition-ref', + } and not _core_utils.is_function_with_inner_schema(choice): + # We should eventually handle 'definition-ref' as well + err_str = f'The core schema type {choice["type"]!r} is not a valid discriminated union variant.' + if choice['type'] == 'list': + err_str += ( + ' If you are making use of a list of union types, make sure the discriminator is applied to the ' + 'union type and not the list (e.g. `list[Annotated[ | , Field(discriminator=...)]]`).' + ) + raise TypeError(err_str) + else: + if choice['type'] == 'tagged-union' and self._is_discriminator_shared(choice): + # In this case, this inner tagged-union is compatible with the outer tagged-union, + # and its choices can be coalesced into the outer TaggedUnionSchema. + subchoices = [x for x in choice['choices'].values() if not isinstance(x, (str, int))] + # Reverse the choices list before extending the stack so that they get handled in the order they occur + self._choices_to_handle.extend(subchoices[::-1]) + return + + inferred_discriminator_values = self._infer_discriminator_values_for_choice(choice, source_name=None) + self._set_unique_choice_for_values(choice, inferred_discriminator_values) + + def _is_discriminator_shared(self, choice: core_schema.TaggedUnionSchema) -> bool: + """This method returns a boolean indicating whether the discriminator for the `choice` + is the same as that being used for the outermost tagged union. This is used to + determine whether this TaggedUnionSchema choice should be "coalesced" into the top level, + or whether it should be treated as a separate (nested) choice. + """ + inner_discriminator = choice['discriminator'] + return inner_discriminator == self.discriminator or ( + isinstance(inner_discriminator, list) + and (self.discriminator in inner_discriminator or [self.discriminator] in inner_discriminator) + ) + + def _infer_discriminator_values_for_choice( # noqa C901 + self, choice: core_schema.CoreSchema, source_name: str | None + ) -> list[str | int]: + """This function recurses over `choice`, extracting all discriminator values that should map to this choice. + + `model_name` is accepted for the purpose of producing useful error messages. + """ + if choice['type'] == 'definitions': + return self._infer_discriminator_values_for_choice(choice['schema'], source_name=source_name) + + elif _core_utils.is_function_with_inner_schema(choice): + return self._infer_discriminator_values_for_choice(choice['schema'], source_name=source_name) + + elif choice['type'] == 'lax-or-strict': + return sorted( + set( + self._infer_discriminator_values_for_choice(choice['lax_schema'], source_name=None) + + self._infer_discriminator_values_for_choice(choice['strict_schema'], source_name=None) + ) + ) + + elif choice['type'] == 'tagged-union': + values: list[str | int] = [] + # Ignore str/int "choices" since these are just references to other choices + subchoices = [x for x in choice['choices'].values() if not isinstance(x, (str, int))] + for subchoice in subchoices: + subchoice_values = self._infer_discriminator_values_for_choice(subchoice, source_name=None) + values.extend(subchoice_values) + return values + + elif choice['type'] == 'union': + values = [] + for subchoice in choice['choices']: + subchoice_schema = subchoice[0] if isinstance(subchoice, tuple) else subchoice + subchoice_values = self._infer_discriminator_values_for_choice(subchoice_schema, source_name=None) + values.extend(subchoice_values) + return values + + elif choice['type'] == 'nullable': + self._should_be_nullable = True + return self._infer_discriminator_values_for_choice(choice['schema'], source_name=None) + + elif choice['type'] == 'model': + return self._infer_discriminator_values_for_choice(choice['schema'], source_name=choice['cls'].__name__) + + elif choice['type'] == 'dataclass': + return self._infer_discriminator_values_for_choice(choice['schema'], source_name=choice['cls'].__name__) + + elif choice['type'] == 'model-fields': + return self._infer_discriminator_values_for_model_choice(choice, source_name=source_name) + + elif choice['type'] == 'dataclass-args': + return self._infer_discriminator_values_for_dataclass_choice(choice, source_name=source_name) + + elif choice['type'] == 'typed-dict': + return self._infer_discriminator_values_for_typed_dict_choice(choice, source_name=source_name) + + elif choice['type'] == 'definition-ref': + schema_ref = choice['schema_ref'] + if schema_ref not in self.definitions: + raise MissingDefinitionForUnionRef(schema_ref) + return self._infer_discriminator_values_for_choice(self.definitions[schema_ref], source_name=source_name) + else: + err_str = f'The core schema type {choice["type"]!r} is not a valid discriminated union variant.' + if choice['type'] == 'list': + err_str += ( + ' If you are making use of a list of union types, make sure the discriminator is applied to the ' + 'union type and not the list (e.g. `list[Annotated[ | , Field(discriminator=...)]]`).' + ) + raise TypeError(err_str) + + def _infer_discriminator_values_for_typed_dict_choice( + self, choice: core_schema.TypedDictSchema, source_name: str | None = None + ) -> list[str | int]: + """This method just extracts the _infer_discriminator_values_for_choice logic specific to TypedDictSchema + for the sake of readability. + """ + source = 'TypedDict' if source_name is None else f'TypedDict {source_name!r}' + field = choice['fields'].get(self.discriminator) + if field is None: + raise PydanticUserError( + f'{source} needs a discriminator field for key {self.discriminator!r}', code='discriminator-no-field' + ) + return self._infer_discriminator_values_for_field(field, source) + + def _infer_discriminator_values_for_model_choice( + self, choice: core_schema.ModelFieldsSchema, source_name: str | None = None + ) -> list[str | int]: + source = 'ModelFields' if source_name is None else f'Model {source_name!r}' + field = choice['fields'].get(self.discriminator) + if field is None: + raise PydanticUserError( + f'{source} needs a discriminator field for key {self.discriminator!r}', code='discriminator-no-field' + ) + return self._infer_discriminator_values_for_field(field, source) + + def _infer_discriminator_values_for_dataclass_choice( + self, choice: core_schema.DataclassArgsSchema, source_name: str | None = None + ) -> list[str | int]: + source = 'DataclassArgs' if source_name is None else f'Dataclass {source_name!r}' + for field in choice['fields']: + if field['name'] == self.discriminator: + break + else: + raise PydanticUserError( + f'{source} needs a discriminator field for key {self.discriminator!r}', code='discriminator-no-field' + ) + return self._infer_discriminator_values_for_field(field, source) + + def _infer_discriminator_values_for_field(self, field: CoreSchemaField, source: str) -> list[str | int]: + if field['type'] == 'computed-field': + # This should never occur as a discriminator, as it is only relevant to serialization + return [] + alias = field.get('validation_alias', self.discriminator) + if not isinstance(alias, str): + raise PydanticUserError( + f'Alias {alias!r} is not supported in a discriminated union', code='discriminator-alias-type' + ) + if self._discriminator_alias is None: + self._discriminator_alias = alias + elif self._discriminator_alias != alias: + raise PydanticUserError( + f'Aliases for discriminator {self.discriminator!r} must be the same ' + f'(got {alias}, {self._discriminator_alias})', + code='discriminator-alias', + ) + return self._infer_discriminator_values_for_inner_schema(field['schema'], source) + + def _infer_discriminator_values_for_inner_schema( + self, schema: core_schema.CoreSchema, source: str + ) -> list[str | int]: + """When inferring discriminator values for a field, we typically extract the expected values from a literal + schema. This function does that, but also handles nested unions and defaults. + """ + if schema['type'] == 'literal': + return schema['expected'] + + elif schema['type'] == 'union': + # Generally when multiple values are allowed they should be placed in a single `Literal`, but + # we add this case to handle the situation where a field is annotated as a `Union` of `Literal`s. + # For example, this lets us handle `Union[Literal['key'], Union[Literal['Key'], Literal['KEY']]]` + values: list[Any] = [] + for choice in schema['choices']: + choice_schema = choice[0] if isinstance(choice, tuple) else choice + choice_values = self._infer_discriminator_values_for_inner_schema(choice_schema, source) + values.extend(choice_values) + return values + + elif schema['type'] == 'default': + # This will happen if the field has a default value; we ignore it while extracting the discriminator values + return self._infer_discriminator_values_for_inner_schema(schema['schema'], source) + + elif schema['type'] == 'function-after': + # After validators don't affect the discriminator values + return self._infer_discriminator_values_for_inner_schema(schema['schema'], source) + + elif schema['type'] in {'function-before', 'function-wrap', 'function-plain'}: + validator_type = repr(schema['type'].split('-')[1]) + raise PydanticUserError( + f'Cannot use a mode={validator_type} validator in the' + f' discriminator field {self.discriminator!r} of {source}', + code='discriminator-validator', + ) + + else: + raise PydanticUserError( + f'{source} needs field {self.discriminator!r} to be of type `Literal`', + code='discriminator-needs-literal', + ) + + def _set_unique_choice_for_values(self, choice: core_schema.CoreSchema, values: Sequence[str | int]) -> None: + """This method updates `self.tagged_union_choices` so that all provided (discriminator) `values` map to the + provided `choice`, validating that none of these values already map to another (different) choice. + """ + for discriminator_value in values: + if discriminator_value in self._tagged_union_choices: + # It is okay if `value` is already in tagged_union_choices as long as it maps to the same value. + # Because tagged_union_choices may map values to other values, we need to walk the choices dict + # until we get to a "real" choice, and confirm that is equal to the one assigned. + existing_choice = self._tagged_union_choices[discriminator_value] + if existing_choice != choice: + raise TypeError( + f'Value {discriminator_value!r} for discriminator ' + f'{self.discriminator!r} mapped to multiple choices' + ) + else: + self._tagged_union_choices[discriminator_value] = choice diff --git a/lib/python3.12/site-packages/pydantic/_internal/_docs_extraction.py b/lib/python3.12/site-packages/pydantic/_internal/_docs_extraction.py new file mode 100644 index 0000000000000000000000000000000000000000..7b5f310fe2211e5be1706ce11a4ee5c8ad327af8 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_docs_extraction.py @@ -0,0 +1,108 @@ +"""Utilities related to attribute docstring extraction.""" + +from __future__ import annotations + +import ast +import inspect +import textwrap +from typing import Any + + +class DocstringVisitor(ast.NodeVisitor): + def __init__(self) -> None: + super().__init__() + + self.target: str | None = None + self.attrs: dict[str, str] = {} + self.previous_node_type: type[ast.AST] | None = None + + def visit(self, node: ast.AST) -> Any: + node_result = super().visit(node) + self.previous_node_type = type(node) + return node_result + + def visit_AnnAssign(self, node: ast.AnnAssign) -> Any: + if isinstance(node.target, ast.Name): + self.target = node.target.id + + def visit_Expr(self, node: ast.Expr) -> Any: + if ( + isinstance(node.value, ast.Constant) + and isinstance(node.value.value, str) + and self.previous_node_type is ast.AnnAssign + ): + docstring = inspect.cleandoc(node.value.value) + if self.target: + self.attrs[self.target] = docstring + self.target = None + + +def _dedent_source_lines(source: list[str]) -> str: + # Required for nested class definitions, e.g. in a function block + dedent_source = textwrap.dedent(''.join(source)) + if dedent_source.startswith((' ', '\t')): + # We are in the case where there's a dedented (usually multiline) string + # at a lower indentation level than the class itself. We wrap our class + # in a function as a workaround. + dedent_source = f'def dedent_workaround():\n{dedent_source}' + return dedent_source + + +def _extract_source_from_frame(cls: type[Any]) -> list[str] | None: + frame = inspect.currentframe() + + while frame: + if inspect.getmodule(frame) is inspect.getmodule(cls): + lnum = frame.f_lineno + try: + lines, _ = inspect.findsource(frame) + except OSError: # pragma: no cover + # Source can't be retrieved (maybe because running in an interactive terminal), + # we don't want to error here. + pass + else: + block_lines = inspect.getblock(lines[lnum - 1 :]) + dedent_source = _dedent_source_lines(block_lines) + try: + block_tree = ast.parse(dedent_source) + except SyntaxError: + pass + else: + stmt = block_tree.body[0] + if isinstance(stmt, ast.FunctionDef) and stmt.name == 'dedent_workaround': + # `_dedent_source_lines` wrapped the class around the workaround function + stmt = stmt.body[0] + if isinstance(stmt, ast.ClassDef) and stmt.name == cls.__name__: + return block_lines + + frame = frame.f_back + + +def extract_docstrings_from_cls(cls: type[Any], use_inspect: bool = False) -> dict[str, str]: + """Map model attributes and their corresponding docstring. + + Args: + cls: The class of the Pydantic model to inspect. + use_inspect: Whether to skip usage of frames to find the object and use + the `inspect` module instead. + + Returns: + A mapping containing attribute names and their corresponding docstring. + """ + if use_inspect: + # Might not work as expected if two classes have the same name in the same source file. + try: + source, _ = inspect.getsourcelines(cls) + except OSError: # pragma: no cover + return {} + else: + source = _extract_source_from_frame(cls) + + if not source: + return {} + + dedent_source = _dedent_source_lines(source) + + visitor = DocstringVisitor() + visitor.visit(ast.parse(dedent_source)) + return visitor.attrs diff --git a/lib/python3.12/site-packages/pydantic/_internal/_fields.py b/lib/python3.12/site-packages/pydantic/_internal/_fields.py new file mode 100644 index 0000000000000000000000000000000000000000..658be3baaa51b46ad9650434ceead39fae1f73e2 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_fields.py @@ -0,0 +1,515 @@ +"""Private logic related to fields (the `Field()` function and `FieldInfo` class), and arguments to `Annotated`.""" + +from __future__ import annotations as _annotations + +import dataclasses +import warnings +from collections.abc import Mapping +from copy import copy +from functools import cache +from inspect import Parameter, ismethoddescriptor, signature +from re import Pattern +from typing import TYPE_CHECKING, Any, Callable, TypeVar + +from pydantic_core import PydanticUndefined +from typing_extensions import TypeIs, get_origin +from typing_inspection import typing_objects +from typing_inspection.introspection import AnnotationSource + +from pydantic import PydanticDeprecatedSince211 +from pydantic.errors import PydanticUserError + +from . import _generics, _typing_extra +from ._config import ConfigWrapper +from ._docs_extraction import extract_docstrings_from_cls +from ._import_utils import import_cached_base_model, import_cached_field_info +from ._namespace_utils import NsResolver +from ._repr import Representation +from ._utils import can_be_positional + +if TYPE_CHECKING: + from annotated_types import BaseMetadata + + from ..fields import FieldInfo + from ..main import BaseModel + from ._dataclasses import PydanticDataclass, StandardDataclass + from ._decorators import DecoratorInfos + + +class PydanticMetadata(Representation): + """Base class for annotation markers like `Strict`.""" + + __slots__ = () + + +def pydantic_general_metadata(**metadata: Any) -> BaseMetadata: + """Create a new `_PydanticGeneralMetadata` class with the given metadata. + + Args: + **metadata: The metadata to add. + + Returns: + The new `_PydanticGeneralMetadata` class. + """ + return _general_metadata_cls()(metadata) # type: ignore + + +@cache +def _general_metadata_cls() -> type[BaseMetadata]: + """Do it this way to avoid importing `annotated_types` at import time.""" + from annotated_types import BaseMetadata + + class _PydanticGeneralMetadata(PydanticMetadata, BaseMetadata): + """Pydantic general metadata like `max_digits`.""" + + def __init__(self, metadata: Any): + self.__dict__ = metadata + + return _PydanticGeneralMetadata # type: ignore + + +def _update_fields_from_docstrings(cls: type[Any], fields: dict[str, FieldInfo], use_inspect: bool = False) -> None: + fields_docs = extract_docstrings_from_cls(cls, use_inspect=use_inspect) + for ann_name, field_info in fields.items(): + if field_info.description is None and ann_name in fields_docs: + field_info.description = fields_docs[ann_name] + + +def collect_model_fields( # noqa: C901 + cls: type[BaseModel], + config_wrapper: ConfigWrapper, + ns_resolver: NsResolver | None, + *, + typevars_map: Mapping[TypeVar, Any] | None = None, +) -> tuple[dict[str, FieldInfo], set[str]]: + """Collect the fields and class variables names of a nascent Pydantic model. + + The fields collection process is *lenient*, meaning it won't error if string annotations + fail to evaluate. If this happens, the original annotation (and assigned value, if any) + is stored on the created `FieldInfo` instance. + + The `rebuild_model_fields()` should be called at a later point (e.g. when rebuilding the model), + and will make use of these stored attributes. + + Args: + cls: BaseModel or dataclass. + config_wrapper: The config wrapper instance. + ns_resolver: Namespace resolver to use when getting model annotations. + typevars_map: A dictionary mapping type variables to their concrete types. + + Returns: + A two-tuple containing model fields and class variables names. + + Raises: + NameError: + - If there is a conflict between a field name and protected namespaces. + - If there is a field other than `root` in `RootModel`. + - If a field shadows an attribute in the parent model. + """ + BaseModel = import_cached_base_model() + FieldInfo_ = import_cached_field_info() + + bases = cls.__bases__ + parent_fields_lookup: dict[str, FieldInfo] = {} + for base in reversed(bases): + if model_fields := getattr(base, '__pydantic_fields__', None): + parent_fields_lookup.update(model_fields) + + type_hints = _typing_extra.get_model_type_hints(cls, ns_resolver=ns_resolver) + + # https://docs.python.org/3/howto/annotations.html#accessing-the-annotations-dict-of-an-object-in-python-3-9-and-older + # annotations is only used for finding fields in parent classes + annotations = cls.__dict__.get('__annotations__', {}) + fields: dict[str, FieldInfo] = {} + + class_vars: set[str] = set() + for ann_name, (ann_type, evaluated) in type_hints.items(): + if ann_name == 'model_config': + # We never want to treat `model_config` as a field + # Note: we may need to change this logic if/when we introduce a `BareModel` class with no + # protected namespaces (where `model_config` might be allowed as a field name) + continue + + for protected_namespace in config_wrapper.protected_namespaces: + ns_violation: bool = False + if isinstance(protected_namespace, Pattern): + ns_violation = protected_namespace.match(ann_name) is not None + elif isinstance(protected_namespace, str): + ns_violation = ann_name.startswith(protected_namespace) + + if ns_violation: + for b in bases: + if hasattr(b, ann_name): + if not (issubclass(b, BaseModel) and ann_name in getattr(b, '__pydantic_fields__', {})): + raise NameError( + f'Field "{ann_name}" conflicts with member {getattr(b, ann_name)}' + f' of protected namespace "{protected_namespace}".' + ) + else: + valid_namespaces = () + for pn in config_wrapper.protected_namespaces: + if isinstance(pn, Pattern): + if not pn.match(ann_name): + valid_namespaces += (f're.compile({pn.pattern})',) + else: + if not ann_name.startswith(pn): + valid_namespaces += (pn,) + + warnings.warn( + f'Field "{ann_name}" in {cls.__name__} has conflict with protected namespace "{protected_namespace}".' + '\n\nYou may be able to resolve this warning by setting' + f" `model_config['protected_namespaces'] = {valid_namespaces}`.", + UserWarning, + ) + if _typing_extra.is_classvar_annotation(ann_type): + class_vars.add(ann_name) + continue + + assigned_value = getattr(cls, ann_name, PydanticUndefined) + + if not is_valid_field_name(ann_name): + continue + if cls.__pydantic_root_model__ and ann_name != 'root': + raise NameError( + f"Unexpected field with name {ann_name!r}; only 'root' is allowed as a field of a `RootModel`" + ) + + # when building a generic model with `MyModel[int]`, the generic_origin check makes sure we don't get + # "... shadows an attribute" warnings + generic_origin = getattr(cls, '__pydantic_generic_metadata__', {}).get('origin') + for base in bases: + dataclass_fields = { + field.name for field in (dataclasses.fields(base) if dataclasses.is_dataclass(base) else ()) + } + if hasattr(base, ann_name): + if base is generic_origin: + # Don't warn when "shadowing" of attributes in parametrized generics + continue + + if ann_name in dataclass_fields: + # Don't warn when inheriting stdlib dataclasses whose fields are "shadowed" by defaults being set + # on the class instance. + continue + + if ann_name not in annotations: + # Don't warn when a field exists in a parent class but has not been defined in the current class + continue + + warnings.warn( + f'Field name "{ann_name}" in "{cls.__qualname__}" shadows an attribute in parent ' + f'"{base.__qualname__}"', + UserWarning, + ) + + if assigned_value is PydanticUndefined: # no assignment, just a plain annotation + if ann_name in annotations or ann_name not in parent_fields_lookup: + # field is either: + # - present in the current model's annotations (and *not* from parent classes) + # - not found on any base classes; this seems to be caused by fields bot getting + # generated due to models not being fully defined while initializing recursive models. + # Nothing stops us from just creating a `FieldInfo` for this type hint, so we do this. + field_info = FieldInfo_.from_annotation(ann_type, _source=AnnotationSource.CLASS) + if not evaluated: + field_info._complete = False + # Store the original annotation that should be used to rebuild + # the field info later: + field_info._original_annotation = ann_type + else: + # The field was present on one of the (possibly multiple) base classes + # copy the field to make sure typevar substitutions don't cause issues with the base classes + field_info = copy(parent_fields_lookup[ann_name]) + + else: # An assigned value is present (either the default value, or a `Field()` function) + _warn_on_nested_alias_in_annotation(ann_type, ann_name) + if isinstance(assigned_value, FieldInfo_) and ismethoddescriptor(assigned_value.default): + # `assigned_value` was fetched using `getattr`, which triggers a call to `__get__` + # for descriptors, so we do the same if the `= field(default=...)` form is used. + # Note that we only do this for method descriptors for now, we might want to + # extend this to any descriptor in the future (by simply checking for + # `hasattr(assigned_value.default, '__get__')`). + assigned_value.default = assigned_value.default.__get__(None, cls) + + # The `from_annotated_attribute()` call below mutates the assigned `Field()`, so make a copy: + original_assignment = ( + assigned_value._copy() if not evaluated and isinstance(assigned_value, FieldInfo_) else assigned_value + ) + + field_info = FieldInfo_.from_annotated_attribute(ann_type, assigned_value, _source=AnnotationSource.CLASS) + # Store the original annotation and assignment value that should be used to rebuild the field info later. + # Note that the assignment is always stored as the annotation might contain a type var that is later + # parameterized with an unknown forward reference (and we'll need it to rebuild the field info): + field_info._original_assignment = original_assignment + if not evaluated: + field_info._complete = False + field_info._original_annotation = ann_type + elif 'final' in field_info._qualifiers and not field_info.is_required(): + warnings.warn( + f'Annotation {ann_name!r} is marked as final and has a default value. Pydantic treats {ann_name!r} as a ' + 'class variable, but it will be considered as a normal field in V3 to be aligned with dataclasses. If you ' + f'still want {ann_name!r} to be considered as a class variable, annotate it as: `ClassVar[] = .`', + category=PydanticDeprecatedSince211, + # Incorrect when `create_model` is used, but the chance that final with a default is used is low in that case: + stacklevel=4, + ) + class_vars.add(ann_name) + continue + + # attributes which are fields are removed from the class namespace: + # 1. To match the behaviour of annotation-only fields + # 2. To avoid false positives in the NameError check above + try: + delattr(cls, ann_name) + except AttributeError: + pass # indicates the attribute was on a parent class + + # Use cls.__dict__['__pydantic_decorators__'] instead of cls.__pydantic_decorators__ + # to make sure the decorators have already been built for this exact class + decorators: DecoratorInfos = cls.__dict__['__pydantic_decorators__'] + if ann_name in decorators.computed_fields: + raise TypeError( + f'Field {ann_name!r} of class {cls.__name__!r} overrides symbol of same name in a parent class. ' + 'This override with a computed_field is incompatible.' + ) + fields[ann_name] = field_info + + if typevars_map: + for field in fields.values(): + if field._complete: + field.apply_typevars_map(typevars_map) + + if config_wrapper.use_attribute_docstrings: + _update_fields_from_docstrings(cls, fields) + return fields, class_vars + + +def _warn_on_nested_alias_in_annotation(ann_type: type[Any], ann_name: str) -> None: + FieldInfo = import_cached_field_info() + + args = getattr(ann_type, '__args__', None) + if args: + for anno_arg in args: + if typing_objects.is_annotated(get_origin(anno_arg)): + for anno_type_arg in _typing_extra.get_args(anno_arg): + if isinstance(anno_type_arg, FieldInfo) and anno_type_arg.alias is not None: + warnings.warn( + f'`alias` specification on field "{ann_name}" must be set on outermost annotation to take effect.', + UserWarning, + ) + return + + +def rebuild_model_fields( + cls: type[BaseModel], + *, + ns_resolver: NsResolver, + typevars_map: Mapping[TypeVar, Any], +) -> dict[str, FieldInfo]: + """Rebuild the (already present) model fields by trying to reevaluate annotations. + + This function should be called whenever a model with incomplete fields is encountered. + + Raises: + NameError: If one of the annotations failed to evaluate. + + Note: + This function *doesn't* mutate the model fields in place, as it can be called during + schema generation, where you don't want to mutate other model's fields. + """ + FieldInfo_ = import_cached_field_info() + + rebuilt_fields: dict[str, FieldInfo] = {} + with ns_resolver.push(cls): + for f_name, field_info in cls.__pydantic_fields__.items(): + if field_info._complete: + rebuilt_fields[f_name] = field_info + else: + existing_desc = field_info.description + ann = _typing_extra.eval_type( + field_info._original_annotation, + *ns_resolver.types_namespace, + ) + ann = _generics.replace_types(ann, typevars_map) + + if (assign := field_info._original_assignment) is PydanticUndefined: + new_field = FieldInfo_.from_annotation(ann, _source=AnnotationSource.CLASS) + else: + new_field = FieldInfo_.from_annotated_attribute(ann, assign, _source=AnnotationSource.CLASS) + # The description might come from the docstring if `use_attribute_docstrings` was `True`: + new_field.description = new_field.description if new_field.description is not None else existing_desc + rebuilt_fields[f_name] = new_field + + return rebuilt_fields + + +def collect_dataclass_fields( + cls: type[StandardDataclass], + *, + ns_resolver: NsResolver | None = None, + typevars_map: dict[Any, Any] | None = None, + config_wrapper: ConfigWrapper | None = None, +) -> dict[str, FieldInfo]: + """Collect the fields of a dataclass. + + Args: + cls: dataclass. + ns_resolver: Namespace resolver to use when getting dataclass annotations. + Defaults to an empty instance. + typevars_map: A dictionary mapping type variables to their concrete types. + config_wrapper: The config wrapper instance. + + Returns: + The dataclass fields. + """ + FieldInfo_ = import_cached_field_info() + + fields: dict[str, FieldInfo] = {} + ns_resolver = ns_resolver or NsResolver() + dataclass_fields = cls.__dataclass_fields__ + + # The logic here is similar to `_typing_extra.get_cls_type_hints`, + # although we do it manually as stdlib dataclasses already have annotations + # collected in each class: + for base in reversed(cls.__mro__): + if not dataclasses.is_dataclass(base): + continue + + with ns_resolver.push(base): + for ann_name, dataclass_field in dataclass_fields.items(): + if ann_name not in base.__dict__.get('__annotations__', {}): + # `__dataclass_fields__`contains every field, even the ones from base classes. + # Only collect the ones defined on `base`. + continue + + globalns, localns = ns_resolver.types_namespace + ann_type, evaluated = _typing_extra.try_eval_type(dataclass_field.type, globalns, localns) + + if _typing_extra.is_classvar_annotation(ann_type): + continue + + if ( + not dataclass_field.init + and dataclass_field.default is dataclasses.MISSING + and dataclass_field.default_factory is dataclasses.MISSING + ): + # TODO: We should probably do something with this so that validate_assignment behaves properly + # Issue: https://github.com/pydantic/pydantic/issues/5470 + continue + + if isinstance(dataclass_field.default, FieldInfo_): + if dataclass_field.default.init_var: + if dataclass_field.default.init is False: + raise PydanticUserError( + f'Dataclass field {ann_name} has init=False and init_var=True, but these are mutually exclusive.', + code='clashing-init-and-init-var', + ) + + # TODO: same note as above re validate_assignment + continue + field_info = FieldInfo_.from_annotated_attribute( + ann_type, dataclass_field.default, _source=AnnotationSource.DATACLASS + ) + field_info._original_assignment = dataclass_field.default + else: + field_info = FieldInfo_.from_annotated_attribute( + ann_type, dataclass_field, _source=AnnotationSource.DATACLASS + ) + field_info._original_assignment = dataclass_field + + if not evaluated: + field_info._complete = False + field_info._original_annotation = ann_type + + fields[ann_name] = field_info + + if field_info.default is not PydanticUndefined and isinstance( + getattr(cls, ann_name, field_info), FieldInfo_ + ): + # We need this to fix the default when the "default" from __dataclass_fields__ is a pydantic.FieldInfo + setattr(cls, ann_name, field_info.default) + + if typevars_map: + for field in fields.values(): + # We don't pass any ns, as `field.annotation` + # was already evaluated. TODO: is this method relevant? + # Can't we juste use `_generics.replace_types`? + field.apply_typevars_map(typevars_map) + + if config_wrapper is not None and config_wrapper.use_attribute_docstrings: + _update_fields_from_docstrings( + cls, + fields, + # We can't rely on the (more reliable) frame inspection method + # for stdlib dataclasses: + use_inspect=not hasattr(cls, '__is_pydantic_dataclass__'), + ) + + return fields + + +def rebuild_dataclass_fields( + cls: type[PydanticDataclass], + *, + config_wrapper: ConfigWrapper, + ns_resolver: NsResolver, + typevars_map: Mapping[TypeVar, Any], +) -> dict[str, FieldInfo]: + """Rebuild the (already present) dataclass fields by trying to reevaluate annotations. + + This function should be called whenever a dataclass with incomplete fields is encountered. + + Raises: + NameError: If one of the annotations failed to evaluate. + + Note: + This function *doesn't* mutate the dataclass fields in place, as it can be called during + schema generation, where you don't want to mutate other dataclass's fields. + """ + FieldInfo_ = import_cached_field_info() + + rebuilt_fields: dict[str, FieldInfo] = {} + with ns_resolver.push(cls): + for f_name, field_info in cls.__pydantic_fields__.items(): + if field_info._complete: + rebuilt_fields[f_name] = field_info + else: + existing_desc = field_info.description + ann = _typing_extra.eval_type( + field_info._original_annotation, + *ns_resolver.types_namespace, + ) + ann = _generics.replace_types(ann, typevars_map) + new_field = FieldInfo_.from_annotated_attribute( + ann, + field_info._original_assignment, + _source=AnnotationSource.DATACLASS, + ) + + # The description might come from the docstring if `use_attribute_docstrings` was `True`: + new_field.description = new_field.description if new_field.description is not None else existing_desc + rebuilt_fields[f_name] = new_field + + return rebuilt_fields + + +def is_valid_field_name(name: str) -> bool: + return not name.startswith('_') + + +def is_valid_privateattr_name(name: str) -> bool: + return name.startswith('_') and not name.startswith('__') + + +def takes_validated_data_argument( + default_factory: Callable[[], Any] | Callable[[dict[str, Any]], Any], +) -> TypeIs[Callable[[dict[str, Any]], Any]]: + """Whether the provided default factory callable has a validated data parameter.""" + try: + sig = signature(default_factory) + except (ValueError, TypeError): + # `inspect.signature` might not be able to infer a signature, e.g. with C objects. + # In this case, we assume no data argument is present: + return False + + parameters = list(sig.parameters.values()) + + return len(parameters) == 1 and can_be_positional(parameters[0]) and parameters[0].default is Parameter.empty diff --git a/lib/python3.12/site-packages/pydantic/_internal/_forward_ref.py b/lib/python3.12/site-packages/pydantic/_internal/_forward_ref.py new file mode 100644 index 0000000000000000000000000000000000000000..231f81d11b87ef05d4765cd61d88d7348eb9e401 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_forward_ref.py @@ -0,0 +1,23 @@ +from __future__ import annotations as _annotations + +from dataclasses import dataclass +from typing import Union + + +@dataclass +class PydanticRecursiveRef: + type_ref: str + + __name__ = 'PydanticRecursiveRef' + __hash__ = object.__hash__ + + def __call__(self) -> None: + """Defining __call__ is necessary for the `typing` module to let you use an instance of + this class as the result of resolving a standard ForwardRef. + """ + + def __or__(self, other): + return Union[self, other] # type: ignore + + def __ror__(self, other): + return Union[other, self] # type: ignore diff --git a/lib/python3.12/site-packages/pydantic/_internal/_generate_schema.py b/lib/python3.12/site-packages/pydantic/_internal/_generate_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..04512289aced75c1923d2eeefa55b9b8b51ede0a --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_generate_schema.py @@ -0,0 +1,2904 @@ +"""Convert python types to pydantic-core schema.""" + +from __future__ import annotations as _annotations + +import collections.abc +import dataclasses +import datetime +import inspect +import os +import pathlib +import re +import sys +import typing +import warnings +from collections.abc import Generator, Iterable, Iterator, Mapping +from contextlib import contextmanager +from copy import copy +from decimal import Decimal +from enum import Enum +from fractions import Fraction +from functools import partial +from inspect import Parameter, _ParameterKind, signature +from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network +from itertools import chain +from operator import attrgetter +from types import FunctionType, GenericAlias, LambdaType, MethodType +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Final, + ForwardRef, + Literal, + TypeVar, + Union, + cast, + overload, +) +from uuid import UUID +from warnings import warn +from zoneinfo import ZoneInfo + +import typing_extensions +from pydantic_core import ( + CoreSchema, + MultiHostUrl, + PydanticCustomError, + PydanticSerializationUnexpectedValue, + PydanticUndefined, + Url, + core_schema, + to_jsonable_python, +) +from typing_extensions import TypeAlias, TypeAliasType, TypedDict, get_args, get_origin, is_typeddict +from typing_inspection import typing_objects +from typing_inspection.introspection import AnnotationSource, get_literal_values, is_union_origin + +from ..aliases import AliasChoices, AliasGenerator, AliasPath +from ..annotated_handlers import GetCoreSchemaHandler, GetJsonSchemaHandler +from ..config import ConfigDict, JsonDict, JsonEncoder, JsonSchemaExtraCallable +from ..errors import PydanticSchemaGenerationError, PydanticUndefinedAnnotation, PydanticUserError +from ..functional_validators import AfterValidator, BeforeValidator, FieldValidatorModes, PlainValidator, WrapValidator +from ..json_schema import JsonSchemaValue +from ..version import version_short +from ..warnings import PydanticDeprecatedSince20 +from . import _decorators, _discriminated_union, _known_annotated_metadata, _repr, _typing_extra +from ._config import ConfigWrapper, ConfigWrapperStack +from ._core_metadata import CoreMetadata, update_core_metadata +from ._core_utils import ( + get_ref, + get_type_ref, + is_list_like_schema_with_items_schema, + validate_core_schema, +) +from ._decorators import ( + Decorator, + DecoratorInfos, + FieldSerializerDecoratorInfo, + FieldValidatorDecoratorInfo, + ModelSerializerDecoratorInfo, + ModelValidatorDecoratorInfo, + RootValidatorDecoratorInfo, + ValidatorDecoratorInfo, + get_attribute_from_bases, + inspect_field_serializer, + inspect_model_serializer, + inspect_validator, +) +from ._docs_extraction import extract_docstrings_from_cls +from ._fields import ( + collect_dataclass_fields, + rebuild_dataclass_fields, + rebuild_model_fields, + takes_validated_data_argument, +) +from ._forward_ref import PydanticRecursiveRef +from ._generics import get_standard_typevars_map, replace_types +from ._import_utils import import_cached_base_model, import_cached_field_info +from ._mock_val_ser import MockCoreSchema +from ._namespace_utils import NamespacesTuple, NsResolver +from ._schema_gather import MissingDefinitionError, gather_schemas_for_cleaning +from ._schema_generation_shared import CallbackGetCoreSchemaHandler +from ._utils import lenient_issubclass, smart_deepcopy + +if TYPE_CHECKING: + from ..fields import ComputedFieldInfo, FieldInfo + from ..main import BaseModel + from ..types import Discriminator + from ._dataclasses import StandardDataclass + from ._schema_generation_shared import GetJsonSchemaFunction + +_SUPPORTS_TYPEDDICT = sys.version_info >= (3, 12) + +FieldDecoratorInfo = Union[ValidatorDecoratorInfo, FieldValidatorDecoratorInfo, FieldSerializerDecoratorInfo] +FieldDecoratorInfoType = TypeVar('FieldDecoratorInfoType', bound=FieldDecoratorInfo) +AnyFieldDecorator = Union[ + Decorator[ValidatorDecoratorInfo], + Decorator[FieldValidatorDecoratorInfo], + Decorator[FieldSerializerDecoratorInfo], +] + +ModifyCoreSchemaWrapHandler: TypeAlias = GetCoreSchemaHandler +GetCoreSchemaFunction: TypeAlias = Callable[[Any, ModifyCoreSchemaWrapHandler], core_schema.CoreSchema] +ParametersCallback: TypeAlias = "Callable[[int, str, Any], Literal['skip'] | None]" + +TUPLE_TYPES: list[type] = [typing.Tuple, tuple] # noqa: UP006 +LIST_TYPES: list[type] = [typing.List, list, collections.abc.MutableSequence] # noqa: UP006 +SET_TYPES: list[type] = [typing.Set, set, collections.abc.MutableSet] # noqa: UP006 +FROZEN_SET_TYPES: list[type] = [typing.FrozenSet, frozenset, collections.abc.Set] # noqa: UP006 +DICT_TYPES: list[type] = [typing.Dict, dict] # noqa: UP006 +IP_TYPES: list[type] = [IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network] +SEQUENCE_TYPES: list[type] = [typing.Sequence, collections.abc.Sequence] +ITERABLE_TYPES: list[type] = [typing.Iterable, collections.abc.Iterable, typing.Generator, collections.abc.Generator] +TYPE_TYPES: list[type] = [typing.Type, type] # noqa: UP006 +PATTERN_TYPES: list[type] = [typing.Pattern, re.Pattern] +PATH_TYPES: list[type] = [ + os.PathLike, + pathlib.Path, + pathlib.PurePath, + pathlib.PosixPath, + pathlib.PurePosixPath, + pathlib.PureWindowsPath, +] +MAPPING_TYPES = [ + typing.Mapping, + typing.MutableMapping, + collections.abc.Mapping, + collections.abc.MutableMapping, + collections.OrderedDict, + typing_extensions.OrderedDict, + typing.DefaultDict, # noqa: UP006 + collections.defaultdict, +] +COUNTER_TYPES = [collections.Counter, typing.Counter] +DEQUE_TYPES: list[type] = [collections.deque, typing.Deque] # noqa: UP006 + +# Note: This does not play very well with type checkers. For example, +# `a: LambdaType = lambda x: x` will raise a type error by Pyright. +ValidateCallSupportedTypes = Union[ + LambdaType, + FunctionType, + MethodType, + partial, +] + +VALIDATE_CALL_SUPPORTED_TYPES = get_args(ValidateCallSupportedTypes) + +_mode_to_validator: dict[ + FieldValidatorModes, type[BeforeValidator | AfterValidator | PlainValidator | WrapValidator] +] = {'before': BeforeValidator, 'after': AfterValidator, 'plain': PlainValidator, 'wrap': WrapValidator} + + +def check_validator_fields_against_field_name( + info: FieldDecoratorInfo, + field: str, +) -> bool: + """Check if field name is in validator fields. + + Args: + info: The field info. + field: The field name to check. + + Returns: + `True` if field name is in validator fields, `False` otherwise. + """ + fields = info.fields + return '*' in fields or field in fields + + +def check_decorator_fields_exist(decorators: Iterable[AnyFieldDecorator], fields: Iterable[str]) -> None: + """Check if the defined fields in decorators exist in `fields` param. + + It ignores the check for a decorator if the decorator has `*` as field or `check_fields=False`. + + Args: + decorators: An iterable of decorators. + fields: An iterable of fields name. + + Raises: + PydanticUserError: If one of the field names does not exist in `fields` param. + """ + fields = set(fields) + for dec in decorators: + if '*' in dec.info.fields: + continue + if dec.info.check_fields is False: + continue + for field in dec.info.fields: + if field not in fields: + raise PydanticUserError( + f'Decorators defined with incorrect fields: {dec.cls_ref}.{dec.cls_var_name}' + " (use check_fields=False if you're inheriting from the model and intended this)", + code='decorator-missing-field', + ) + + +def filter_field_decorator_info_by_field( + validator_functions: Iterable[Decorator[FieldDecoratorInfoType]], field: str +) -> list[Decorator[FieldDecoratorInfoType]]: + return [dec for dec in validator_functions if check_validator_fields_against_field_name(dec.info, field)] + + +def apply_each_item_validators( + schema: core_schema.CoreSchema, + each_item_validators: list[Decorator[ValidatorDecoratorInfo]], + field_name: str | None, +) -> core_schema.CoreSchema: + # This V1 compatibility shim should eventually be removed + + # fail early if each_item_validators is empty + if not each_item_validators: + return schema + + # push down any `each_item=True` validators + # note that this won't work for any Annotated types that get wrapped by a function validator + # but that's okay because that didn't exist in V1 + if schema['type'] == 'nullable': + schema['schema'] = apply_each_item_validators(schema['schema'], each_item_validators, field_name) + return schema + elif schema['type'] == 'tuple': + if (variadic_item_index := schema.get('variadic_item_index')) is not None: + schema['items_schema'][variadic_item_index] = apply_validators( + schema['items_schema'][variadic_item_index], + each_item_validators, + field_name, + ) + elif is_list_like_schema_with_items_schema(schema): + inner_schema = schema.get('items_schema', core_schema.any_schema()) + schema['items_schema'] = apply_validators(inner_schema, each_item_validators, field_name) + elif schema['type'] == 'dict': + inner_schema = schema.get('values_schema', core_schema.any_schema()) + schema['values_schema'] = apply_validators(inner_schema, each_item_validators, field_name) + else: + raise TypeError( + f'`@validator(..., each_item=True)` cannot be applied to fields with a schema of {schema["type"]}' + ) + return schema + + +def _extract_json_schema_info_from_field_info( + info: FieldInfo | ComputedFieldInfo, +) -> tuple[JsonDict | None, JsonDict | JsonSchemaExtraCallable | None]: + json_schema_updates = { + 'title': info.title, + 'description': info.description, + 'deprecated': bool(info.deprecated) or info.deprecated == '' or None, + 'examples': to_jsonable_python(info.examples), + } + json_schema_updates = {k: v for k, v in json_schema_updates.items() if v is not None} + return (json_schema_updates or None, info.json_schema_extra) + + +JsonEncoders = dict[type[Any], JsonEncoder] + + +def _add_custom_serialization_from_json_encoders( + json_encoders: JsonEncoders | None, tp: Any, schema: CoreSchema +) -> CoreSchema: + """Iterate over the json_encoders and add the first matching encoder to the schema. + + Args: + json_encoders: A dictionary of types and their encoder functions. + tp: The type to check for a matching encoder. + schema: The schema to add the encoder to. + """ + if not json_encoders: + return schema + if 'serialization' in schema: + return schema + # Check the class type and its superclasses for a matching encoder + # Decimal.__class__.__mro__ (and probably other cases) doesn't include Decimal itself + # if the type is a GenericAlias (e.g. from list[int]) we need to use __class__ instead of .__mro__ + for base in (tp, *getattr(tp, '__mro__', tp.__class__.__mro__)[:-1]): + encoder = json_encoders.get(base) + if encoder is None: + continue + + warnings.warn( + f'`json_encoders` is deprecated. See https://docs.pydantic.dev/{version_short()}/concepts/serialization/#custom-serializers for alternatives', + PydanticDeprecatedSince20, + ) + + # TODO: in theory we should check that the schema accepts a serialization key + schema['serialization'] = core_schema.plain_serializer_function_ser_schema(encoder, when_used='json') + return schema + + return schema + + +def _get_first_non_null(a: Any, b: Any) -> Any: + """Return the first argument if it is not None, otherwise return the second argument. + + Use case: serialization_alias (argument a) and alias (argument b) are both defined, and serialization_alias is ''. + This function will return serialization_alias, which is the first argument, even though it is an empty string. + """ + return a if a is not None else b + + +class InvalidSchemaError(Exception): + """The core schema is invalid.""" + + +class GenerateSchema: + """Generate core schema for a Pydantic model, dataclass and types like `str`, `datetime`, ... .""" + + __slots__ = ( + '_config_wrapper_stack', + '_ns_resolver', + '_typevars_map', + 'field_name_stack', + 'model_type_stack', + 'defs', + ) + + def __init__( + self, + config_wrapper: ConfigWrapper, + ns_resolver: NsResolver | None = None, + typevars_map: Mapping[TypeVar, Any] | None = None, + ) -> None: + # we need a stack for recursing into nested models + self._config_wrapper_stack = ConfigWrapperStack(config_wrapper) + self._ns_resolver = ns_resolver or NsResolver() + self._typevars_map = typevars_map + self.field_name_stack = _FieldNameStack() + self.model_type_stack = _ModelTypeStack() + self.defs = _Definitions() + + def __init_subclass__(cls) -> None: + super().__init_subclass__() + warnings.warn( + 'Subclassing `GenerateSchema` is not supported. The API is highly subject to change in minor versions.', + UserWarning, + stacklevel=2, + ) + + @property + def _config_wrapper(self) -> ConfigWrapper: + return self._config_wrapper_stack.tail + + @property + def _types_namespace(self) -> NamespacesTuple: + return self._ns_resolver.types_namespace + + @property + def _arbitrary_types(self) -> bool: + return self._config_wrapper.arbitrary_types_allowed + + # the following methods can be overridden but should be considered + # unstable / private APIs + def _list_schema(self, items_type: Any) -> CoreSchema: + return core_schema.list_schema(self.generate_schema(items_type)) + + def _dict_schema(self, keys_type: Any, values_type: Any) -> CoreSchema: + return core_schema.dict_schema(self.generate_schema(keys_type), self.generate_schema(values_type)) + + def _set_schema(self, items_type: Any) -> CoreSchema: + return core_schema.set_schema(self.generate_schema(items_type)) + + def _frozenset_schema(self, items_type: Any) -> CoreSchema: + return core_schema.frozenset_schema(self.generate_schema(items_type)) + + def _enum_schema(self, enum_type: type[Enum]) -> CoreSchema: + cases: list[Any] = list(enum_type.__members__.values()) + + enum_ref = get_type_ref(enum_type) + description = None if not enum_type.__doc__ else inspect.cleandoc(enum_type.__doc__) + if ( + description == 'An enumeration.' + ): # This is the default value provided by enum.EnumMeta.__new__; don't use it + description = None + js_updates = {'title': enum_type.__name__, 'description': description} + js_updates = {k: v for k, v in js_updates.items() if v is not None} + + sub_type: Literal['str', 'int', 'float'] | None = None + if issubclass(enum_type, int): + sub_type = 'int' + value_ser_type: core_schema.SerSchema = core_schema.simple_ser_schema('int') + elif issubclass(enum_type, str): + # this handles `StrEnum` (3.11 only), and also `Foobar(str, Enum)` + sub_type = 'str' + value_ser_type = core_schema.simple_ser_schema('str') + elif issubclass(enum_type, float): + sub_type = 'float' + value_ser_type = core_schema.simple_ser_schema('float') + else: + # TODO this is an ugly hack, how do we trigger an Any schema for serialization? + value_ser_type = core_schema.plain_serializer_function_ser_schema(lambda x: x) + + if cases: + + def get_json_schema(schema: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue: + json_schema = handler(schema) + original_schema = handler.resolve_ref_schema(json_schema) + original_schema.update(js_updates) + return json_schema + + # we don't want to add the missing to the schema if it's the default one + default_missing = getattr(enum_type._missing_, '__func__', None) is Enum._missing_.__func__ # pyright: ignore[reportFunctionMemberAccess] + enum_schema = core_schema.enum_schema( + enum_type, + cases, + sub_type=sub_type, + missing=None if default_missing else enum_type._missing_, + ref=enum_ref, + metadata={'pydantic_js_functions': [get_json_schema]}, + ) + + if self._config_wrapper.use_enum_values: + enum_schema = core_schema.no_info_after_validator_function( + attrgetter('value'), enum_schema, serialization=value_ser_type + ) + + return enum_schema + + else: + + def get_json_schema_no_cases(_, handler: GetJsonSchemaHandler) -> JsonSchemaValue: + json_schema = handler(core_schema.enum_schema(enum_type, cases, sub_type=sub_type, ref=enum_ref)) + original_schema = handler.resolve_ref_schema(json_schema) + original_schema.update(js_updates) + return json_schema + + # Use an isinstance check for enums with no cases. + # The most important use case for this is creating TypeVar bounds for generics that should + # be restricted to enums. This is more consistent than it might seem at first, since you can only + # subclass enum.Enum (or subclasses of enum.Enum) if all parent classes have no cases. + # We use the get_json_schema function when an Enum subclass has been declared with no cases + # so that we can still generate a valid json schema. + return core_schema.is_instance_schema( + enum_type, + metadata={'pydantic_js_functions': [get_json_schema_no_cases]}, + ) + + def _ip_schema(self, tp: Any) -> CoreSchema: + from ._validators import IP_VALIDATOR_LOOKUP, IpType + + ip_type_json_schema_format: dict[type[IpType], str] = { + IPv4Address: 'ipv4', + IPv4Network: 'ipv4network', + IPv4Interface: 'ipv4interface', + IPv6Address: 'ipv6', + IPv6Network: 'ipv6network', + IPv6Interface: 'ipv6interface', + } + + def ser_ip(ip: Any, info: core_schema.SerializationInfo) -> str | IpType: + if not isinstance(ip, (tp, str)): + raise PydanticSerializationUnexpectedValue( + f"Expected `{tp}` but got `{type(ip)}` with value `'{ip}'` - serialized value may not be as expected." + ) + if info.mode == 'python': + return ip + return str(ip) + + return core_schema.lax_or_strict_schema( + lax_schema=core_schema.no_info_plain_validator_function(IP_VALIDATOR_LOOKUP[tp]), + strict_schema=core_schema.json_or_python_schema( + json_schema=core_schema.no_info_after_validator_function(tp, core_schema.str_schema()), + python_schema=core_schema.is_instance_schema(tp), + ), + serialization=core_schema.plain_serializer_function_ser_schema(ser_ip, info_arg=True, when_used='always'), + metadata={ + 'pydantic_js_functions': [lambda _1, _2: {'type': 'string', 'format': ip_type_json_schema_format[tp]}] + }, + ) + + def _path_schema(self, tp: Any, path_type: Any) -> CoreSchema: + if tp is os.PathLike and (path_type not in {str, bytes} and not typing_objects.is_any(path_type)): + raise PydanticUserError( + '`os.PathLike` can only be used with `str`, `bytes` or `Any`', code='schema-for-unknown-type' + ) + + path_constructor = pathlib.PurePath if tp is os.PathLike else tp + strict_inner_schema = ( + core_schema.bytes_schema(strict=True) if (path_type is bytes) else core_schema.str_schema(strict=True) + ) + lax_inner_schema = core_schema.bytes_schema() if (path_type is bytes) else core_schema.str_schema() + + def path_validator(input_value: str | bytes) -> os.PathLike[Any]: # type: ignore + try: + if path_type is bytes: + if isinstance(input_value, bytes): + try: + input_value = input_value.decode() + except UnicodeDecodeError as e: + raise PydanticCustomError('bytes_type', 'Input must be valid bytes') from e + else: + raise PydanticCustomError('bytes_type', 'Input must be bytes') + elif not isinstance(input_value, str): + raise PydanticCustomError('path_type', 'Input is not a valid path') + + return path_constructor(input_value) # type: ignore + except TypeError as e: + raise PydanticCustomError('path_type', 'Input is not a valid path') from e + + def ser_path(path: Any, info: core_schema.SerializationInfo) -> str | os.PathLike[Any]: + if not isinstance(path, (tp, str)): + raise PydanticSerializationUnexpectedValue( + f"Expected `{tp}` but got `{type(path)}` with value `'{path}'` - serialized value may not be as expected." + ) + if info.mode == 'python': + return path + return str(path) + + instance_schema = core_schema.json_or_python_schema( + json_schema=core_schema.no_info_after_validator_function(path_validator, lax_inner_schema), + python_schema=core_schema.is_instance_schema(tp), + ) + + schema = core_schema.lax_or_strict_schema( + lax_schema=core_schema.union_schema( + [ + instance_schema, + core_schema.no_info_after_validator_function(path_validator, strict_inner_schema), + ], + custom_error_type='path_type', + custom_error_message=f'Input is not a valid path for {tp}', + ), + strict_schema=instance_schema, + serialization=core_schema.plain_serializer_function_ser_schema(ser_path, info_arg=True, when_used='always'), + metadata={'pydantic_js_functions': [lambda source, handler: {**handler(source), 'format': 'path'}]}, + ) + return schema + + def _deque_schema(self, items_type: Any) -> CoreSchema: + from ._serializers import serialize_sequence_via_list + from ._validators import deque_validator + + item_type_schema = self.generate_schema(items_type) + + # we have to use a lax list schema here, because we need to validate the deque's + # items via a list schema, but it's ok if the deque itself is not a list + list_schema = core_schema.list_schema(item_type_schema, strict=False) + + check_instance = core_schema.json_or_python_schema( + json_schema=list_schema, + python_schema=core_schema.is_instance_schema(collections.deque, cls_repr='Deque'), + ) + + lax_schema = core_schema.no_info_wrap_validator_function(deque_validator, list_schema) + + return core_schema.lax_or_strict_schema( + lax_schema=lax_schema, + strict_schema=core_schema.chain_schema([check_instance, lax_schema]), + serialization=core_schema.wrap_serializer_function_ser_schema( + serialize_sequence_via_list, schema=item_type_schema, info_arg=True + ), + ) + + def _mapping_schema(self, tp: Any, keys_type: Any, values_type: Any) -> CoreSchema: + from ._validators import MAPPING_ORIGIN_MAP, defaultdict_validator, get_defaultdict_default_default_factory + + mapped_origin = MAPPING_ORIGIN_MAP[tp] + keys_schema = self.generate_schema(keys_type) + values_schema = self.generate_schema(values_type) + dict_schema = core_schema.dict_schema(keys_schema, values_schema, strict=False) + + if mapped_origin is dict: + schema = dict_schema + else: + check_instance = core_schema.json_or_python_schema( + json_schema=dict_schema, + python_schema=core_schema.is_instance_schema(mapped_origin), + ) + + if tp is collections.defaultdict: + default_default_factory = get_defaultdict_default_default_factory(values_type) + coerce_instance_wrap = partial( + core_schema.no_info_wrap_validator_function, + partial(defaultdict_validator, default_default_factory=default_default_factory), + ) + else: + coerce_instance_wrap = partial(core_schema.no_info_after_validator_function, mapped_origin) + + lax_schema = coerce_instance_wrap(dict_schema) + strict_schema = core_schema.chain_schema([check_instance, lax_schema]) + + schema = core_schema.lax_or_strict_schema( + lax_schema=lax_schema, + strict_schema=strict_schema, + serialization=core_schema.wrap_serializer_function_ser_schema( + lambda v, h: h(v), schema=dict_schema, info_arg=False + ), + ) + + return schema + + def _fraction_schema(self) -> CoreSchema: + """Support for [`fractions.Fraction`][fractions.Fraction].""" + from ._validators import fraction_validator + + # TODO: note, this is a fairly common pattern, re lax / strict for attempted type coercion, + # can we use a helper function to reduce boilerplate? + return core_schema.lax_or_strict_schema( + lax_schema=core_schema.no_info_plain_validator_function(fraction_validator), + strict_schema=core_schema.json_or_python_schema( + json_schema=core_schema.no_info_plain_validator_function(fraction_validator), + python_schema=core_schema.is_instance_schema(Fraction), + ), + # use str serialization to guarantee round trip behavior + serialization=core_schema.to_string_ser_schema(when_used='always'), + metadata={'pydantic_js_functions': [lambda _1, _2: {'type': 'string', 'format': 'fraction'}]}, + ) + + def _arbitrary_type_schema(self, tp: Any) -> CoreSchema: + if not isinstance(tp, type): + warn( + f'{tp!r} is not a Python type (it may be an instance of an object),' + ' Pydantic will allow any object with no validation since we cannot even' + ' enforce that the input is an instance of the given type.' + ' To get rid of this error wrap the type with `pydantic.SkipValidation`.', + UserWarning, + ) + return core_schema.any_schema() + return core_schema.is_instance_schema(tp) + + def _unknown_type_schema(self, obj: Any) -> CoreSchema: + raise PydanticSchemaGenerationError( + f'Unable to generate pydantic-core schema for {obj!r}. ' + 'Set `arbitrary_types_allowed=True` in the model_config to ignore this error' + ' or implement `__get_pydantic_core_schema__` on your type to fully support it.' + '\n\nIf you got this error by calling handler() within' + ' `__get_pydantic_core_schema__` then you likely need to call' + ' `handler.generate_schema()` since we do not call' + ' `__get_pydantic_core_schema__` on `` otherwise to avoid infinite recursion.' + ) + + def _apply_discriminator_to_union( + self, schema: CoreSchema, discriminator: str | Discriminator | None + ) -> CoreSchema: + if discriminator is None: + return schema + try: + return _discriminated_union.apply_discriminator( + schema, + discriminator, + self.defs._definitions, + ) + except _discriminated_union.MissingDefinitionForUnionRef: + # defer until defs are resolved + _discriminated_union.set_discriminator_in_metadata( + schema, + discriminator, + ) + return schema + + def clean_schema(self, schema: CoreSchema) -> CoreSchema: + schema = self.defs.finalize_schema(schema) + schema = validate_core_schema(schema) + return schema + + def _add_js_function(self, metadata_schema: CoreSchema, js_function: Callable[..., Any]) -> None: + metadata = metadata_schema.get('metadata', {}) + pydantic_js_functions = metadata.setdefault('pydantic_js_functions', []) + # because of how we generate core schemas for nested generic models + # we can end up adding `BaseModel.__get_pydantic_json_schema__` multiple times + # this check may fail to catch duplicates if the function is a `functools.partial` + # or something like that, but if it does it'll fail by inserting the duplicate + if js_function not in pydantic_js_functions: + pydantic_js_functions.append(js_function) + metadata_schema['metadata'] = metadata + + def generate_schema( + self, + obj: Any, + ) -> core_schema.CoreSchema: + """Generate core schema. + + Args: + obj: The object to generate core schema for. + + Returns: + The generated core schema. + + Raises: + PydanticUndefinedAnnotation: + If it is not possible to evaluate forward reference. + PydanticSchemaGenerationError: + If it is not possible to generate pydantic-core schema. + TypeError: + - If `alias_generator` returns a disallowed type (must be str, AliasPath or AliasChoices). + - If V1 style validator with `each_item=True` applied on a wrong field. + PydanticUserError: + - If `typing.TypedDict` is used instead of `typing_extensions.TypedDict` on Python < 3.12. + - If `__modify_schema__` method is used instead of `__get_pydantic_json_schema__`. + """ + schema = self._generate_schema_from_get_schema_method(obj, obj) + + if schema is None: + schema = self._generate_schema_inner(obj) + + metadata_js_function = _extract_get_pydantic_json_schema(obj) + if metadata_js_function is not None: + metadata_schema = resolve_original_schema(schema, self.defs) + if metadata_schema: + self._add_js_function(metadata_schema, metadata_js_function) + + schema = _add_custom_serialization_from_json_encoders(self._config_wrapper.json_encoders, obj, schema) + + return schema + + def _model_schema(self, cls: type[BaseModel]) -> core_schema.CoreSchema: + """Generate schema for a Pydantic model.""" + BaseModel_ = import_cached_base_model() + + with self.defs.get_schema_or_ref(cls) as (model_ref, maybe_schema): + if maybe_schema is not None: + return maybe_schema + + schema = cls.__dict__.get('__pydantic_core_schema__') + if schema is not None and not isinstance(schema, MockCoreSchema): + if schema['type'] == 'definitions': + schema = self.defs.unpack_definitions(schema) + ref = get_ref(schema) + if ref: + return self.defs.create_definition_reference_schema(schema) + else: + return schema + + config_wrapper = ConfigWrapper(cls.model_config, check=False) + + with self._config_wrapper_stack.push(config_wrapper), self._ns_resolver.push(cls): + core_config = self._config_wrapper.core_config(title=cls.__name__) + + if cls.__pydantic_fields_complete__ or cls is BaseModel_: + fields = getattr(cls, '__pydantic_fields__', {}) + else: + if not hasattr(cls, '__pydantic_fields__'): + # This happens when we have a loop in the schema generation: + # class Base[T](BaseModel): + # t: T + # + # class Other(BaseModel): + # b: 'Base[Other]' + # When we build fields for `Other`, we evaluate the forward annotation. + # At this point, `Other` doesn't have the model fields set. We create + # `Base[Other]`; model fields are successfully built, and we try to generate + # a schema for `t: Other`. As `Other.__pydantic_fields__` aren't set, we abort. + raise PydanticUndefinedAnnotation( + name=cls.__name__, + message=f'Class {cls.__name__!r} is not defined', + ) + try: + fields = rebuild_model_fields( + cls, + ns_resolver=self._ns_resolver, + typevars_map=self._typevars_map or {}, + ) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + + decorators = cls.__pydantic_decorators__ + computed_fields = decorators.computed_fields + check_decorator_fields_exist( + chain( + decorators.field_validators.values(), + decorators.field_serializers.values(), + decorators.validators.values(), + ), + {*fields.keys(), *computed_fields.keys()}, + ) + + model_validators = decorators.model_validators.values() + + extras_schema = None + extras_keys_schema = None + if core_config.get('extra_fields_behavior') == 'allow': + assert cls.__mro__[0] is cls + assert cls.__mro__[-1] is object + for candidate_cls in cls.__mro__[:-1]: + extras_annotation = getattr(candidate_cls, '__annotations__', {}).get( + '__pydantic_extra__', None + ) + if extras_annotation is not None: + if isinstance(extras_annotation, str): + extras_annotation = _typing_extra.eval_type_backport( + _typing_extra._make_forward_ref( + extras_annotation, is_argument=False, is_class=True + ), + *self._types_namespace, + ) + tp = get_origin(extras_annotation) + if tp not in DICT_TYPES: + raise PydanticSchemaGenerationError( + 'The type annotation for `__pydantic_extra__` must be `dict[str, ...]`' + ) + extra_keys_type, extra_items_type = self._get_args_resolving_forward_refs( + extras_annotation, + required=True, + ) + if extra_keys_type is not str: + extras_keys_schema = self.generate_schema(extra_keys_type) + if not typing_objects.is_any(extra_items_type): + extras_schema = self.generate_schema(extra_items_type) + if extras_keys_schema is not None or extras_schema is not None: + break + + generic_origin: type[BaseModel] | None = getattr(cls, '__pydantic_generic_metadata__', {}).get('origin') + + if cls.__pydantic_root_model__: + root_field = self._common_field_schema('root', fields['root'], decorators) + inner_schema = root_field['schema'] + inner_schema = apply_model_validators(inner_schema, model_validators, 'inner') + model_schema = core_schema.model_schema( + cls, + inner_schema, + generic_origin=generic_origin, + custom_init=getattr(cls, '__pydantic_custom_init__', None), + root_model=True, + post_init=getattr(cls, '__pydantic_post_init__', None), + config=core_config, + ref=model_ref, + ) + else: + fields_schema: core_schema.CoreSchema = core_schema.model_fields_schema( + {k: self._generate_md_field_schema(k, v, decorators) for k, v in fields.items()}, + computed_fields=[ + self._computed_field_schema(d, decorators.field_serializers) + for d in computed_fields.values() + ], + extras_schema=extras_schema, + extras_keys_schema=extras_keys_schema, + model_name=cls.__name__, + ) + inner_schema = apply_validators(fields_schema, decorators.root_validators.values(), None) + inner_schema = apply_model_validators(inner_schema, model_validators, 'inner') + + model_schema = core_schema.model_schema( + cls, + inner_schema, + generic_origin=generic_origin, + custom_init=getattr(cls, '__pydantic_custom_init__', None), + root_model=False, + post_init=getattr(cls, '__pydantic_post_init__', None), + config=core_config, + ref=model_ref, + ) + + schema = self._apply_model_serializers(model_schema, decorators.model_serializers.values()) + schema = apply_model_validators(schema, model_validators, 'outer') + return self.defs.create_definition_reference_schema(schema) + + def _resolve_self_type(self, obj: Any) -> Any: + obj = self.model_type_stack.get() + if obj is None: + raise PydanticUserError('`typing.Self` is invalid in this context', code='invalid-self-type') + return obj + + def _generate_schema_from_get_schema_method(self, obj: Any, source: Any) -> core_schema.CoreSchema | None: + BaseModel_ = import_cached_base_model() + + get_schema = getattr(obj, '__get_pydantic_core_schema__', None) + is_base_model_get_schema = ( + getattr(get_schema, '__func__', None) is BaseModel_.__get_pydantic_core_schema__.__func__ # pyright: ignore[reportFunctionMemberAccess] + ) + + if ( + get_schema is not None + # BaseModel.__get_pydantic_core_schema__ is defined for backwards compatibility, + # to allow existing code to call `super().__get_pydantic_core_schema__` in Pydantic + # model that overrides `__get_pydantic_core_schema__`. However, it raises a deprecation + # warning stating that the method will be removed, and during the core schema gen we actually + # don't call the method: + and not is_base_model_get_schema + ): + # Some referenceable types might have a `__get_pydantic_core_schema__` method + # defined on it by users (e.g. on a dataclass). This generally doesn't play well + # as these types are already recognized by the `GenerateSchema` class and isn't ideal + # as we might end up calling `get_schema_or_ref` (expensive) on types that are actually + # not referenceable: + with self.defs.get_schema_or_ref(obj) as (_, maybe_schema): + if maybe_schema is not None: + return maybe_schema + + if obj is source: + ref_mode = 'unpack' + else: + ref_mode = 'to-def' + schema = get_schema( + source, CallbackGetCoreSchemaHandler(self._generate_schema_inner, self, ref_mode=ref_mode) + ) + if schema['type'] == 'definitions': + schema = self.defs.unpack_definitions(schema) + + ref = get_ref(schema) + if ref: + return self.defs.create_definition_reference_schema(schema) + + # Note: if schema is of type `'definition-ref'`, we might want to copy it as a + # safety measure (because these are inlined in place -- i.e. mutated directly) + return schema + + if get_schema is None and (validators := getattr(obj, '__get_validators__', None)) is not None: + from pydantic.v1 import BaseModel as BaseModelV1 + + if issubclass(obj, BaseModelV1): + warn( + f'Mixing V1 models and V2 models (or constructs, like `TypeAdapter`) is not supported. Please upgrade `{obj.__name__}` to V2.', + UserWarning, + ) + else: + warn( + '`__get_validators__` is deprecated and will be removed, use `__get_pydantic_core_schema__` instead.', + PydanticDeprecatedSince20, + ) + return core_schema.chain_schema([core_schema.with_info_plain_validator_function(v) for v in validators()]) + + def _resolve_forward_ref(self, obj: Any) -> Any: + # we assume that types_namespace has the target of forward references in its scope, + # but this could fail, for example, if calling Validator on an imported type which contains + # forward references to other types only defined in the module from which it was imported + # `Validator(SomeImportedTypeAliasWithAForwardReference)` + # or the equivalent for BaseModel + # class Model(BaseModel): + # x: SomeImportedTypeAliasWithAForwardReference + try: + obj = _typing_extra.eval_type_backport(obj, *self._types_namespace) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + + # if obj is still a ForwardRef, it means we can't evaluate it, raise PydanticUndefinedAnnotation + if isinstance(obj, ForwardRef): + raise PydanticUndefinedAnnotation(obj.__forward_arg__, f'Unable to evaluate forward reference {obj}') + + if self._typevars_map: + obj = replace_types(obj, self._typevars_map) + + return obj + + @overload + def _get_args_resolving_forward_refs(self, obj: Any, required: Literal[True]) -> tuple[Any, ...]: ... + + @overload + def _get_args_resolving_forward_refs(self, obj: Any) -> tuple[Any, ...] | None: ... + + def _get_args_resolving_forward_refs(self, obj: Any, required: bool = False) -> tuple[Any, ...] | None: + args = get_args(obj) + if args: + if isinstance(obj, GenericAlias): + # PEP 585 generic aliases don't convert args to ForwardRefs, unlike `typing.List/Dict` etc. + args = (_typing_extra._make_forward_ref(a) if isinstance(a, str) else a for a in args) + args = tuple(self._resolve_forward_ref(a) if isinstance(a, ForwardRef) else a for a in args) + elif required: # pragma: no cover + raise TypeError(f'Expected {obj} to have generic parameters but it had none') + return args + + def _get_first_arg_or_any(self, obj: Any) -> Any: + args = self._get_args_resolving_forward_refs(obj) + if not args: + return Any + return args[0] + + def _get_first_two_args_or_any(self, obj: Any) -> tuple[Any, Any]: + args = self._get_args_resolving_forward_refs(obj) + if not args: + return (Any, Any) + if len(args) < 2: + origin = get_origin(obj) + raise TypeError(f'Expected two type arguments for {origin}, got 1') + return args[0], args[1] + + def _generate_schema_inner(self, obj: Any) -> core_schema.CoreSchema: + if typing_objects.is_self(obj): + obj = self._resolve_self_type(obj) + + if typing_objects.is_annotated(get_origin(obj)): + return self._annotated_schema(obj) + + if isinstance(obj, dict): + # we assume this is already a valid schema + return obj # type: ignore[return-value] + + if isinstance(obj, str): + obj = ForwardRef(obj) + + if isinstance(obj, ForwardRef): + return self.generate_schema(self._resolve_forward_ref(obj)) + + BaseModel = import_cached_base_model() + + if lenient_issubclass(obj, BaseModel): + with self.model_type_stack.push(obj): + return self._model_schema(obj) + + if isinstance(obj, PydanticRecursiveRef): + return core_schema.definition_reference_schema(schema_ref=obj.type_ref) + + return self.match_type(obj) + + def match_type(self, obj: Any) -> core_schema.CoreSchema: # noqa: C901 + """Main mapping of types to schemas. + + The general structure is a series of if statements starting with the simple cases + (non-generic primitive types) and then handling generics and other more complex cases. + + Each case either generates a schema directly, calls into a public user-overridable method + (like `GenerateSchema.tuple_variable_schema`) or calls into a private method that handles some + boilerplate before calling into the user-facing method (e.g. `GenerateSchema._tuple_schema`). + + The idea is that we'll evolve this into adding more and more user facing methods over time + as they get requested and we figure out what the right API for them is. + """ + if obj is str: + return core_schema.str_schema() + elif obj is bytes: + return core_schema.bytes_schema() + elif obj is int: + return core_schema.int_schema() + elif obj is float: + return core_schema.float_schema() + elif obj is bool: + return core_schema.bool_schema() + elif obj is complex: + return core_schema.complex_schema() + elif typing_objects.is_any(obj) or obj is object: + return core_schema.any_schema() + elif obj is datetime.date: + return core_schema.date_schema() + elif obj is datetime.datetime: + return core_schema.datetime_schema() + elif obj is datetime.time: + return core_schema.time_schema() + elif obj is datetime.timedelta: + return core_schema.timedelta_schema() + elif obj is Decimal: + return core_schema.decimal_schema() + elif obj is UUID: + return core_schema.uuid_schema() + elif obj is Url: + return core_schema.url_schema() + elif obj is Fraction: + return self._fraction_schema() + elif obj is MultiHostUrl: + return core_schema.multi_host_url_schema() + elif obj is None or obj is _typing_extra.NoneType: + return core_schema.none_schema() + elif obj in IP_TYPES: + return self._ip_schema(obj) + elif obj in TUPLE_TYPES: + return self._tuple_schema(obj) + elif obj in LIST_TYPES: + return self._list_schema(Any) + elif obj in SET_TYPES: + return self._set_schema(Any) + elif obj in FROZEN_SET_TYPES: + return self._frozenset_schema(Any) + elif obj in SEQUENCE_TYPES: + return self._sequence_schema(Any) + elif obj in ITERABLE_TYPES: + return self._iterable_schema(obj) + elif obj in DICT_TYPES: + return self._dict_schema(Any, Any) + elif obj in PATH_TYPES: + return self._path_schema(obj, Any) + elif obj in DEQUE_TYPES: + return self._deque_schema(Any) + elif obj in MAPPING_TYPES: + return self._mapping_schema(obj, Any, Any) + elif obj in COUNTER_TYPES: + return self._mapping_schema(obj, Any, int) + elif typing_objects.is_typealiastype(obj): + return self._type_alias_type_schema(obj) + elif obj is type: + return self._type_schema() + elif _typing_extra.is_callable(obj): + return core_schema.callable_schema() + elif typing_objects.is_literal(get_origin(obj)): + return self._literal_schema(obj) + elif is_typeddict(obj): + return self._typed_dict_schema(obj, None) + elif _typing_extra.is_namedtuple(obj): + return self._namedtuple_schema(obj, None) + elif typing_objects.is_newtype(obj): + # NewType, can't use isinstance because it fails <3.10 + return self.generate_schema(obj.__supertype__) + elif obj in PATTERN_TYPES: + return self._pattern_schema(obj) + elif _typing_extra.is_hashable(obj): + return self._hashable_schema() + elif isinstance(obj, typing.TypeVar): + return self._unsubstituted_typevar_schema(obj) + elif _typing_extra.is_finalvar(obj): + if obj is Final: + return core_schema.any_schema() + return self.generate_schema( + self._get_first_arg_or_any(obj), + ) + elif isinstance(obj, VALIDATE_CALL_SUPPORTED_TYPES): + return self._call_schema(obj) + elif inspect.isclass(obj) and issubclass(obj, Enum): + return self._enum_schema(obj) + elif obj is ZoneInfo: + return self._zoneinfo_schema() + + # dataclasses.is_dataclass coerces dc instances to types, but we only handle + # the case of a dc type here + if dataclasses.is_dataclass(obj): + return self._dataclass_schema(obj, None) # pyright: ignore[reportArgumentType] + + origin = get_origin(obj) + if origin is not None: + return self._match_generic_type(obj, origin) + + if self._arbitrary_types: + return self._arbitrary_type_schema(obj) + return self._unknown_type_schema(obj) + + def _match_generic_type(self, obj: Any, origin: Any) -> CoreSchema: # noqa: C901 + # Need to handle generic dataclasses before looking for the schema properties because attribute accesses + # on _GenericAlias delegate to the origin type, so lose the information about the concrete parametrization + # As a result, currently, there is no way to cache the schema for generic dataclasses. This may be possible + # to resolve by modifying the value returned by `Generic.__class_getitem__`, but that is a dangerous game. + if dataclasses.is_dataclass(origin): + return self._dataclass_schema(obj, origin) # pyright: ignore[reportArgumentType] + if _typing_extra.is_namedtuple(origin): + return self._namedtuple_schema(obj, origin) + + schema = self._generate_schema_from_get_schema_method(origin, obj) + if schema is not None: + return schema + + if typing_objects.is_typealiastype(origin): + return self._type_alias_type_schema(obj) + elif is_union_origin(origin): + return self._union_schema(obj) + elif origin in TUPLE_TYPES: + return self._tuple_schema(obj) + elif origin in LIST_TYPES: + return self._list_schema(self._get_first_arg_or_any(obj)) + elif origin in SET_TYPES: + return self._set_schema(self._get_first_arg_or_any(obj)) + elif origin in FROZEN_SET_TYPES: + return self._frozenset_schema(self._get_first_arg_or_any(obj)) + elif origin in DICT_TYPES: + return self._dict_schema(*self._get_first_two_args_or_any(obj)) + elif origin in PATH_TYPES: + return self._path_schema(origin, self._get_first_arg_or_any(obj)) + elif origin in DEQUE_TYPES: + return self._deque_schema(self._get_first_arg_or_any(obj)) + elif origin in MAPPING_TYPES: + return self._mapping_schema(origin, *self._get_first_two_args_or_any(obj)) + elif origin in COUNTER_TYPES: + return self._mapping_schema(origin, self._get_first_arg_or_any(obj), int) + elif is_typeddict(origin): + return self._typed_dict_schema(obj, origin) + elif origin in TYPE_TYPES: + return self._subclass_schema(obj) + elif origin in SEQUENCE_TYPES: + return self._sequence_schema(self._get_first_arg_or_any(obj)) + elif origin in ITERABLE_TYPES: + return self._iterable_schema(obj) + elif origin in PATTERN_TYPES: + return self._pattern_schema(obj) + + if self._arbitrary_types: + return self._arbitrary_type_schema(origin) + return self._unknown_type_schema(obj) + + def _generate_td_field_schema( + self, + name: str, + field_info: FieldInfo, + decorators: DecoratorInfos, + *, + required: bool = True, + ) -> core_schema.TypedDictField: + """Prepare a TypedDictField to represent a model or typeddict field.""" + common_field = self._common_field_schema(name, field_info, decorators) + return core_schema.typed_dict_field( + common_field['schema'], + required=False if not field_info.is_required() else required, + serialization_exclude=common_field['serialization_exclude'], + validation_alias=common_field['validation_alias'], + serialization_alias=common_field['serialization_alias'], + metadata=common_field['metadata'], + ) + + def _generate_md_field_schema( + self, + name: str, + field_info: FieldInfo, + decorators: DecoratorInfos, + ) -> core_schema.ModelField: + """Prepare a ModelField to represent a model field.""" + common_field = self._common_field_schema(name, field_info, decorators) + return core_schema.model_field( + common_field['schema'], + serialization_exclude=common_field['serialization_exclude'], + validation_alias=common_field['validation_alias'], + serialization_alias=common_field['serialization_alias'], + frozen=common_field['frozen'], + metadata=common_field['metadata'], + ) + + def _generate_dc_field_schema( + self, + name: str, + field_info: FieldInfo, + decorators: DecoratorInfos, + ) -> core_schema.DataclassField: + """Prepare a DataclassField to represent the parameter/field, of a dataclass.""" + common_field = self._common_field_schema(name, field_info, decorators) + return core_schema.dataclass_field( + name, + common_field['schema'], + init=field_info.init, + init_only=field_info.init_var or None, + kw_only=None if field_info.kw_only else False, + serialization_exclude=common_field['serialization_exclude'], + validation_alias=common_field['validation_alias'], + serialization_alias=common_field['serialization_alias'], + frozen=common_field['frozen'], + metadata=common_field['metadata'], + ) + + @staticmethod + def _apply_alias_generator_to_field_info( + alias_generator: Callable[[str], str] | AliasGenerator, field_info: FieldInfo, field_name: str + ) -> None: + """Apply an alias_generator to aliases on a FieldInfo instance if appropriate. + + Args: + alias_generator: A callable that takes a string and returns a string, or an AliasGenerator instance. + field_info: The FieldInfo instance to which the alias_generator is (maybe) applied. + field_name: The name of the field from which to generate the alias. + """ + # Apply an alias_generator if + # 1. An alias is not specified + # 2. An alias is specified, but the priority is <= 1 + if ( + field_info.alias_priority is None + or field_info.alias_priority <= 1 + or field_info.alias is None + or field_info.validation_alias is None + or field_info.serialization_alias is None + ): + alias, validation_alias, serialization_alias = None, None, None + + if isinstance(alias_generator, AliasGenerator): + alias, validation_alias, serialization_alias = alias_generator.generate_aliases(field_name) + elif isinstance(alias_generator, Callable): + alias = alias_generator(field_name) + if not isinstance(alias, str): + raise TypeError(f'alias_generator {alias_generator} must return str, not {alias.__class__}') + + # if priority is not set, we set to 1 + # which supports the case where the alias_generator from a child class is used + # to generate an alias for a field in a parent class + if field_info.alias_priority is None or field_info.alias_priority <= 1: + field_info.alias_priority = 1 + + # if the priority is 1, then we set the aliases to the generated alias + if field_info.alias_priority == 1: + field_info.serialization_alias = _get_first_non_null(serialization_alias, alias) + field_info.validation_alias = _get_first_non_null(validation_alias, alias) + field_info.alias = alias + + # if any of the aliases are not set, then we set them to the corresponding generated alias + if field_info.alias is None: + field_info.alias = alias + if field_info.serialization_alias is None: + field_info.serialization_alias = _get_first_non_null(serialization_alias, alias) + if field_info.validation_alias is None: + field_info.validation_alias = _get_first_non_null(validation_alias, alias) + + @staticmethod + def _apply_alias_generator_to_computed_field_info( + alias_generator: Callable[[str], str] | AliasGenerator, + computed_field_info: ComputedFieldInfo, + computed_field_name: str, + ): + """Apply an alias_generator to alias on a ComputedFieldInfo instance if appropriate. + + Args: + alias_generator: A callable that takes a string and returns a string, or an AliasGenerator instance. + computed_field_info: The ComputedFieldInfo instance to which the alias_generator is (maybe) applied. + computed_field_name: The name of the computed field from which to generate the alias. + """ + # Apply an alias_generator if + # 1. An alias is not specified + # 2. An alias is specified, but the priority is <= 1 + + if ( + computed_field_info.alias_priority is None + or computed_field_info.alias_priority <= 1 + or computed_field_info.alias is None + ): + alias, validation_alias, serialization_alias = None, None, None + + if isinstance(alias_generator, AliasGenerator): + alias, validation_alias, serialization_alias = alias_generator.generate_aliases(computed_field_name) + elif isinstance(alias_generator, Callable): + alias = alias_generator(computed_field_name) + if not isinstance(alias, str): + raise TypeError(f'alias_generator {alias_generator} must return str, not {alias.__class__}') + + # if priority is not set, we set to 1 + # which supports the case where the alias_generator from a child class is used + # to generate an alias for a field in a parent class + if computed_field_info.alias_priority is None or computed_field_info.alias_priority <= 1: + computed_field_info.alias_priority = 1 + + # if the priority is 1, then we set the aliases to the generated alias + # note that we use the serialization_alias with priority over alias, as computed_field + # aliases are used for serialization only (not validation) + if computed_field_info.alias_priority == 1: + computed_field_info.alias = _get_first_non_null(serialization_alias, alias) + + @staticmethod + def _apply_field_title_generator_to_field_info( + config_wrapper: ConfigWrapper, field_info: FieldInfo | ComputedFieldInfo, field_name: str + ) -> None: + """Apply a field_title_generator on a FieldInfo or ComputedFieldInfo instance if appropriate + Args: + config_wrapper: The config of the model + field_info: The FieldInfo or ComputedField instance to which the title_generator is (maybe) applied. + field_name: The name of the field from which to generate the title. + """ + field_title_generator = field_info.field_title_generator or config_wrapper.field_title_generator + + if field_title_generator is None: + return + + if field_info.title is None: + title = field_title_generator(field_name, field_info) # type: ignore + if not isinstance(title, str): + raise TypeError(f'field_title_generator {field_title_generator} must return str, not {title.__class__}') + + field_info.title = title + + def _common_field_schema( # C901 + self, name: str, field_info: FieldInfo, decorators: DecoratorInfos + ) -> _CommonField: + source_type, annotations = field_info.annotation, field_info.metadata + + def set_discriminator(schema: CoreSchema) -> CoreSchema: + schema = self._apply_discriminator_to_union(schema, field_info.discriminator) + return schema + + # Convert `@field_validator` decorators to `Before/After/Plain/WrapValidator` instances: + validators_from_decorators = [] + for decorator in filter_field_decorator_info_by_field(decorators.field_validators.values(), name): + validators_from_decorators.append(_mode_to_validator[decorator.info.mode]._from_decorator(decorator)) + + with self.field_name_stack.push(name): + if field_info.discriminator is not None: + schema = self._apply_annotations( + source_type, annotations + validators_from_decorators, transform_inner_schema=set_discriminator + ) + else: + schema = self._apply_annotations( + source_type, + annotations + validators_from_decorators, + ) + + # This V1 compatibility shim should eventually be removed + # push down any `each_item=True` validators + # note that this won't work for any Annotated types that get wrapped by a function validator + # but that's okay because that didn't exist in V1 + this_field_validators = filter_field_decorator_info_by_field(decorators.validators.values(), name) + if _validators_require_validate_default(this_field_validators): + field_info.validate_default = True + each_item_validators = [v for v in this_field_validators if v.info.each_item is True] + this_field_validators = [v for v in this_field_validators if v not in each_item_validators] + schema = apply_each_item_validators(schema, each_item_validators, name) + + schema = apply_validators(schema, this_field_validators, name) + + # the default validator needs to go outside of any other validators + # so that it is the topmost validator for the field validator + # which uses it to check if the field has a default value or not + if not field_info.is_required(): + schema = wrap_default(field_info, schema) + + schema = self._apply_field_serializers( + schema, filter_field_decorator_info_by_field(decorators.field_serializers.values(), name) + ) + self._apply_field_title_generator_to_field_info(self._config_wrapper, field_info, name) + + pydantic_js_updates, pydantic_js_extra = _extract_json_schema_info_from_field_info(field_info) + core_metadata: dict[str, Any] = {} + update_core_metadata( + core_metadata, pydantic_js_updates=pydantic_js_updates, pydantic_js_extra=pydantic_js_extra + ) + + alias_generator = self._config_wrapper.alias_generator + if alias_generator is not None: + self._apply_alias_generator_to_field_info(alias_generator, field_info, name) + + if isinstance(field_info.validation_alias, (AliasChoices, AliasPath)): + validation_alias = field_info.validation_alias.convert_to_aliases() + else: + validation_alias = field_info.validation_alias + + return _common_field( + schema, + serialization_exclude=True if field_info.exclude else None, + validation_alias=validation_alias, + serialization_alias=field_info.serialization_alias, + frozen=field_info.frozen, + metadata=core_metadata, + ) + + def _union_schema(self, union_type: Any) -> core_schema.CoreSchema: + """Generate schema for a Union.""" + args = self._get_args_resolving_forward_refs(union_type, required=True) + choices: list[CoreSchema] = [] + nullable = False + for arg in args: + if arg is None or arg is _typing_extra.NoneType: + nullable = True + else: + choices.append(self.generate_schema(arg)) + + if len(choices) == 1: + s = choices[0] + else: + choices_with_tags: list[CoreSchema | tuple[CoreSchema, str]] = [] + for choice in choices: + tag = cast(CoreMetadata, choice.get('metadata', {})).get('pydantic_internal_union_tag_key') + if tag is not None: + choices_with_tags.append((choice, tag)) + else: + choices_with_tags.append(choice) + s = core_schema.union_schema(choices_with_tags) + + if nullable: + s = core_schema.nullable_schema(s) + return s + + def _type_alias_type_schema(self, obj: TypeAliasType) -> CoreSchema: + with self.defs.get_schema_or_ref(obj) as (ref, maybe_schema): + if maybe_schema is not None: + return maybe_schema + + origin: TypeAliasType = get_origin(obj) or obj + typevars_map = get_standard_typevars_map(obj) + + with self._ns_resolver.push(origin): + try: + annotation = _typing_extra.eval_type(origin.__value__, *self._types_namespace) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + annotation = replace_types(annotation, typevars_map) + schema = self.generate_schema(annotation) + assert schema['type'] != 'definitions' + schema['ref'] = ref # type: ignore + return self.defs.create_definition_reference_schema(schema) + + def _literal_schema(self, literal_type: Any) -> CoreSchema: + """Generate schema for a Literal.""" + expected = list(get_literal_values(literal_type, type_check=False, unpack_type_aliases='eager')) + assert expected, f'literal "expected" cannot be empty, obj={literal_type}' + schema = core_schema.literal_schema(expected) + + if self._config_wrapper.use_enum_values and any(isinstance(v, Enum) for v in expected): + schema = core_schema.no_info_after_validator_function( + lambda v: v.value if isinstance(v, Enum) else v, schema + ) + + return schema + + def _typed_dict_schema(self, typed_dict_cls: Any, origin: Any) -> core_schema.CoreSchema: + """Generate a core schema for a `TypedDict` class. + + To be able to build a `DecoratorInfos` instance for the `TypedDict` class (which will include + validators, serializers, etc.), we need to have access to the original bases of the class + (see https://docs.python.org/3/library/types.html#types.get_original_bases). + However, the `__orig_bases__` attribute was only added in 3.12 (https://github.com/python/cpython/pull/103698). + + For this reason, we require Python 3.12 (or using the `typing_extensions` backport). + """ + FieldInfo = import_cached_field_info() + + with ( + self.model_type_stack.push(typed_dict_cls), + self.defs.get_schema_or_ref(typed_dict_cls) as ( + typed_dict_ref, + maybe_schema, + ), + ): + if maybe_schema is not None: + return maybe_schema + + typevars_map = get_standard_typevars_map(typed_dict_cls) + if origin is not None: + typed_dict_cls = origin + + if not _SUPPORTS_TYPEDDICT and type(typed_dict_cls).__module__ == 'typing': + raise PydanticUserError( + 'Please use `typing_extensions.TypedDict` instead of `typing.TypedDict` on Python < 3.12.', + code='typed-dict-version', + ) + + try: + # if a typed dictionary class doesn't have config, we use the parent's config, hence a default of `None` + # see https://github.com/pydantic/pydantic/issues/10917 + config: ConfigDict | None = get_attribute_from_bases(typed_dict_cls, '__pydantic_config__') + except AttributeError: + config = None + + with self._config_wrapper_stack.push(config): + core_config = self._config_wrapper.core_config(title=typed_dict_cls.__name__) + + required_keys: frozenset[str] = typed_dict_cls.__required_keys__ + + fields: dict[str, core_schema.TypedDictField] = {} + + decorators = DecoratorInfos.build(typed_dict_cls) + + if self._config_wrapper.use_attribute_docstrings: + field_docstrings = extract_docstrings_from_cls(typed_dict_cls, use_inspect=True) + else: + field_docstrings = None + + try: + annotations = _typing_extra.get_cls_type_hints(typed_dict_cls, ns_resolver=self._ns_resolver) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + + readonly_fields: list[str] = [] + + for field_name, annotation in annotations.items(): + field_info = FieldInfo.from_annotation(annotation, _source=AnnotationSource.TYPED_DICT) + field_info.annotation = replace_types(field_info.annotation, typevars_map) + + required = ( + field_name in required_keys or 'required' in field_info._qualifiers + ) and 'not_required' not in field_info._qualifiers + if 'read_only' in field_info._qualifiers: + readonly_fields.append(field_name) + + if ( + field_docstrings is not None + and field_info.description is None + and field_name in field_docstrings + ): + field_info.description = field_docstrings[field_name] + self._apply_field_title_generator_to_field_info(self._config_wrapper, field_info, field_name) + fields[field_name] = self._generate_td_field_schema( + field_name, field_info, decorators, required=required + ) + + if readonly_fields: + fields_repr = ', '.join(repr(f) for f in readonly_fields) + plural = len(readonly_fields) >= 2 + warnings.warn( + f'Item{"s" if plural else ""} {fields_repr} on TypedDict class {typed_dict_cls.__name__!r} ' + f'{"are" if plural else "is"} using the `ReadOnly` qualifier. Pydantic will not protect items ' + 'from any mutation on dictionary instances.', + UserWarning, + ) + + td_schema = core_schema.typed_dict_schema( + fields, + cls=typed_dict_cls, + computed_fields=[ + self._computed_field_schema(d, decorators.field_serializers) + for d in decorators.computed_fields.values() + ], + ref=typed_dict_ref, + config=core_config, + ) + + schema = self._apply_model_serializers(td_schema, decorators.model_serializers.values()) + schema = apply_model_validators(schema, decorators.model_validators.values(), 'all') + return self.defs.create_definition_reference_schema(schema) + + def _namedtuple_schema(self, namedtuple_cls: Any, origin: Any) -> core_schema.CoreSchema: + """Generate schema for a NamedTuple.""" + with ( + self.model_type_stack.push(namedtuple_cls), + self.defs.get_schema_or_ref(namedtuple_cls) as ( + namedtuple_ref, + maybe_schema, + ), + ): + if maybe_schema is not None: + return maybe_schema + typevars_map = get_standard_typevars_map(namedtuple_cls) + if origin is not None: + namedtuple_cls = origin + + try: + annotations = _typing_extra.get_cls_type_hints(namedtuple_cls, ns_resolver=self._ns_resolver) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + if not annotations: + # annotations is empty, happens if namedtuple_cls defined via collections.namedtuple(...) + annotations: dict[str, Any] = {k: Any for k in namedtuple_cls._fields} + + if typevars_map: + annotations = { + field_name: replace_types(annotation, typevars_map) + for field_name, annotation in annotations.items() + } + + arguments_schema = core_schema.arguments_schema( + [ + self._generate_parameter_schema( + field_name, + annotation, + source=AnnotationSource.NAMED_TUPLE, + default=namedtuple_cls._field_defaults.get(field_name, Parameter.empty), + ) + for field_name, annotation in annotations.items() + ], + metadata={'pydantic_js_prefer_positional_arguments': True}, + ) + schema = core_schema.call_schema(arguments_schema, namedtuple_cls, ref=namedtuple_ref) + return self.defs.create_definition_reference_schema(schema) + + def _generate_parameter_schema( + self, + name: str, + annotation: type[Any], + source: AnnotationSource, + default: Any = Parameter.empty, + mode: Literal['positional_only', 'positional_or_keyword', 'keyword_only'] | None = None, + ) -> core_schema.ArgumentsParameter: + """Generate the definition of a field in a namedtuple or a parameter in a function signature. + + This definition is meant to be used for the `'arguments'` core schema, which will be replaced + in V3 by the `'arguments-v3`'. + """ + FieldInfo = import_cached_field_info() + + if default is Parameter.empty: + field = FieldInfo.from_annotation(annotation, _source=source) + else: + field = FieldInfo.from_annotated_attribute(annotation, default, _source=source) + assert field.annotation is not None, 'field.annotation should not be None when generating a schema' + with self.field_name_stack.push(name): + schema = self._apply_annotations(field.annotation, [field]) + + if not field.is_required(): + schema = wrap_default(field, schema) + + parameter_schema = core_schema.arguments_parameter(name, schema) + if mode is not None: + parameter_schema['mode'] = mode + if field.alias is not None: + parameter_schema['alias'] = field.alias + else: + alias_generator = self._config_wrapper.alias_generator + if isinstance(alias_generator, AliasGenerator) and alias_generator.alias is not None: + parameter_schema['alias'] = alias_generator.alias(name) + elif callable(alias_generator): + parameter_schema['alias'] = alias_generator(name) + return parameter_schema + + def _generate_parameter_v3_schema( + self, + name: str, + annotation: Any, + source: AnnotationSource, + mode: Literal[ + 'positional_only', + 'positional_or_keyword', + 'keyword_only', + 'var_args', + 'var_kwargs_uniform', + 'var_kwargs_unpacked_typed_dict', + ], + default: Any = Parameter.empty, + ) -> core_schema.ArgumentsV3Parameter: + """Generate the definition of a parameter in a function signature. + + This definition is meant to be used for the `'arguments-v3'` core schema, which will replace + the `'arguments`' schema in V3. + """ + FieldInfo = import_cached_field_info() + + if default is Parameter.empty: + field = FieldInfo.from_annotation(annotation, _source=source) + else: + field = FieldInfo.from_annotated_attribute(annotation, default, _source=source) + + with self.field_name_stack.push(name): + schema = self._apply_annotations(field.annotation, [field]) + + if not field.is_required(): + schema = wrap_default(field, schema) + + parameter_schema = core_schema.arguments_v3_parameter( + name=name, + schema=schema, + mode=mode, + ) + if field.alias is not None: + parameter_schema['alias'] = field.alias + else: + alias_generator = self._config_wrapper.alias_generator + if isinstance(alias_generator, AliasGenerator) and alias_generator.alias is not None: + parameter_schema['alias'] = alias_generator.alias(name) + elif callable(alias_generator): + parameter_schema['alias'] = alias_generator(name) + + return parameter_schema + + def _tuple_schema(self, tuple_type: Any) -> core_schema.CoreSchema: + """Generate schema for a Tuple, e.g. `tuple[int, str]` or `tuple[int, ...]`.""" + # TODO: do we really need to resolve type vars here? + typevars_map = get_standard_typevars_map(tuple_type) + params = self._get_args_resolving_forward_refs(tuple_type) + + if typevars_map and params: + params = tuple(replace_types(param, typevars_map) for param in params) + + # NOTE: subtle difference: `tuple[()]` gives `params=()`, whereas `typing.Tuple[()]` gives `params=((),)` + # This is only true for <3.11, on Python 3.11+ `typing.Tuple[()]` gives `params=()` + if not params: + if tuple_type in TUPLE_TYPES: + return core_schema.tuple_schema([core_schema.any_schema()], variadic_item_index=0) + else: + # special case for `tuple[()]` which means `tuple[]` - an empty tuple + return core_schema.tuple_schema([]) + elif params[-1] is Ellipsis: + if len(params) == 2: + return core_schema.tuple_schema([self.generate_schema(params[0])], variadic_item_index=0) + else: + # TODO: something like https://github.com/pydantic/pydantic/issues/5952 + raise ValueError('Variable tuples can only have one type') + elif len(params) == 1 and params[0] == (): + # special case for `tuple[()]` which means `tuple[]` - an empty tuple + # NOTE: This conditional can be removed when we drop support for Python 3.10. + return core_schema.tuple_schema([]) + else: + return core_schema.tuple_schema([self.generate_schema(param) for param in params]) + + def _type_schema(self) -> core_schema.CoreSchema: + return core_schema.custom_error_schema( + core_schema.is_instance_schema(type), + custom_error_type='is_type', + custom_error_message='Input should be a type', + ) + + def _zoneinfo_schema(self) -> core_schema.CoreSchema: + """Generate schema for a zone_info.ZoneInfo object""" + from ._validators import validate_str_is_valid_iana_tz + + metadata = {'pydantic_js_functions': [lambda _1, _2: {'type': 'string', 'format': 'zoneinfo'}]} + return core_schema.no_info_plain_validator_function( + validate_str_is_valid_iana_tz, + serialization=core_schema.to_string_ser_schema(), + metadata=metadata, + ) + + def _union_is_subclass_schema(self, union_type: Any) -> core_schema.CoreSchema: + """Generate schema for `type[Union[X, ...]]`.""" + args = self._get_args_resolving_forward_refs(union_type, required=True) + return core_schema.union_schema([self.generate_schema(type[args]) for args in args]) + + def _subclass_schema(self, type_: Any) -> core_schema.CoreSchema: + """Generate schema for a type, e.g. `type[int]`.""" + type_param = self._get_first_arg_or_any(type_) + + # Assume `type[Annotated[, ...]]` is equivalent to `type[]`: + type_param = _typing_extra.annotated_type(type_param) or type_param + + if typing_objects.is_any(type_param): + return self._type_schema() + elif typing_objects.is_typealiastype(type_param): + return self.generate_schema(type[type_param.__value__]) + elif typing_objects.is_typevar(type_param): + if type_param.__bound__: + if is_union_origin(get_origin(type_param.__bound__)): + return self._union_is_subclass_schema(type_param.__bound__) + return core_schema.is_subclass_schema(type_param.__bound__) + elif type_param.__constraints__: + return core_schema.union_schema([self.generate_schema(type[c]) for c in type_param.__constraints__]) + else: + return self._type_schema() + elif is_union_origin(get_origin(type_param)): + return self._union_is_subclass_schema(type_param) + else: + if typing_objects.is_self(type_param): + type_param = self._resolve_self_type(type_param) + if _typing_extra.is_generic_alias(type_param): + raise PydanticUserError( + 'Subscripting `type[]` with an already parametrized type is not supported. ' + f'Instead of using type[{type_param!r}], use type[{_repr.display_as_type(get_origin(type_param))}].', + code=None, + ) + if not inspect.isclass(type_param): + # when using type[None], this doesn't type convert to type[NoneType], and None isn't a class + # so we handle it manually here + if type_param is None: + return core_schema.is_subclass_schema(_typing_extra.NoneType) + raise TypeError(f'Expected a class, got {type_param!r}') + return core_schema.is_subclass_schema(type_param) + + def _sequence_schema(self, items_type: Any) -> core_schema.CoreSchema: + """Generate schema for a Sequence, e.g. `Sequence[int]`.""" + from ._serializers import serialize_sequence_via_list + + item_type_schema = self.generate_schema(items_type) + list_schema = core_schema.list_schema(item_type_schema) + + json_schema = smart_deepcopy(list_schema) + python_schema = core_schema.is_instance_schema(typing.Sequence, cls_repr='Sequence') + if not typing_objects.is_any(items_type): + from ._validators import sequence_validator + + python_schema = core_schema.chain_schema( + [python_schema, core_schema.no_info_wrap_validator_function(sequence_validator, list_schema)], + ) + + serialization = core_schema.wrap_serializer_function_ser_schema( + serialize_sequence_via_list, schema=item_type_schema, info_arg=True + ) + return core_schema.json_or_python_schema( + json_schema=json_schema, python_schema=python_schema, serialization=serialization + ) + + def _iterable_schema(self, type_: Any) -> core_schema.GeneratorSchema: + """Generate a schema for an `Iterable`.""" + item_type = self._get_first_arg_or_any(type_) + + return core_schema.generator_schema(self.generate_schema(item_type)) + + def _pattern_schema(self, pattern_type: Any) -> core_schema.CoreSchema: + from . import _validators + + metadata = {'pydantic_js_functions': [lambda _1, _2: {'type': 'string', 'format': 'regex'}]} + ser = core_schema.plain_serializer_function_ser_schema( + attrgetter('pattern'), when_used='json', return_schema=core_schema.str_schema() + ) + if pattern_type is typing.Pattern or pattern_type is re.Pattern: + # bare type + return core_schema.no_info_plain_validator_function( + _validators.pattern_either_validator, serialization=ser, metadata=metadata + ) + + param = self._get_args_resolving_forward_refs( + pattern_type, + required=True, + )[0] + if param is str: + return core_schema.no_info_plain_validator_function( + _validators.pattern_str_validator, serialization=ser, metadata=metadata + ) + elif param is bytes: + return core_schema.no_info_plain_validator_function( + _validators.pattern_bytes_validator, serialization=ser, metadata=metadata + ) + else: + raise PydanticSchemaGenerationError(f'Unable to generate pydantic-core schema for {pattern_type!r}.') + + def _hashable_schema(self) -> core_schema.CoreSchema: + return core_schema.custom_error_schema( + schema=core_schema.json_or_python_schema( + json_schema=core_schema.chain_schema( + [core_schema.any_schema(), core_schema.is_instance_schema(collections.abc.Hashable)] + ), + python_schema=core_schema.is_instance_schema(collections.abc.Hashable), + ), + custom_error_type='is_hashable', + custom_error_message='Input should be hashable', + ) + + def _dataclass_schema( + self, dataclass: type[StandardDataclass], origin: type[StandardDataclass] | None + ) -> core_schema.CoreSchema: + """Generate schema for a dataclass.""" + with ( + self.model_type_stack.push(dataclass), + self.defs.get_schema_or_ref(dataclass) as ( + dataclass_ref, + maybe_schema, + ), + ): + if maybe_schema is not None: + return maybe_schema + + schema = dataclass.__dict__.get('__pydantic_core_schema__') + if schema is not None and not isinstance(schema, MockCoreSchema): + if schema['type'] == 'definitions': + schema = self.defs.unpack_definitions(schema) + ref = get_ref(schema) + if ref: + return self.defs.create_definition_reference_schema(schema) + else: + return schema + + typevars_map = get_standard_typevars_map(dataclass) + if origin is not None: + dataclass = origin + + # if (plain) dataclass doesn't have config, we use the parent's config, hence a default of `None` + # (Pydantic dataclasses have an empty dict config by default). + # see https://github.com/pydantic/pydantic/issues/10917 + config = getattr(dataclass, '__pydantic_config__', None) + + from ..dataclasses import is_pydantic_dataclass + + with self._ns_resolver.push(dataclass), self._config_wrapper_stack.push(config): + if is_pydantic_dataclass(dataclass): + if dataclass.__pydantic_fields_complete__(): + # Copy the field info instances to avoid mutating the `FieldInfo` instances + # of the generic dataclass generic origin (e.g. `apply_typevars_map` below). + # Note that we don't apply `deepcopy` on `__pydantic_fields__` because we + # don't want to copy the `FieldInfo` attributes: + fields = { + f_name: copy(field_info) for f_name, field_info in dataclass.__pydantic_fields__.items() + } + if typevars_map: + for field in fields.values(): + field.apply_typevars_map(typevars_map, *self._types_namespace) + else: + try: + fields = rebuild_dataclass_fields( + dataclass, + config_wrapper=self._config_wrapper, + ns_resolver=self._ns_resolver, + typevars_map=typevars_map or {}, + ) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + else: + fields = collect_dataclass_fields( + dataclass, + typevars_map=typevars_map, + config_wrapper=self._config_wrapper, + ) + + if self._config_wrapper.extra == 'allow': + # disallow combination of init=False on a dataclass field and extra='allow' on a dataclass + for field_name, field in fields.items(): + if field.init is False: + raise PydanticUserError( + f'Field {field_name} has `init=False` and dataclass has config setting `extra="allow"`. ' + f'This combination is not allowed.', + code='dataclass-init-false-extra-allow', + ) + + decorators = dataclass.__dict__.get('__pydantic_decorators__') or DecoratorInfos.build(dataclass) + # Move kw_only=False args to the start of the list, as this is how vanilla dataclasses work. + # Note that when kw_only is missing or None, it is treated as equivalent to kw_only=True + args = sorted( + (self._generate_dc_field_schema(k, v, decorators) for k, v in fields.items()), + key=lambda a: a.get('kw_only') is not False, + ) + has_post_init = hasattr(dataclass, '__post_init__') + has_slots = hasattr(dataclass, '__slots__') + + args_schema = core_schema.dataclass_args_schema( + dataclass.__name__, + args, + computed_fields=[ + self._computed_field_schema(d, decorators.field_serializers) + for d in decorators.computed_fields.values() + ], + collect_init_only=has_post_init, + ) + + inner_schema = apply_validators(args_schema, decorators.root_validators.values(), None) + + model_validators = decorators.model_validators.values() + inner_schema = apply_model_validators(inner_schema, model_validators, 'inner') + + core_config = self._config_wrapper.core_config(title=dataclass.__name__) + + dc_schema = core_schema.dataclass_schema( + dataclass, + inner_schema, + generic_origin=origin, + post_init=has_post_init, + ref=dataclass_ref, + fields=[field.name for field in dataclasses.fields(dataclass)], + slots=has_slots, + config=core_config, + # we don't use a custom __setattr__ for dataclasses, so we must + # pass along the frozen config setting to the pydantic-core schema + frozen=self._config_wrapper_stack.tail.frozen, + ) + schema = self._apply_model_serializers(dc_schema, decorators.model_serializers.values()) + schema = apply_model_validators(schema, model_validators, 'outer') + return self.defs.create_definition_reference_schema(schema) + + def _call_schema(self, function: ValidateCallSupportedTypes) -> core_schema.CallSchema: + """Generate schema for a Callable. + + TODO support functional validators once we support them in Config + """ + arguments_schema = self._arguments_schema(function) + + return_schema: core_schema.CoreSchema | None = None + config_wrapper = self._config_wrapper + if config_wrapper.validate_return: + sig = signature(function) + return_hint = sig.return_annotation + if return_hint is not sig.empty: + globalns, localns = self._types_namespace + type_hints = _typing_extra.get_function_type_hints( + function, globalns=globalns, localns=localns, include_keys={'return'} + ) + return_schema = self.generate_schema(type_hints['return']) + + return core_schema.call_schema( + arguments_schema, + function, + return_schema=return_schema, + ) + + def _arguments_schema( + self, function: ValidateCallSupportedTypes, parameters_callback: ParametersCallback | None = None + ) -> core_schema.ArgumentsSchema: + """Generate schema for a Signature.""" + mode_lookup: dict[_ParameterKind, Literal['positional_only', 'positional_or_keyword', 'keyword_only']] = { + Parameter.POSITIONAL_ONLY: 'positional_only', + Parameter.POSITIONAL_OR_KEYWORD: 'positional_or_keyword', + Parameter.KEYWORD_ONLY: 'keyword_only', + } + + sig = signature(function) + globalns, localns = self._types_namespace + type_hints = _typing_extra.get_function_type_hints(function, globalns=globalns, localns=localns) + + arguments_list: list[core_schema.ArgumentsParameter] = [] + var_args_schema: core_schema.CoreSchema | None = None + var_kwargs_schema: core_schema.CoreSchema | None = None + var_kwargs_mode: core_schema.VarKwargsMode | None = None + + for i, (name, p) in enumerate(sig.parameters.items()): + if p.annotation is sig.empty: + annotation = typing.cast(Any, Any) + else: + annotation = type_hints[name] + + if parameters_callback is not None: + result = parameters_callback(i, name, annotation) + if result == 'skip': + continue + + parameter_mode = mode_lookup.get(p.kind) + if parameter_mode is not None: + arg_schema = self._generate_parameter_schema( + name, annotation, AnnotationSource.FUNCTION, p.default, parameter_mode + ) + arguments_list.append(arg_schema) + elif p.kind == Parameter.VAR_POSITIONAL: + var_args_schema = self.generate_schema(annotation) + else: + assert p.kind == Parameter.VAR_KEYWORD, p.kind + + unpack_type = _typing_extra.unpack_type(annotation) + if unpack_type is not None: + origin = get_origin(unpack_type) or unpack_type + if not is_typeddict(origin): + raise PydanticUserError( + f'Expected a `TypedDict` class inside `Unpack[...]`, got {unpack_type!r}', + code='unpack-typed-dict', + ) + non_pos_only_param_names = { + name for name, p in sig.parameters.items() if p.kind != Parameter.POSITIONAL_ONLY + } + overlapping_params = non_pos_only_param_names.intersection(origin.__annotations__) + if overlapping_params: + raise PydanticUserError( + f'Typed dictionary {origin.__name__!r} overlaps with parameter' + f'{"s" if len(overlapping_params) >= 2 else ""} ' + f'{", ".join(repr(p) for p in sorted(overlapping_params))}', + code='overlapping-unpack-typed-dict', + ) + + var_kwargs_mode = 'unpacked-typed-dict' + var_kwargs_schema = self._typed_dict_schema(unpack_type, get_origin(unpack_type)) + else: + var_kwargs_mode = 'uniform' + var_kwargs_schema = self.generate_schema(annotation) + + return core_schema.arguments_schema( + arguments_list, + var_args_schema=var_args_schema, + var_kwargs_mode=var_kwargs_mode, + var_kwargs_schema=var_kwargs_schema, + validate_by_name=self._config_wrapper.validate_by_name, + ) + + def _arguments_v3_schema( + self, function: ValidateCallSupportedTypes, parameters_callback: ParametersCallback | None = None + ) -> core_schema.ArgumentsV3Schema: + mode_lookup: dict[ + _ParameterKind, Literal['positional_only', 'positional_or_keyword', 'var_args', 'keyword_only'] + ] = { + Parameter.POSITIONAL_ONLY: 'positional_only', + Parameter.POSITIONAL_OR_KEYWORD: 'positional_or_keyword', + Parameter.VAR_POSITIONAL: 'var_args', + Parameter.KEYWORD_ONLY: 'keyword_only', + } + + sig = signature(function) + globalns, localns = self._types_namespace + type_hints = _typing_extra.get_function_type_hints(function, globalns=globalns, localns=localns) + + parameters_list: list[core_schema.ArgumentsV3Parameter] = [] + + for i, (name, p) in enumerate(sig.parameters.items()): + if parameters_callback is not None: + result = parameters_callback(i, name, p.annotation) + if result == 'skip': + continue + + if p.annotation is Parameter.empty: + annotation = typing.cast(Any, Any) + else: + annotation = type_hints[name] + + parameter_mode = mode_lookup.get(p.kind) + if parameter_mode is None: + assert p.kind == Parameter.VAR_KEYWORD, p.kind + + unpack_type = _typing_extra.unpack_type(annotation) + if unpack_type is not None: + origin = get_origin(unpack_type) or unpack_type + if not is_typeddict(origin): + raise PydanticUserError( + f'Expected a `TypedDict` class inside `Unpack[...]`, got {unpack_type!r}', + code='unpack-typed-dict', + ) + non_pos_only_param_names = { + name for name, p in sig.parameters.items() if p.kind != Parameter.POSITIONAL_ONLY + } + overlapping_params = non_pos_only_param_names.intersection(origin.__annotations__) + if overlapping_params: + raise PydanticUserError( + f'Typed dictionary {origin.__name__!r} overlaps with parameter' + f'{"s" if len(overlapping_params) >= 2 else ""} ' + f'{", ".join(repr(p) for p in sorted(overlapping_params))}', + code='overlapping-unpack-typed-dict', + ) + parameter_mode = 'var_kwargs_unpacked_typed_dict' + annotation = unpack_type + else: + parameter_mode = 'var_kwargs_uniform' + + parameters_list.append( + self._generate_parameter_v3_schema( + name, annotation, AnnotationSource.FUNCTION, parameter_mode, default=p.default + ) + ) + + return core_schema.arguments_v3_schema( + parameters_list, + validate_by_name=self._config_wrapper.validate_by_name, + ) + + def _unsubstituted_typevar_schema(self, typevar: typing.TypeVar) -> core_schema.CoreSchema: + try: + has_default = typevar.has_default() + except AttributeError: + # Happens if using `typing.TypeVar` (and not `typing_extensions`) on Python < 3.13 + pass + else: + if has_default: + return self.generate_schema(typevar.__default__) + + if constraints := typevar.__constraints__: + return self._union_schema(typing.Union[constraints]) + + if bound := typevar.__bound__: + schema = self.generate_schema(bound) + schema['serialization'] = core_schema.wrap_serializer_function_ser_schema( + lambda x, h: h(x), + schema=core_schema.any_schema(), + ) + return schema + + return core_schema.any_schema() + + def _computed_field_schema( + self, + d: Decorator[ComputedFieldInfo], + field_serializers: dict[str, Decorator[FieldSerializerDecoratorInfo]], + ) -> core_schema.ComputedField: + if d.info.return_type is not PydanticUndefined: + return_type = d.info.return_type + else: + try: + # Do not pass in globals as the function could be defined in a different module. + # Instead, let `get_callable_return_type` infer the globals to use, but still pass + # in locals that may contain a parent/rebuild namespace: + return_type = _decorators.get_callable_return_type(d.func, localns=self._types_namespace.locals) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + if return_type is PydanticUndefined: + raise PydanticUserError( + 'Computed field is missing return type annotation or specifying `return_type`' + ' to the `@computed_field` decorator (e.g. `@computed_field(return_type=int | str)`)', + code='model-field-missing-annotation', + ) + + return_type = replace_types(return_type, self._typevars_map) + # Create a new ComputedFieldInfo so that different type parametrizations of the same + # generic model's computed field can have different return types. + d.info = dataclasses.replace(d.info, return_type=return_type) + return_type_schema = self.generate_schema(return_type) + # Apply serializers to computed field if there exist + return_type_schema = self._apply_field_serializers( + return_type_schema, + filter_field_decorator_info_by_field(field_serializers.values(), d.cls_var_name), + ) + + alias_generator = self._config_wrapper.alias_generator + if alias_generator is not None: + self._apply_alias_generator_to_computed_field_info( + alias_generator=alias_generator, computed_field_info=d.info, computed_field_name=d.cls_var_name + ) + self._apply_field_title_generator_to_field_info(self._config_wrapper, d.info, d.cls_var_name) + + pydantic_js_updates, pydantic_js_extra = _extract_json_schema_info_from_field_info(d.info) + core_metadata: dict[str, Any] = {} + update_core_metadata( + core_metadata, + pydantic_js_updates={'readOnly': True, **(pydantic_js_updates if pydantic_js_updates else {})}, + pydantic_js_extra=pydantic_js_extra, + ) + return core_schema.computed_field( + d.cls_var_name, return_schema=return_type_schema, alias=d.info.alias, metadata=core_metadata + ) + + def _annotated_schema(self, annotated_type: Any) -> core_schema.CoreSchema: + """Generate schema for an Annotated type, e.g. `Annotated[int, Field(...)]` or `Annotated[int, Gt(0)]`.""" + FieldInfo = import_cached_field_info() + source_type, *annotations = self._get_args_resolving_forward_refs( + annotated_type, + required=True, + ) + schema = self._apply_annotations(source_type, annotations) + # put the default validator last so that TypeAdapter.get_default_value() works + # even if there are function validators involved + for annotation in annotations: + if isinstance(annotation, FieldInfo): + schema = wrap_default(annotation, schema) + return schema + + def _apply_annotations( + self, + source_type: Any, + annotations: list[Any], + transform_inner_schema: Callable[[CoreSchema], CoreSchema] = lambda x: x, + ) -> CoreSchema: + """Apply arguments from `Annotated` or from `FieldInfo` to a schema. + + This gets called by `GenerateSchema._annotated_schema` but differs from it in that it does + not expect `source_type` to be an `Annotated` object, it expects it to be the first argument of that + (in other words, `GenerateSchema._annotated_schema` just unpacks `Annotated`, this process it). + """ + annotations = list(_known_annotated_metadata.expand_grouped_metadata(annotations)) + + pydantic_js_annotation_functions: list[GetJsonSchemaFunction] = [] + + def inner_handler(obj: Any) -> CoreSchema: + schema = self._generate_schema_from_get_schema_method(obj, source_type) + + if schema is None: + schema = self._generate_schema_inner(obj) + + metadata_js_function = _extract_get_pydantic_json_schema(obj) + if metadata_js_function is not None: + metadata_schema = resolve_original_schema(schema, self.defs) + if metadata_schema is not None: + self._add_js_function(metadata_schema, metadata_js_function) + return transform_inner_schema(schema) + + get_inner_schema = CallbackGetCoreSchemaHandler(inner_handler, self) + + for annotation in annotations: + if annotation is None: + continue + get_inner_schema = self._get_wrapped_inner_schema( + get_inner_schema, annotation, pydantic_js_annotation_functions + ) + + schema = get_inner_schema(source_type) + if pydantic_js_annotation_functions: + core_metadata = schema.setdefault('metadata', {}) + update_core_metadata(core_metadata, pydantic_js_annotation_functions=pydantic_js_annotation_functions) + return _add_custom_serialization_from_json_encoders(self._config_wrapper.json_encoders, source_type, schema) + + def _apply_single_annotation(self, schema: core_schema.CoreSchema, metadata: Any) -> core_schema.CoreSchema: + FieldInfo = import_cached_field_info() + + if isinstance(metadata, FieldInfo): + for field_metadata in metadata.metadata: + schema = self._apply_single_annotation(schema, field_metadata) + + if metadata.discriminator is not None: + schema = self._apply_discriminator_to_union(schema, metadata.discriminator) + return schema + + if schema['type'] == 'nullable': + # for nullable schemas, metadata is automatically applied to the inner schema + inner = schema.get('schema', core_schema.any_schema()) + inner = self._apply_single_annotation(inner, metadata) + if inner: + schema['schema'] = inner + return schema + + original_schema = schema + ref = schema.get('ref') + if ref is not None: + schema = schema.copy() + new_ref = ref + f'_{repr(metadata)}' + if (existing := self.defs.get_schema_from_ref(new_ref)) is not None: + return existing + schema['ref'] = new_ref # pyright: ignore[reportGeneralTypeIssues] + elif schema['type'] == 'definition-ref': + ref = schema['schema_ref'] + if (referenced_schema := self.defs.get_schema_from_ref(ref)) is not None: + schema = referenced_schema.copy() + new_ref = ref + f'_{repr(metadata)}' + if (existing := self.defs.get_schema_from_ref(new_ref)) is not None: + return existing + schema['ref'] = new_ref # pyright: ignore[reportGeneralTypeIssues] + + maybe_updated_schema = _known_annotated_metadata.apply_known_metadata(metadata, schema) + + if maybe_updated_schema is not None: + return maybe_updated_schema + return original_schema + + def _apply_single_annotation_json_schema( + self, schema: core_schema.CoreSchema, metadata: Any + ) -> core_schema.CoreSchema: + FieldInfo = import_cached_field_info() + + if isinstance(metadata, FieldInfo): + for field_metadata in metadata.metadata: + schema = self._apply_single_annotation_json_schema(schema, field_metadata) + + pydantic_js_updates, pydantic_js_extra = _extract_json_schema_info_from_field_info(metadata) + core_metadata = schema.setdefault('metadata', {}) + update_core_metadata( + core_metadata, pydantic_js_updates=pydantic_js_updates, pydantic_js_extra=pydantic_js_extra + ) + return schema + + def _get_wrapped_inner_schema( + self, + get_inner_schema: GetCoreSchemaHandler, + annotation: Any, + pydantic_js_annotation_functions: list[GetJsonSchemaFunction], + ) -> CallbackGetCoreSchemaHandler: + annotation_get_schema: GetCoreSchemaFunction | None = getattr(annotation, '__get_pydantic_core_schema__', None) + + def new_handler(source: Any) -> core_schema.CoreSchema: + if annotation_get_schema is not None: + schema = annotation_get_schema(source, get_inner_schema) + else: + schema = get_inner_schema(source) + schema = self._apply_single_annotation(schema, annotation) + schema = self._apply_single_annotation_json_schema(schema, annotation) + + metadata_js_function = _extract_get_pydantic_json_schema(annotation) + if metadata_js_function is not None: + pydantic_js_annotation_functions.append(metadata_js_function) + return schema + + return CallbackGetCoreSchemaHandler(new_handler, self) + + def _apply_field_serializers( + self, + schema: core_schema.CoreSchema, + serializers: list[Decorator[FieldSerializerDecoratorInfo]], + ) -> core_schema.CoreSchema: + """Apply field serializers to a schema.""" + if serializers: + schema = copy(schema) + if schema['type'] == 'definitions': + inner_schema = schema['schema'] + schema['schema'] = self._apply_field_serializers(inner_schema, serializers) + return schema + elif 'ref' in schema: + schema = self.defs.create_definition_reference_schema(schema) + + # use the last serializer to make it easy to override a serializer set on a parent model + serializer = serializers[-1] + is_field_serializer, info_arg = inspect_field_serializer(serializer.func, serializer.info.mode) + + if serializer.info.return_type is not PydanticUndefined: + return_type = serializer.info.return_type + else: + try: + # Do not pass in globals as the function could be defined in a different module. + # Instead, let `get_callable_return_type` infer the globals to use, but still pass + # in locals that may contain a parent/rebuild namespace: + return_type = _decorators.get_callable_return_type( + serializer.func, localns=self._types_namespace.locals + ) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + + if return_type is PydanticUndefined: + return_schema = None + else: + return_schema = self.generate_schema(return_type) + + if serializer.info.mode == 'wrap': + schema['serialization'] = core_schema.wrap_serializer_function_ser_schema( + serializer.func, + is_field_serializer=is_field_serializer, + info_arg=info_arg, + return_schema=return_schema, + when_used=serializer.info.when_used, + ) + else: + assert serializer.info.mode == 'plain' + schema['serialization'] = core_schema.plain_serializer_function_ser_schema( + serializer.func, + is_field_serializer=is_field_serializer, + info_arg=info_arg, + return_schema=return_schema, + when_used=serializer.info.when_used, + ) + return schema + + def _apply_model_serializers( + self, schema: core_schema.CoreSchema, serializers: Iterable[Decorator[ModelSerializerDecoratorInfo]] + ) -> core_schema.CoreSchema: + """Apply model serializers to a schema.""" + ref: str | None = schema.pop('ref', None) # type: ignore + if serializers: + serializer = list(serializers)[-1] + info_arg = inspect_model_serializer(serializer.func, serializer.info.mode) + + if serializer.info.return_type is not PydanticUndefined: + return_type = serializer.info.return_type + else: + try: + # Do not pass in globals as the function could be defined in a different module. + # Instead, let `get_callable_return_type` infer the globals to use, but still pass + # in locals that may contain a parent/rebuild namespace: + return_type = _decorators.get_callable_return_type( + serializer.func, localns=self._types_namespace.locals + ) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + + if return_type is PydanticUndefined: + return_schema = None + else: + return_schema = self.generate_schema(return_type) + + if serializer.info.mode == 'wrap': + ser_schema: core_schema.SerSchema = core_schema.wrap_serializer_function_ser_schema( + serializer.func, + info_arg=info_arg, + return_schema=return_schema, + when_used=serializer.info.when_used, + ) + else: + # plain + ser_schema = core_schema.plain_serializer_function_ser_schema( + serializer.func, + info_arg=info_arg, + return_schema=return_schema, + when_used=serializer.info.when_used, + ) + schema['serialization'] = ser_schema + if ref: + schema['ref'] = ref # type: ignore + return schema + + +_VALIDATOR_F_MATCH: Mapping[ + tuple[FieldValidatorModes, Literal['no-info', 'with-info']], + Callable[[Callable[..., Any], core_schema.CoreSchema, str | None], core_schema.CoreSchema], +] = { + ('before', 'no-info'): lambda f, schema, _: core_schema.no_info_before_validator_function(f, schema), + ('after', 'no-info'): lambda f, schema, _: core_schema.no_info_after_validator_function(f, schema), + ('plain', 'no-info'): lambda f, _1, _2: core_schema.no_info_plain_validator_function(f), + ('wrap', 'no-info'): lambda f, schema, _: core_schema.no_info_wrap_validator_function(f, schema), + ('before', 'with-info'): lambda f, schema, field_name: core_schema.with_info_before_validator_function( + f, schema, field_name=field_name + ), + ('after', 'with-info'): lambda f, schema, field_name: core_schema.with_info_after_validator_function( + f, schema, field_name=field_name + ), + ('plain', 'with-info'): lambda f, _, field_name: core_schema.with_info_plain_validator_function( + f, field_name=field_name + ), + ('wrap', 'with-info'): lambda f, schema, field_name: core_schema.with_info_wrap_validator_function( + f, schema, field_name=field_name + ), +} + + +# TODO V3: this function is only used for deprecated decorators. It should +# be removed once we drop support for those. +def apply_validators( + schema: core_schema.CoreSchema, + validators: Iterable[Decorator[RootValidatorDecoratorInfo]] + | Iterable[Decorator[ValidatorDecoratorInfo]] + | Iterable[Decorator[FieldValidatorDecoratorInfo]], + field_name: str | None, +) -> core_schema.CoreSchema: + """Apply validators to a schema. + + Args: + schema: The schema to apply validators on. + validators: An iterable of validators. + field_name: The name of the field if validators are being applied to a model field. + + Returns: + The updated schema. + """ + for validator in validators: + info_arg = inspect_validator(validator.func, validator.info.mode) + val_type = 'with-info' if info_arg else 'no-info' + + schema = _VALIDATOR_F_MATCH[(validator.info.mode, val_type)](validator.func, schema, field_name) + return schema + + +def _validators_require_validate_default(validators: Iterable[Decorator[ValidatorDecoratorInfo]]) -> bool: + """In v1, if any of the validators for a field had `always=True`, the default value would be validated. + + This serves as an auxiliary function for re-implementing that logic, by looping over a provided + collection of (v1-style) ValidatorDecoratorInfo's and checking if any of them have `always=True`. + + We should be able to drop this function and the associated logic calling it once we drop support + for v1-style validator decorators. (Or we can extend it and keep it if we add something equivalent + to the v1-validator `always` kwarg to `field_validator`.) + """ + for validator in validators: + if validator.info.always: + return True + return False + + +def apply_model_validators( + schema: core_schema.CoreSchema, + validators: Iterable[Decorator[ModelValidatorDecoratorInfo]], + mode: Literal['inner', 'outer', 'all'], +) -> core_schema.CoreSchema: + """Apply model validators to a schema. + + If mode == 'inner', only "before" validators are applied + If mode == 'outer', validators other than "before" are applied + If mode == 'all', all validators are applied + + Args: + schema: The schema to apply validators on. + validators: An iterable of validators. + mode: The validator mode. + + Returns: + The updated schema. + """ + ref: str | None = schema.pop('ref', None) # type: ignore + for validator in validators: + if mode == 'inner' and validator.info.mode != 'before': + continue + if mode == 'outer' and validator.info.mode == 'before': + continue + info_arg = inspect_validator(validator.func, validator.info.mode) + if validator.info.mode == 'wrap': + if info_arg: + schema = core_schema.with_info_wrap_validator_function(function=validator.func, schema=schema) + else: + schema = core_schema.no_info_wrap_validator_function(function=validator.func, schema=schema) + elif validator.info.mode == 'before': + if info_arg: + schema = core_schema.with_info_before_validator_function(function=validator.func, schema=schema) + else: + schema = core_schema.no_info_before_validator_function(function=validator.func, schema=schema) + else: + assert validator.info.mode == 'after' + if info_arg: + schema = core_schema.with_info_after_validator_function(function=validator.func, schema=schema) + else: + schema = core_schema.no_info_after_validator_function(function=validator.func, schema=schema) + if ref: + schema['ref'] = ref # type: ignore + return schema + + +def wrap_default(field_info: FieldInfo, schema: core_schema.CoreSchema) -> core_schema.CoreSchema: + """Wrap schema with default schema if default value or `default_factory` are available. + + Args: + field_info: The field info object. + schema: The schema to apply default on. + + Returns: + Updated schema by default value or `default_factory`. + """ + if field_info.default_factory: + return core_schema.with_default_schema( + schema, + default_factory=field_info.default_factory, + default_factory_takes_data=takes_validated_data_argument(field_info.default_factory), + validate_default=field_info.validate_default, + ) + elif field_info.default is not PydanticUndefined: + return core_schema.with_default_schema( + schema, default=field_info.default, validate_default=field_info.validate_default + ) + else: + return schema + + +def _extract_get_pydantic_json_schema(tp: Any) -> GetJsonSchemaFunction | None: + """Extract `__get_pydantic_json_schema__` from a type, handling the deprecated `__modify_schema__`.""" + js_modify_function = getattr(tp, '__get_pydantic_json_schema__', None) + + if hasattr(tp, '__modify_schema__'): + BaseModel = import_cached_base_model() + + has_custom_v2_modify_js_func = ( + js_modify_function is not None + and BaseModel.__get_pydantic_json_schema__.__func__ # type: ignore + not in (js_modify_function, getattr(js_modify_function, '__func__', None)) + ) + + if not has_custom_v2_modify_js_func: + cls_name = getattr(tp, '__name__', None) + raise PydanticUserError( + f'The `__modify_schema__` method is not supported in Pydantic v2. ' + f'Use `__get_pydantic_json_schema__` instead{f" in class `{cls_name}`" if cls_name else ""}.', + code='custom-json-schema', + ) + + if (origin := get_origin(tp)) is not None: + # Generic aliases proxy attribute access to the origin, *except* dunder attributes, + # such as `__get_pydantic_json_schema__`, hence the explicit check. + return _extract_get_pydantic_json_schema(origin) + + if js_modify_function is None: + return None + + return js_modify_function + + +class _CommonField(TypedDict): + schema: core_schema.CoreSchema + validation_alias: str | list[str | int] | list[list[str | int]] | None + serialization_alias: str | None + serialization_exclude: bool | None + frozen: bool | None + metadata: dict[str, Any] + + +def _common_field( + schema: core_schema.CoreSchema, + *, + validation_alias: str | list[str | int] | list[list[str | int]] | None = None, + serialization_alias: str | None = None, + serialization_exclude: bool | None = None, + frozen: bool | None = None, + metadata: Any = None, +) -> _CommonField: + return { + 'schema': schema, + 'validation_alias': validation_alias, + 'serialization_alias': serialization_alias, + 'serialization_exclude': serialization_exclude, + 'frozen': frozen, + 'metadata': metadata, + } + + +def resolve_original_schema(schema: CoreSchema, definitions: _Definitions) -> CoreSchema | None: + if schema['type'] == 'definition-ref': + return definitions.get_schema_from_ref(schema['schema_ref']) + elif schema['type'] == 'definitions': + return schema['schema'] + else: + return schema + + +def _inlining_behavior( + def_ref: core_schema.DefinitionReferenceSchema, +) -> Literal['inline', 'keep', 'preserve_metadata']: + """Determine the inlining behavior of the `'definition-ref'` schema. + + - If no `'serialization'` schema and no metadata is attached, the schema can safely be inlined. + - If it has metadata but only related to the deferred discriminator application, it can be inlined + provided that such metadata is kept. + - Otherwise, the schema should not be inlined. Doing so would remove the `'serialization'` schema or metadata. + """ + if 'serialization' in def_ref: + return 'keep' + metadata = def_ref.get('metadata') + if not metadata: + return 'inline' + if len(metadata) == 1 and 'pydantic_internal_union_discriminator' in metadata: + return 'preserve_metadata' + return 'keep' + + +class _Definitions: + """Keeps track of references and definitions.""" + + _recursively_seen: set[str] + """A set of recursively seen references. + + When a referenceable type is encountered, the `get_schema_or_ref` context manager is + entered to compute the reference. If the type references itself by some way (e.g. for + a dataclass a Pydantic model, the class can be referenced as a field annotation), + entering the context manager again will yield a `'definition-ref'` schema that should + short-circuit the normal generation process, as the reference was already in this set. + """ + + _definitions: dict[str, core_schema.CoreSchema] + """A mapping of references to their corresponding schema. + + When a schema for a referenceable type is generated, it is stored in this mapping. If the + same type is encountered again, the reference is yielded by the `get_schema_or_ref` context + manager. + """ + + def __init__(self) -> None: + self._recursively_seen = set() + self._definitions = {} + + @contextmanager + def get_schema_or_ref(self, tp: Any, /) -> Generator[tuple[str, core_schema.DefinitionReferenceSchema | None]]: + """Get a definition for `tp` if one exists. + + If a definition exists, a tuple of `(ref_string, CoreSchema)` is returned. + If no definition exists yet, a tuple of `(ref_string, None)` is returned. + + Note that the returned `CoreSchema` will always be a `DefinitionReferenceSchema`, + not the actual definition itself. + + This should be called for any type that can be identified by reference. + This includes any recursive types. + + At present the following types can be named/recursive: + + - Pydantic model + - Pydantic and stdlib dataclasses + - Typed dictionaries + - Named tuples + - `TypeAliasType` instances + - Enums + """ + ref = get_type_ref(tp) + # return the reference if we're either (1) in a cycle or (2) it the reference was already encountered: + if ref in self._recursively_seen or ref in self._definitions: + yield (ref, core_schema.definition_reference_schema(ref)) + else: + self._recursively_seen.add(ref) + try: + yield (ref, None) + finally: + self._recursively_seen.discard(ref) + + def get_schema_from_ref(self, ref: str) -> CoreSchema | None: + """Resolve the schema from the given reference.""" + return self._definitions.get(ref) + + def create_definition_reference_schema(self, schema: CoreSchema) -> core_schema.DefinitionReferenceSchema: + """Store the schema as a definition and return a `'definition-reference'` schema pointing to it. + + The schema must have a reference attached to it. + """ + ref = schema['ref'] # pyright: ignore + self._definitions[ref] = schema + return core_schema.definition_reference_schema(ref) + + def unpack_definitions(self, schema: core_schema.DefinitionsSchema) -> CoreSchema: + """Store the definitions of the `'definitions'` core schema and return the inner core schema.""" + for def_schema in schema['definitions']: + self._definitions[def_schema['ref']] = def_schema # pyright: ignore + return schema['schema'] + + def finalize_schema(self, schema: CoreSchema) -> CoreSchema: + """Finalize the core schema. + + This traverses the core schema and referenced definitions, replaces `'definition-ref'` schemas + by the referenced definition if possible, and applies deferred discriminators. + """ + definitions = self._definitions + try: + gather_result = gather_schemas_for_cleaning( + schema, + definitions=definitions, + ) + except MissingDefinitionError as e: + raise InvalidSchemaError from e + + remaining_defs: dict[str, CoreSchema] = {} + + # Note: this logic doesn't play well when core schemas with deferred discriminator metadata + # and references are encountered. See the `test_deferred_discriminated_union_and_references()` test. + for ref, inlinable_def_ref in gather_result['collected_references'].items(): + if inlinable_def_ref is not None and (inlining_behavior := _inlining_behavior(inlinable_def_ref)) != 'keep': + if inlining_behavior == 'inline': + # `ref` was encountered, and only once: + # - `inlinable_def_ref` is a `'definition-ref'` schema and is guaranteed to be + # the only one. Transform it into the definition it points to. + # - Do not store the definition in the `remaining_defs`. + inlinable_def_ref.clear() # pyright: ignore[reportAttributeAccessIssue] + inlinable_def_ref.update(self._resolve_definition(ref, definitions)) # pyright: ignore + elif inlining_behavior == 'preserve_metadata': + # `ref` was encountered, and only once, but contains discriminator metadata. + # We will do the same thing as if `inlining_behavior` was `'inline'`, but make + # sure to keep the metadata for the deferred discriminator application logic below. + meta = inlinable_def_ref.pop('metadata') + inlinable_def_ref.clear() # pyright: ignore[reportAttributeAccessIssue] + inlinable_def_ref.update(self._resolve_definition(ref, definitions)) # pyright: ignore + inlinable_def_ref['metadata'] = meta + else: + # `ref` was encountered, at least two times (or only once, but with metadata or a serialization schema): + # - Do not inline the `'definition-ref'` schemas (they are not provided in the gather result anyway). + # - Store the the definition in the `remaining_defs` + remaining_defs[ref] = self._resolve_definition(ref, definitions) + + for cs in gather_result['deferred_discriminator_schemas']: + discriminator: str | None = cs['metadata'].pop('pydantic_internal_union_discriminator', None) # pyright: ignore[reportTypedDictNotRequiredAccess] + if discriminator is None: + # This can happen in rare scenarios, when a deferred schema is present multiple times in the + # gather result (e.g. when using the `Sequence` type -- see `test_sequence_discriminated_union()`). + # In this case, a previous loop iteration applied the discriminator and so we can just skip it here. + continue + applied = _discriminated_union.apply_discriminator(cs.copy(), discriminator, remaining_defs) + # Mutate the schema directly to have the discriminator applied + cs.clear() # pyright: ignore[reportAttributeAccessIssue] + cs.update(applied) # pyright: ignore + + if remaining_defs: + schema = core_schema.definitions_schema(schema=schema, definitions=[*remaining_defs.values()]) + return schema + + def _resolve_definition(self, ref: str, definitions: dict[str, CoreSchema]) -> CoreSchema: + definition = definitions[ref] + if definition['type'] != 'definition-ref': + return definition + + # Some `'definition-ref'` schemas might act as "intermediate" references (e.g. when using + # a PEP 695 type alias (which is referenceable) that references another PEP 695 type alias): + visited: set[str] = set() + while definition['type'] == 'definition-ref' and _inlining_behavior(definition) == 'inline': + schema_ref = definition['schema_ref'] + if schema_ref in visited: + raise PydanticUserError( + f'{ref} contains a circular reference to itself.', code='circular-reference-schema' + ) + visited.add(schema_ref) + definition = definitions[schema_ref] + return {**definition, 'ref': ref} # pyright: ignore[reportReturnType] + + +class _FieldNameStack: + __slots__ = ('_stack',) + + def __init__(self) -> None: + self._stack: list[str] = [] + + @contextmanager + def push(self, field_name: str) -> Iterator[None]: + self._stack.append(field_name) + yield + self._stack.pop() + + def get(self) -> str | None: + if self._stack: + return self._stack[-1] + else: + return None + + +class _ModelTypeStack: + __slots__ = ('_stack',) + + def __init__(self) -> None: + self._stack: list[type] = [] + + @contextmanager + def push(self, type_obj: type) -> Iterator[None]: + self._stack.append(type_obj) + yield + self._stack.pop() + + def get(self) -> type | None: + if self._stack: + return self._stack[-1] + else: + return None diff --git a/lib/python3.12/site-packages/pydantic/_internal/_generics.py b/lib/python3.12/site-packages/pydantic/_internal/_generics.py new file mode 100644 index 0000000000000000000000000000000000000000..80136766612aefce6293c550acd3044434f2b277 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_generics.py @@ -0,0 +1,547 @@ +from __future__ import annotations + +import sys +import types +import typing +from collections import ChainMap +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar +from itertools import zip_longest +from types import prepare_class +from typing import TYPE_CHECKING, Annotated, Any, TypeVar +from weakref import WeakValueDictionary + +import typing_extensions +from typing_inspection import typing_objects +from typing_inspection.introspection import is_union_origin + +from . import _typing_extra +from ._core_utils import get_type_ref +from ._forward_ref import PydanticRecursiveRef +from ._utils import all_identical, is_model_class + +if sys.version_info >= (3, 10): + from typing import _UnionGenericAlias # type: ignore[attr-defined] + +if TYPE_CHECKING: + from ..main import BaseModel + +GenericTypesCacheKey = tuple[Any, Any, tuple[Any, ...]] + +# Note: We want to remove LimitedDict, but to do this, we'd need to improve the handling of generics caching. +# Right now, to handle recursive generics, we some types must remain cached for brief periods without references. +# By chaining the WeakValuesDict with a LimitedDict, we have a way to retain caching for all types with references, +# while also retaining a limited number of types even without references. This is generally enough to build +# specific recursive generic models without losing required items out of the cache. + +KT = TypeVar('KT') +VT = TypeVar('VT') +_LIMITED_DICT_SIZE = 100 + + +class LimitedDict(dict[KT, VT]): + def __init__(self, size_limit: int = _LIMITED_DICT_SIZE) -> None: + self.size_limit = size_limit + super().__init__() + + def __setitem__(self, key: KT, value: VT, /) -> None: + super().__setitem__(key, value) + if len(self) > self.size_limit: + excess = len(self) - self.size_limit + self.size_limit // 10 + to_remove = list(self.keys())[:excess] + for k in to_remove: + del self[k] + + +# weak dictionaries allow the dynamically created parametrized versions of generic models to get collected +# once they are no longer referenced by the caller. +GenericTypesCache = WeakValueDictionary[GenericTypesCacheKey, 'type[BaseModel]'] + +if TYPE_CHECKING: + + class DeepChainMap(ChainMap[KT, VT]): # type: ignore + ... + +else: + + class DeepChainMap(ChainMap): + """Variant of ChainMap that allows direct updates to inner scopes. + + Taken from https://docs.python.org/3/library/collections.html#collections.ChainMap, + with some light modifications for this use case. + """ + + def clear(self) -> None: + for mapping in self.maps: + mapping.clear() + + def __setitem__(self, key: KT, value: VT) -> None: + for mapping in self.maps: + mapping[key] = value + + def __delitem__(self, key: KT) -> None: + hit = False + for mapping in self.maps: + if key in mapping: + del mapping[key] + hit = True + if not hit: + raise KeyError(key) + + +# Despite the fact that LimitedDict _seems_ no longer necessary, I'm very nervous to actually remove it +# and discover later on that we need to re-add all this infrastructure... +# _GENERIC_TYPES_CACHE = DeepChainMap(GenericTypesCache(), LimitedDict()) + +_GENERIC_TYPES_CACHE: ContextVar[GenericTypesCache | None] = ContextVar('_GENERIC_TYPES_CACHE', default=None) + + +class PydanticGenericMetadata(typing_extensions.TypedDict): + origin: type[BaseModel] | None # analogous to typing._GenericAlias.__origin__ + args: tuple[Any, ...] # analogous to typing._GenericAlias.__args__ + parameters: tuple[TypeVar, ...] # analogous to typing.Generic.__parameters__ + + +def create_generic_submodel( + model_name: str, origin: type[BaseModel], args: tuple[Any, ...], params: tuple[Any, ...] +) -> type[BaseModel]: + """Dynamically create a submodel of a provided (generic) BaseModel. + + This is used when producing concrete parametrizations of generic models. This function + only *creates* the new subclass; the schema/validators/serialization must be updated to + reflect a concrete parametrization elsewhere. + + Args: + model_name: The name of the newly created model. + origin: The base class for the new model to inherit from. + args: A tuple of generic metadata arguments. + params: A tuple of generic metadata parameters. + + Returns: + The created submodel. + """ + namespace: dict[str, Any] = {'__module__': origin.__module__} + bases = (origin,) + meta, ns, kwds = prepare_class(model_name, bases) + namespace.update(ns) + created_model = meta( + model_name, + bases, + namespace, + __pydantic_generic_metadata__={ + 'origin': origin, + 'args': args, + 'parameters': params, + }, + __pydantic_reset_parent_namespace__=False, + **kwds, + ) + + model_module, called_globally = _get_caller_frame_info(depth=3) + if called_globally: # create global reference and therefore allow pickling + object_by_reference = None + reference_name = model_name + reference_module_globals = sys.modules[created_model.__module__].__dict__ + while object_by_reference is not created_model: + object_by_reference = reference_module_globals.setdefault(reference_name, created_model) + reference_name += '_' + + return created_model + + +def _get_caller_frame_info(depth: int = 2) -> tuple[str | None, bool]: + """Used inside a function to check whether it was called globally. + + Args: + depth: The depth to get the frame. + + Returns: + A tuple contains `module_name` and `called_globally`. + + Raises: + RuntimeError: If the function is not called inside a function. + """ + try: + previous_caller_frame = sys._getframe(depth) + except ValueError as e: + raise RuntimeError('This function must be used inside another function') from e + except AttributeError: # sys module does not have _getframe function, so there's nothing we can do about it + return None, False + frame_globals = previous_caller_frame.f_globals + return frame_globals.get('__name__'), previous_caller_frame.f_locals is frame_globals + + +DictValues: type[Any] = {}.values().__class__ + + +def iter_contained_typevars(v: Any) -> Iterator[TypeVar]: + """Recursively iterate through all subtypes and type args of `v` and yield any typevars that are found. + + This is inspired as an alternative to directly accessing the `__parameters__` attribute of a GenericAlias, + since __parameters__ of (nested) generic BaseModel subclasses won't show up in that list. + """ + if isinstance(v, TypeVar): + yield v + elif is_model_class(v): + yield from v.__pydantic_generic_metadata__['parameters'] + elif isinstance(v, (DictValues, list)): + for var in v: + yield from iter_contained_typevars(var) + else: + args = get_args(v) + for arg in args: + yield from iter_contained_typevars(arg) + + +def get_args(v: Any) -> Any: + pydantic_generic_metadata: PydanticGenericMetadata | None = getattr(v, '__pydantic_generic_metadata__', None) + if pydantic_generic_metadata: + return pydantic_generic_metadata.get('args') + return typing_extensions.get_args(v) + + +def get_origin(v: Any) -> Any: + pydantic_generic_metadata: PydanticGenericMetadata | None = getattr(v, '__pydantic_generic_metadata__', None) + if pydantic_generic_metadata: + return pydantic_generic_metadata.get('origin') + return typing_extensions.get_origin(v) + + +def get_standard_typevars_map(cls: Any) -> dict[TypeVar, Any] | None: + """Package a generic type's typevars and parametrization (if present) into a dictionary compatible with the + `replace_types` function. Specifically, this works with standard typing generics and typing._GenericAlias. + """ + origin = get_origin(cls) + if origin is None: + return None + if not hasattr(origin, '__parameters__'): + return None + + # In this case, we know that cls is a _GenericAlias, and origin is the generic type + # So it is safe to access cls.__args__ and origin.__parameters__ + args: tuple[Any, ...] = cls.__args__ # type: ignore + parameters: tuple[TypeVar, ...] = origin.__parameters__ + return dict(zip(parameters, args)) + + +def get_model_typevars_map(cls: type[BaseModel]) -> dict[TypeVar, Any]: + """Package a generic BaseModel's typevars and concrete parametrization (if present) into a dictionary compatible + with the `replace_types` function. + + Since BaseModel.__class_getitem__ does not produce a typing._GenericAlias, and the BaseModel generic info is + stored in the __pydantic_generic_metadata__ attribute, we need special handling here. + """ + # TODO: This could be unified with `get_standard_typevars_map` if we stored the generic metadata + # in the __origin__, __args__, and __parameters__ attributes of the model. + generic_metadata = cls.__pydantic_generic_metadata__ + origin = generic_metadata['origin'] + args = generic_metadata['args'] + if not args: + # No need to go into `iter_contained_typevars`: + return {} + return dict(zip(iter_contained_typevars(origin), args)) + + +def replace_types(type_: Any, type_map: Mapping[TypeVar, Any] | None) -> Any: + """Return type with all occurrences of `type_map` keys recursively replaced with their values. + + Args: + type_: The class or generic alias. + type_map: Mapping from `TypeVar` instance to concrete types. + + Returns: + A new type representing the basic structure of `type_` with all + `typevar_map` keys recursively replaced. + + Example: + ```python + from typing import List, Union + + from pydantic._internal._generics import replace_types + + replace_types(tuple[str, Union[List[str], float]], {str: int}) + #> tuple[int, Union[List[int], float]] + ``` + """ + if not type_map: + return type_ + + type_args = get_args(type_) + origin_type = get_origin(type_) + + if typing_objects.is_annotated(origin_type): + annotated_type, *annotations = type_args + annotated_type = replace_types(annotated_type, type_map) + # TODO remove parentheses when we drop support for Python 3.10: + return Annotated[(annotated_type, *annotations)] + + # Having type args is a good indicator that this is a typing special form + # instance or a generic alias of some sort. + if type_args: + resolved_type_args = tuple(replace_types(arg, type_map) for arg in type_args) + if all_identical(type_args, resolved_type_args): + # If all arguments are the same, there is no need to modify the + # type or create a new object at all + return type_ + + if ( + origin_type is not None + and isinstance(type_, _typing_extra.typing_base) + and not isinstance(origin_type, _typing_extra.typing_base) + and getattr(type_, '_name', None) is not None + ): + # In python < 3.9 generic aliases don't exist so any of these like `list`, + # `type` or `collections.abc.Callable` need to be translated. + # See: https://www.python.org/dev/peps/pep-0585 + origin_type = getattr(typing, type_._name) + assert origin_type is not None + + if is_union_origin(origin_type): + if any(typing_objects.is_any(arg) for arg in resolved_type_args): + # `Any | T` ~ `Any`: + resolved_type_args = (Any,) + # `Never | T` ~ `T`: + resolved_type_args = tuple( + arg + for arg in resolved_type_args + if not (typing_objects.is_noreturn(arg) or typing_objects.is_never(arg)) + ) + + # PEP-604 syntax (Ex.: list | str) is represented with a types.UnionType object that does not have __getitem__. + # We also cannot use isinstance() since we have to compare types. + if sys.version_info >= (3, 10) and origin_type is types.UnionType: + return _UnionGenericAlias(origin_type, resolved_type_args) + # NotRequired[T] and Required[T] don't support tuple type resolved_type_args, hence the condition below + return origin_type[resolved_type_args[0] if len(resolved_type_args) == 1 else resolved_type_args] + + # We handle pydantic generic models separately as they don't have the same + # semantics as "typing" classes or generic aliases + + if not origin_type and is_model_class(type_): + parameters = type_.__pydantic_generic_metadata__['parameters'] + if not parameters: + return type_ + resolved_type_args = tuple(replace_types(t, type_map) for t in parameters) + if all_identical(parameters, resolved_type_args): + return type_ + return type_[resolved_type_args] + + # Handle special case for typehints that can have lists as arguments. + # `typing.Callable[[int, str], int]` is an example for this. + if isinstance(type_, list): + resolved_list = [replace_types(element, type_map) for element in type_] + if all_identical(type_, resolved_list): + return type_ + return resolved_list + + # If all else fails, we try to resolve the type directly and otherwise just + # return the input with no modifications. + return type_map.get(type_, type_) + + +def map_generic_model_arguments(cls: type[BaseModel], args: tuple[Any, ...]) -> dict[TypeVar, Any]: + """Return a mapping between the parameters of a generic model and the provided arguments during parameterization. + + Raises: + TypeError: If the number of arguments does not match the parameters (i.e. if providing too few or too many arguments). + + Example: + ```python {test="skip" lint="skip"} + class Model[T, U, V = int](BaseModel): ... + + map_generic_model_arguments(Model, (str, bytes)) + #> {T: str, U: bytes, V: int} + + map_generic_model_arguments(Model, (str,)) + #> TypeError: Too few arguments for ; actual 1, expected at least 2 + + map_generic_model_arguments(Model, (str, bytes, int, complex)) + #> TypeError: Too many arguments for ; actual 4, expected 3 + ``` + + Note: + This function is analogous to the private `typing._check_generic_specialization` function. + """ + parameters = cls.__pydantic_generic_metadata__['parameters'] + expected_len = len(parameters) + typevars_map: dict[TypeVar, Any] = {} + + _missing = object() + for parameter, argument in zip_longest(parameters, args, fillvalue=_missing): + if parameter is _missing: + raise TypeError(f'Too many arguments for {cls}; actual {len(args)}, expected {expected_len}') + + if argument is _missing: + param = typing.cast(TypeVar, parameter) + try: + has_default = param.has_default() + except AttributeError: + # Happens if using `typing.TypeVar` (and not `typing_extensions`) on Python < 3.13. + has_default = False + if has_default: + # The default might refer to other type parameters. For an example, see: + # https://typing.readthedocs.io/en/latest/spec/generics.html#type-parameters-as-parameters-to-generics + typevars_map[param] = replace_types(param.__default__, typevars_map) + else: + expected_len -= sum(hasattr(p, 'has_default') and p.has_default() for p in parameters) + raise TypeError(f'Too few arguments for {cls}; actual {len(args)}, expected at least {expected_len}') + else: + param = typing.cast(TypeVar, parameter) + typevars_map[param] = argument + + return typevars_map + + +_generic_recursion_cache: ContextVar[set[str] | None] = ContextVar('_generic_recursion_cache', default=None) + + +@contextmanager +def generic_recursion_self_type( + origin: type[BaseModel], args: tuple[Any, ...] +) -> Iterator[PydanticRecursiveRef | None]: + """This contextmanager should be placed around the recursive calls used to build a generic type, + and accept as arguments the generic origin type and the type arguments being passed to it. + + If the same origin and arguments are observed twice, it implies that a self-reference placeholder + can be used while building the core schema, and will produce a schema_ref that will be valid in the + final parent schema. + """ + previously_seen_type_refs = _generic_recursion_cache.get() + if previously_seen_type_refs is None: + previously_seen_type_refs = set() + token = _generic_recursion_cache.set(previously_seen_type_refs) + else: + token = None + + try: + type_ref = get_type_ref(origin, args_override=args) + if type_ref in previously_seen_type_refs: + self_type = PydanticRecursiveRef(type_ref=type_ref) + yield self_type + else: + previously_seen_type_refs.add(type_ref) + yield + previously_seen_type_refs.remove(type_ref) + finally: + if token: + _generic_recursion_cache.reset(token) + + +def recursively_defined_type_refs() -> set[str]: + visited = _generic_recursion_cache.get() + if not visited: + return set() # not in a generic recursion, so there are no types + + return visited.copy() # don't allow modifications + + +def get_cached_generic_type_early(parent: type[BaseModel], typevar_values: Any) -> type[BaseModel] | None: + """The use of a two-stage cache lookup approach was necessary to have the highest performance possible for + repeated calls to `__class_getitem__` on generic types (which may happen in tighter loops during runtime), + while still ensuring that certain alternative parametrizations ultimately resolve to the same type. + + As a concrete example, this approach was necessary to make Model[List[T]][int] equal to Model[List[int]]. + The approach could be modified to not use two different cache keys at different points, but the + _early_cache_key is optimized to be as quick to compute as possible (for repeated-access speed), and the + _late_cache_key is optimized to be as "correct" as possible, so that two types that will ultimately be the + same after resolving the type arguments will always produce cache hits. + + If we wanted to move to only using a single cache key per type, we would either need to always use the + slower/more computationally intensive logic associated with _late_cache_key, or would need to accept + that Model[List[T]][int] is a different type than Model[List[T]][int]. Because we rely on subclass relationships + during validation, I think it is worthwhile to ensure that types that are functionally equivalent are actually + equal. + """ + generic_types_cache = _GENERIC_TYPES_CACHE.get() + if generic_types_cache is None: + generic_types_cache = GenericTypesCache() + _GENERIC_TYPES_CACHE.set(generic_types_cache) + return generic_types_cache.get(_early_cache_key(parent, typevar_values)) + + +def get_cached_generic_type_late( + parent: type[BaseModel], typevar_values: Any, origin: type[BaseModel], args: tuple[Any, ...] +) -> type[BaseModel] | None: + """See the docstring of `get_cached_generic_type_early` for more information about the two-stage cache lookup.""" + generic_types_cache = _GENERIC_TYPES_CACHE.get() + if ( + generic_types_cache is None + ): # pragma: no cover (early cache is guaranteed to run first and initialize the cache) + generic_types_cache = GenericTypesCache() + _GENERIC_TYPES_CACHE.set(generic_types_cache) + cached = generic_types_cache.get(_late_cache_key(origin, args, typevar_values)) + if cached is not None: + set_cached_generic_type(parent, typevar_values, cached, origin, args) + return cached + + +def set_cached_generic_type( + parent: type[BaseModel], + typevar_values: tuple[Any, ...], + type_: type[BaseModel], + origin: type[BaseModel] | None = None, + args: tuple[Any, ...] | None = None, +) -> None: + """See the docstring of `get_cached_generic_type_early` for more information about why items are cached with + two different keys. + """ + generic_types_cache = _GENERIC_TYPES_CACHE.get() + if ( + generic_types_cache is None + ): # pragma: no cover (cache lookup is guaranteed to run first and initialize the cache) + generic_types_cache = GenericTypesCache() + _GENERIC_TYPES_CACHE.set(generic_types_cache) + generic_types_cache[_early_cache_key(parent, typevar_values)] = type_ + if len(typevar_values) == 1: + generic_types_cache[_early_cache_key(parent, typevar_values[0])] = type_ + if origin and args: + generic_types_cache[_late_cache_key(origin, args, typevar_values)] = type_ + + +def _union_orderings_key(typevar_values: Any) -> Any: + """This is intended to help differentiate between Union types with the same arguments in different order. + + Thanks to caching internal to the `typing` module, it is not possible to distinguish between + List[Union[int, float]] and List[Union[float, int]] (and similarly for other "parent" origins besides List) + because `typing` considers Union[int, float] to be equal to Union[float, int]. + + However, you _can_ distinguish between (top-level) Union[int, float] vs. Union[float, int]. + Because we parse items as the first Union type that is successful, we get slightly more consistent behavior + if we make an effort to distinguish the ordering of items in a union. It would be best if we could _always_ + get the exact-correct order of items in the union, but that would require a change to the `typing` module itself. + (See https://github.com/python/cpython/issues/86483 for reference.) + """ + if isinstance(typevar_values, tuple): + args_data = [] + for value in typevar_values: + args_data.append(_union_orderings_key(value)) + return tuple(args_data) + elif typing_objects.is_union(typing_extensions.get_origin(typevar_values)): + return get_args(typevar_values) + else: + return () + + +def _early_cache_key(cls: type[BaseModel], typevar_values: Any) -> GenericTypesCacheKey: + """This is intended for minimal computational overhead during lookups of cached types. + + Note that this is overly simplistic, and it's possible that two different cls/typevar_values + inputs would ultimately result in the same type being created in BaseModel.__class_getitem__. + To handle this, we have a fallback _late_cache_key that is checked later if the _early_cache_key + lookup fails, and should result in a cache hit _precisely_ when the inputs to __class_getitem__ + would result in the same type. + """ + return cls, typevar_values, _union_orderings_key(typevar_values) + + +def _late_cache_key(origin: type[BaseModel], args: tuple[Any, ...], typevar_values: Any) -> GenericTypesCacheKey: + """This is intended for use later in the process of creating a new type, when we have more information + about the exact args that will be passed. If it turns out that a different set of inputs to + __class_getitem__ resulted in the same inputs to the generic type creation process, we can still + return the cached type, and update the cache with the _early_cache_key as well. + """ + # The _union_orderings_key is placed at the start here to ensure there cannot be a collision with an + # _early_cache_key, as that function will always produce a BaseModel subclass as the first item in the key, + # whereas this function will always produce a tuple as the first item in the key. + return _union_orderings_key(typevar_values), origin, args diff --git a/lib/python3.12/site-packages/pydantic/_internal/_git.py b/lib/python3.12/site-packages/pydantic/_internal/_git.py new file mode 100644 index 0000000000000000000000000000000000000000..96dcda28f801cd19989838b029f6f2b250dedc52 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_git.py @@ -0,0 +1,27 @@ +"""Git utilities, adopted from mypy's git utilities (https://github.com/python/mypy/blob/master/mypy/git.py).""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + + +def is_git_repo(dir: Path) -> bool: + """Is the given directory version-controlled with git?""" + return dir.joinpath('.git').exists() + + +def have_git() -> bool: # pragma: no cover + """Can we run the git executable?""" + try: + subprocess.check_output(['git', '--help']) + return True + except subprocess.CalledProcessError: + return False + except OSError: + return False + + +def git_revision(dir: Path) -> str: + """Get the SHA-1 of the HEAD of a git repository.""" + return subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD'], cwd=dir).decode('utf-8').strip() diff --git a/lib/python3.12/site-packages/pydantic/_internal/_import_utils.py b/lib/python3.12/site-packages/pydantic/_internal/_import_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..638102f771dd674799255cc6ef69cacb2c184d13 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_import_utils.py @@ -0,0 +1,20 @@ +from functools import cache +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pydantic import BaseModel + from pydantic.fields import FieldInfo + + +@cache +def import_cached_base_model() -> type['BaseModel']: + from pydantic import BaseModel + + return BaseModel + + +@cache +def import_cached_field_info() -> type['FieldInfo']: + from pydantic.fields import FieldInfo + + return FieldInfo diff --git a/lib/python3.12/site-packages/pydantic/_internal/_internal_dataclass.py b/lib/python3.12/site-packages/pydantic/_internal/_internal_dataclass.py new file mode 100644 index 0000000000000000000000000000000000000000..33e152cc8d178486ef016285855495a86d5991d8 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_internal_dataclass.py @@ -0,0 +1,7 @@ +import sys + +# `slots` is available on Python >= 3.10 +if sys.version_info >= (3, 10): + slots_true = {'slots': True} +else: + slots_true = {} diff --git a/lib/python3.12/site-packages/pydantic/_internal/_known_annotated_metadata.py b/lib/python3.12/site-packages/pydantic/_internal/_known_annotated_metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..c127e27ac16af910d03b4d92934a89b40cb53d3f --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_known_annotated_metadata.py @@ -0,0 +1,393 @@ +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Iterable +from copy import copy +from functools import lru_cache, partial +from typing import TYPE_CHECKING, Any + +from pydantic_core import CoreSchema, PydanticCustomError, ValidationError, to_jsonable_python +from pydantic_core import core_schema as cs + +from ._fields import PydanticMetadata +from ._import_utils import import_cached_field_info + +if TYPE_CHECKING: + pass + +STRICT = {'strict'} +FAIL_FAST = {'fail_fast'} +LENGTH_CONSTRAINTS = {'min_length', 'max_length'} +INEQUALITY = {'le', 'ge', 'lt', 'gt'} +NUMERIC_CONSTRAINTS = {'multiple_of', *INEQUALITY} +ALLOW_INF_NAN = {'allow_inf_nan'} + +STR_CONSTRAINTS = { + *LENGTH_CONSTRAINTS, + *STRICT, + 'strip_whitespace', + 'to_lower', + 'to_upper', + 'pattern', + 'coerce_numbers_to_str', +} +BYTES_CONSTRAINTS = {*LENGTH_CONSTRAINTS, *STRICT} + +LIST_CONSTRAINTS = {*LENGTH_CONSTRAINTS, *STRICT, *FAIL_FAST} +TUPLE_CONSTRAINTS = {*LENGTH_CONSTRAINTS, *STRICT, *FAIL_FAST} +SET_CONSTRAINTS = {*LENGTH_CONSTRAINTS, *STRICT, *FAIL_FAST} +DICT_CONSTRAINTS = {*LENGTH_CONSTRAINTS, *STRICT} +GENERATOR_CONSTRAINTS = {*LENGTH_CONSTRAINTS, *STRICT} +SEQUENCE_CONSTRAINTS = {*LENGTH_CONSTRAINTS, *FAIL_FAST} + +FLOAT_CONSTRAINTS = {*NUMERIC_CONSTRAINTS, *ALLOW_INF_NAN, *STRICT} +DECIMAL_CONSTRAINTS = {'max_digits', 'decimal_places', *FLOAT_CONSTRAINTS} +INT_CONSTRAINTS = {*NUMERIC_CONSTRAINTS, *ALLOW_INF_NAN, *STRICT} +BOOL_CONSTRAINTS = STRICT +UUID_CONSTRAINTS = STRICT + +DATE_TIME_CONSTRAINTS = {*NUMERIC_CONSTRAINTS, *STRICT} +TIMEDELTA_CONSTRAINTS = {*NUMERIC_CONSTRAINTS, *STRICT} +TIME_CONSTRAINTS = {*NUMERIC_CONSTRAINTS, *STRICT} +LAX_OR_STRICT_CONSTRAINTS = STRICT +ENUM_CONSTRAINTS = STRICT +COMPLEX_CONSTRAINTS = STRICT + +UNION_CONSTRAINTS = {'union_mode'} +URL_CONSTRAINTS = { + 'max_length', + 'allowed_schemes', + 'host_required', + 'default_host', + 'default_port', + 'default_path', +} + +TEXT_SCHEMA_TYPES = ('str', 'bytes', 'url', 'multi-host-url') +SEQUENCE_SCHEMA_TYPES = ('list', 'tuple', 'set', 'frozenset', 'generator', *TEXT_SCHEMA_TYPES) +NUMERIC_SCHEMA_TYPES = ('float', 'int', 'date', 'time', 'timedelta', 'datetime') + +CONSTRAINTS_TO_ALLOWED_SCHEMAS: dict[str, set[str]] = defaultdict(set) + +constraint_schema_pairings: list[tuple[set[str], tuple[str, ...]]] = [ + (STR_CONSTRAINTS, TEXT_SCHEMA_TYPES), + (BYTES_CONSTRAINTS, ('bytes',)), + (LIST_CONSTRAINTS, ('list',)), + (TUPLE_CONSTRAINTS, ('tuple',)), + (SET_CONSTRAINTS, ('set', 'frozenset')), + (DICT_CONSTRAINTS, ('dict',)), + (GENERATOR_CONSTRAINTS, ('generator',)), + (FLOAT_CONSTRAINTS, ('float',)), + (INT_CONSTRAINTS, ('int',)), + (DATE_TIME_CONSTRAINTS, ('date', 'time', 'datetime', 'timedelta')), + # TODO: this is a bit redundant, we could probably avoid some of these + (STRICT, (*TEXT_SCHEMA_TYPES, *SEQUENCE_SCHEMA_TYPES, *NUMERIC_SCHEMA_TYPES, 'typed-dict', 'model')), + (UNION_CONSTRAINTS, ('union',)), + (URL_CONSTRAINTS, ('url', 'multi-host-url')), + (BOOL_CONSTRAINTS, ('bool',)), + (UUID_CONSTRAINTS, ('uuid',)), + (LAX_OR_STRICT_CONSTRAINTS, ('lax-or-strict',)), + (ENUM_CONSTRAINTS, ('enum',)), + (DECIMAL_CONSTRAINTS, ('decimal',)), + (COMPLEX_CONSTRAINTS, ('complex',)), +] + +for constraints, schemas in constraint_schema_pairings: + for c in constraints: + CONSTRAINTS_TO_ALLOWED_SCHEMAS[c].update(schemas) + + +def as_jsonable_value(v: Any) -> Any: + if type(v) not in (int, str, float, bytes, bool, type(None)): + return to_jsonable_python(v) + return v + + +def expand_grouped_metadata(annotations: Iterable[Any]) -> Iterable[Any]: + """Expand the annotations. + + Args: + annotations: An iterable of annotations. + + Returns: + An iterable of expanded annotations. + + Example: + ```python + from annotated_types import Ge, Len + + from pydantic._internal._known_annotated_metadata import expand_grouped_metadata + + print(list(expand_grouped_metadata([Ge(4), Len(5)]))) + #> [Ge(ge=4), MinLen(min_length=5)] + ``` + """ + import annotated_types as at + + FieldInfo = import_cached_field_info() + + for annotation in annotations: + if isinstance(annotation, at.GroupedMetadata): + yield from annotation + elif isinstance(annotation, FieldInfo): + yield from annotation.metadata + # this is a bit problematic in that it results in duplicate metadata + # all of our "consumers" can handle it, but it is not ideal + # we probably should split up FieldInfo into: + # - annotated types metadata + # - individual metadata known only to Pydantic + annotation = copy(annotation) + annotation.metadata = [] + yield annotation + else: + yield annotation + + +@lru_cache +def _get_at_to_constraint_map() -> dict[type, str]: + """Return a mapping of annotated types to constraints. + + Normally, we would define a mapping like this in the module scope, but we can't do that + because we don't permit module level imports of `annotated_types`, in an attempt to speed up + the import time of `pydantic`. We still only want to have this dictionary defined in one place, + so we use this function to cache the result. + """ + import annotated_types as at + + return { + at.Gt: 'gt', + at.Ge: 'ge', + at.Lt: 'lt', + at.Le: 'le', + at.MultipleOf: 'multiple_of', + at.MinLen: 'min_length', + at.MaxLen: 'max_length', + } + + +def apply_known_metadata(annotation: Any, schema: CoreSchema) -> CoreSchema | None: # noqa: C901 + """Apply `annotation` to `schema` if it is an annotation we know about (Gt, Le, etc.). + Otherwise return `None`. + + This does not handle all known annotations. If / when it does, it can always + return a CoreSchema and return the unmodified schema if the annotation should be ignored. + + Assumes that GroupedMetadata has already been expanded via `expand_grouped_metadata`. + + Args: + annotation: The annotation. + schema: The schema. + + Returns: + An updated schema with annotation if it is an annotation we know about, `None` otherwise. + + Raises: + PydanticCustomError: If `Predicate` fails. + """ + import annotated_types as at + + from ._validators import NUMERIC_VALIDATOR_LOOKUP, forbid_inf_nan_check + + schema = schema.copy() + schema_update, other_metadata = collect_known_metadata([annotation]) + schema_type = schema['type'] + + chain_schema_constraints: set[str] = { + 'pattern', + 'strip_whitespace', + 'to_lower', + 'to_upper', + 'coerce_numbers_to_str', + } + chain_schema_steps: list[CoreSchema] = [] + + for constraint, value in schema_update.items(): + if constraint not in CONSTRAINTS_TO_ALLOWED_SCHEMAS: + raise ValueError(f'Unknown constraint {constraint}') + allowed_schemas = CONSTRAINTS_TO_ALLOWED_SCHEMAS[constraint] + + # if it becomes necessary to handle more than one constraint + # in this recursive case with function-after or function-wrap, we should refactor + # this is a bit challenging because we sometimes want to apply constraints to the inner schema, + # whereas other times we want to wrap the existing schema with a new one that enforces a new constraint. + if schema_type in {'function-before', 'function-wrap', 'function-after'} and constraint == 'strict': + schema['schema'] = apply_known_metadata(annotation, schema['schema']) # type: ignore # schema is function schema + return schema + + # if we're allowed to apply constraint directly to the schema, like le to int, do that + if schema_type in allowed_schemas: + if constraint == 'union_mode' and schema_type == 'union': + schema['mode'] = value # type: ignore # schema is UnionSchema + else: + schema[constraint] = value + continue + + # else, apply a function after validator to the schema to enforce the corresponding constraint + if constraint in chain_schema_constraints: + + def _apply_constraint_with_incompatibility_info( + value: Any, handler: cs.ValidatorFunctionWrapHandler + ) -> Any: + try: + x = handler(value) + except ValidationError as ve: + # if the error is about the type, it's likely that the constraint is incompatible the type of the field + # for example, the following invalid schema wouldn't be caught during schema build, but rather at this point + # with a cryptic 'string_type' error coming from the string validator, + # that we'd rather express as a constraint incompatibility error (TypeError) + # Annotated[list[int], Field(pattern='abc')] + if 'type' in ve.errors()[0]['type']: + raise TypeError( + f"Unable to apply constraint '{constraint}' to supplied value {value} for schema of type '{schema_type}'" # noqa: B023 + ) + raise ve + return x + + chain_schema_steps.append( + cs.no_info_wrap_validator_function( + _apply_constraint_with_incompatibility_info, cs.str_schema(**{constraint: value}) + ) + ) + elif constraint in NUMERIC_VALIDATOR_LOOKUP: + if constraint in LENGTH_CONSTRAINTS: + inner_schema = schema + while inner_schema['type'] in {'function-before', 'function-wrap', 'function-after'}: + inner_schema = inner_schema['schema'] # type: ignore + inner_schema_type = inner_schema['type'] + if inner_schema_type == 'list' or ( + inner_schema_type == 'json-or-python' and inner_schema['json_schema']['type'] == 'list' # type: ignore + ): + js_constraint_key = 'minItems' if constraint == 'min_length' else 'maxItems' + else: + js_constraint_key = 'minLength' if constraint == 'min_length' else 'maxLength' + else: + js_constraint_key = constraint + + schema = cs.no_info_after_validator_function( + partial(NUMERIC_VALIDATOR_LOOKUP[constraint], **{constraint: value}), schema + ) + metadata = schema.get('metadata', {}) + if (existing_json_schema_updates := metadata.get('pydantic_js_updates')) is not None: + metadata['pydantic_js_updates'] = { + **existing_json_schema_updates, + **{js_constraint_key: as_jsonable_value(value)}, + } + else: + metadata['pydantic_js_updates'] = {js_constraint_key: as_jsonable_value(value)} + schema['metadata'] = metadata + elif constraint == 'allow_inf_nan' and value is False: + schema = cs.no_info_after_validator_function( + forbid_inf_nan_check, + schema, + ) + else: + # It's rare that we'd get here, but it's possible if we add a new constraint and forget to handle it + # Most constraint errors are caught at runtime during attempted application + raise RuntimeError(f"Unable to apply constraint '{constraint}' to schema of type '{schema_type}'") + + for annotation in other_metadata: + if (annotation_type := type(annotation)) in (at_to_constraint_map := _get_at_to_constraint_map()): + constraint = at_to_constraint_map[annotation_type] + validator = NUMERIC_VALIDATOR_LOOKUP.get(constraint) + if validator is None: + raise ValueError(f'Unknown constraint {constraint}') + schema = cs.no_info_after_validator_function( + partial(validator, {constraint: getattr(annotation, constraint)}), schema + ) + continue + elif isinstance(annotation, (at.Predicate, at.Not)): + predicate_name = f'{annotation.func.__qualname__}' if hasattr(annotation.func, '__qualname__') else '' + + def val_func(v: Any) -> Any: + predicate_satisfied = annotation.func(v) # noqa: B023 + + # annotation.func may also raise an exception, let it pass through + if isinstance(annotation, at.Predicate): # noqa: B023 + if not predicate_satisfied: + raise PydanticCustomError( + 'predicate_failed', + f'Predicate {predicate_name} failed', # type: ignore # noqa: B023 + ) + else: + if predicate_satisfied: + raise PydanticCustomError( + 'not_operation_failed', + f'Not of {predicate_name} failed', # type: ignore # noqa: B023 + ) + + return v + + schema = cs.no_info_after_validator_function(val_func, schema) + else: + # ignore any other unknown metadata + return None + + if chain_schema_steps: + chain_schema_steps = [schema] + chain_schema_steps + return cs.chain_schema(chain_schema_steps) + + return schema + + +def collect_known_metadata(annotations: Iterable[Any]) -> tuple[dict[str, Any], list[Any]]: + """Split `annotations` into known metadata and unknown annotations. + + Args: + annotations: An iterable of annotations. + + Returns: + A tuple contains a dict of known metadata and a list of unknown annotations. + + Example: + ```python + from annotated_types import Gt, Len + + from pydantic._internal._known_annotated_metadata import collect_known_metadata + + print(collect_known_metadata([Gt(1), Len(42), ...])) + #> ({'gt': 1, 'min_length': 42}, [Ellipsis]) + ``` + """ + annotations = expand_grouped_metadata(annotations) + + res: dict[str, Any] = {} + remaining: list[Any] = [] + + for annotation in annotations: + # isinstance(annotation, PydanticMetadata) also covers ._fields:_PydanticGeneralMetadata + if isinstance(annotation, PydanticMetadata): + res.update(annotation.__dict__) + # we don't use dataclasses.asdict because that recursively calls asdict on the field values + elif (annotation_type := type(annotation)) in (at_to_constraint_map := _get_at_to_constraint_map()): + constraint = at_to_constraint_map[annotation_type] + res[constraint] = getattr(annotation, constraint) + elif isinstance(annotation, type) and issubclass(annotation, PydanticMetadata): + # also support PydanticMetadata classes being used without initialisation, + # e.g. `Annotated[int, Strict]` as well as `Annotated[int, Strict()]` + res.update({k: v for k, v in vars(annotation).items() if not k.startswith('_')}) + else: + remaining.append(annotation) + # Nones can sneak in but pydantic-core will reject them + # it'd be nice to clean things up so we don't put in None (we probably don't _need_ to, it was just easier) + # but this is simple enough to kick that can down the road + res = {k: v for k, v in res.items() if v is not None} + return res, remaining + + +def check_metadata(metadata: dict[str, Any], allowed: Iterable[str], source_type: Any) -> None: + """A small utility function to validate that the given metadata can be applied to the target. + More than saving lines of code, this gives us a consistent error message for all of our internal implementations. + + Args: + metadata: A dict of metadata. + allowed: An iterable of allowed metadata. + source_type: The source type. + + Raises: + TypeError: If there is metadatas that can't be applied on source type. + """ + unknown = metadata.keys() - set(allowed) + if unknown: + raise TypeError( + f'The following constraints cannot be applied to {source_type!r}: {", ".join([f"{k!r}" for k in unknown])}' + ) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_mock_val_ser.py b/lib/python3.12/site-packages/pydantic/_internal/_mock_val_ser.py new file mode 100644 index 0000000000000000000000000000000000000000..9125ab32e5348b6e873c4b1939a887ebde35d97a --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_mock_val_ser.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, TypeVar, Union + +from pydantic_core import CoreSchema, SchemaSerializer, SchemaValidator + +from ..errors import PydanticErrorCodes, PydanticUserError +from ..plugin._schema_validator import PluggableSchemaValidator + +if TYPE_CHECKING: + from ..dataclasses import PydanticDataclass + from ..main import BaseModel + from ..type_adapter import TypeAdapter + + +ValSer = TypeVar('ValSer', bound=Union[SchemaValidator, PluggableSchemaValidator, SchemaSerializer]) +T = TypeVar('T') + + +class MockCoreSchema(Mapping[str, Any]): + """Mocker for `pydantic_core.CoreSchema` which optionally attempts to + rebuild the thing it's mocking when one of its methods is accessed and raises an error if that fails. + """ + + __slots__ = '_error_message', '_code', '_attempt_rebuild', '_built_memo' + + def __init__( + self, + error_message: str, + *, + code: PydanticErrorCodes, + attempt_rebuild: Callable[[], CoreSchema | None] | None = None, + ) -> None: + self._error_message = error_message + self._code: PydanticErrorCodes = code + self._attempt_rebuild = attempt_rebuild + self._built_memo: CoreSchema | None = None + + def __getitem__(self, key: str) -> Any: + return self._get_built().__getitem__(key) + + def __len__(self) -> int: + return self._get_built().__len__() + + def __iter__(self) -> Iterator[str]: + return self._get_built().__iter__() + + def _get_built(self) -> CoreSchema: + if self._built_memo is not None: + return self._built_memo + + if self._attempt_rebuild: + schema = self._attempt_rebuild() + if schema is not None: + self._built_memo = schema + return schema + raise PydanticUserError(self._error_message, code=self._code) + + def rebuild(self) -> CoreSchema | None: + self._built_memo = None + if self._attempt_rebuild: + schema = self._attempt_rebuild() + if schema is not None: + return schema + else: + raise PydanticUserError(self._error_message, code=self._code) + return None + + +class MockValSer(Generic[ValSer]): + """Mocker for `pydantic_core.SchemaValidator` or `pydantic_core.SchemaSerializer` which optionally attempts to + rebuild the thing it's mocking when one of its methods is accessed and raises an error if that fails. + """ + + __slots__ = '_error_message', '_code', '_val_or_ser', '_attempt_rebuild' + + def __init__( + self, + error_message: str, + *, + code: PydanticErrorCodes, + val_or_ser: Literal['validator', 'serializer'], + attempt_rebuild: Callable[[], ValSer | None] | None = None, + ) -> None: + self._error_message = error_message + self._val_or_ser = SchemaValidator if val_or_ser == 'validator' else SchemaSerializer + self._code: PydanticErrorCodes = code + self._attempt_rebuild = attempt_rebuild + + def __getattr__(self, item: str) -> None: + __tracebackhide__ = True + if self._attempt_rebuild: + val_ser = self._attempt_rebuild() + if val_ser is not None: + return getattr(val_ser, item) + + # raise an AttributeError if `item` doesn't exist + getattr(self._val_or_ser, item) + raise PydanticUserError(self._error_message, code=self._code) + + def rebuild(self) -> ValSer | None: + if self._attempt_rebuild: + val_ser = self._attempt_rebuild() + if val_ser is not None: + return val_ser + else: + raise PydanticUserError(self._error_message, code=self._code) + return None + + +def set_type_adapter_mocks(adapter: TypeAdapter) -> None: + """Set `core_schema`, `validator` and `serializer` to mock core types on a type adapter instance. + + Args: + adapter: The type adapter instance to set the mocks on + """ + type_repr = str(adapter._type) + undefined_type_error_message = ( + f'`TypeAdapter[{type_repr}]` is not fully defined; you should define `{type_repr}` and all referenced types,' + f' then call `.rebuild()` on the instance.' + ) + + def attempt_rebuild_fn(attr_fn: Callable[[TypeAdapter], T]) -> Callable[[], T | None]: + def handler() -> T | None: + if adapter.rebuild(raise_errors=False, _parent_namespace_depth=5) is not False: + return attr_fn(adapter) + return None + + return handler + + adapter.core_schema = MockCoreSchema( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + attempt_rebuild=attempt_rebuild_fn(lambda ta: ta.core_schema), + ) + adapter.validator = MockValSer( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + val_or_ser='validator', + attempt_rebuild=attempt_rebuild_fn(lambda ta: ta.validator), + ) + adapter.serializer = MockValSer( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + val_or_ser='serializer', + attempt_rebuild=attempt_rebuild_fn(lambda ta: ta.serializer), + ) + + +def set_model_mocks(cls: type[BaseModel], undefined_name: str = 'all referenced types') -> None: + """Set `__pydantic_core_schema__`, `__pydantic_validator__` and `__pydantic_serializer__` to mock core types on a model. + + Args: + cls: The model class to set the mocks on + undefined_name: Name of the undefined thing, used in error messages + """ + undefined_type_error_message = ( + f'`{cls.__name__}` is not fully defined; you should define {undefined_name},' + f' then call `{cls.__name__}.model_rebuild()`.' + ) + + def attempt_rebuild_fn(attr_fn: Callable[[type[BaseModel]], T]) -> Callable[[], T | None]: + def handler() -> T | None: + if cls.model_rebuild(raise_errors=False, _parent_namespace_depth=5) is not False: + return attr_fn(cls) + return None + + return handler + + cls.__pydantic_core_schema__ = MockCoreSchema( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_core_schema__), + ) + cls.__pydantic_validator__ = MockValSer( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + val_or_ser='validator', + attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_validator__), + ) + cls.__pydantic_serializer__ = MockValSer( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + val_or_ser='serializer', + attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_serializer__), + ) + + +def set_dataclass_mocks(cls: type[PydanticDataclass], undefined_name: str = 'all referenced types') -> None: + """Set `__pydantic_validator__` and `__pydantic_serializer__` to `MockValSer`s on a dataclass. + + Args: + cls: The model class to set the mocks on + undefined_name: Name of the undefined thing, used in error messages + """ + from ..dataclasses import rebuild_dataclass + + undefined_type_error_message = ( + f'`{cls.__name__}` is not fully defined; you should define {undefined_name},' + f' then call `pydantic.dataclasses.rebuild_dataclass({cls.__name__})`.' + ) + + def attempt_rebuild_fn(attr_fn: Callable[[type[PydanticDataclass]], T]) -> Callable[[], T | None]: + def handler() -> T | None: + if rebuild_dataclass(cls, raise_errors=False, _parent_namespace_depth=5) is not False: + return attr_fn(cls) + return None + + return handler + + cls.__pydantic_core_schema__ = MockCoreSchema( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_core_schema__), + ) + cls.__pydantic_validator__ = MockValSer( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + val_or_ser='validator', + attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_validator__), + ) + cls.__pydantic_serializer__ = MockValSer( # pyright: ignore[reportAttributeAccessIssue] + undefined_type_error_message, + code='class-not-fully-defined', + val_or_ser='serializer', + attempt_rebuild=attempt_rebuild_fn(lambda c: c.__pydantic_serializer__), + ) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_model_construction.py b/lib/python3.12/site-packages/pydantic/_internal/_model_construction.py new file mode 100644 index 0000000000000000000000000000000000000000..fd5d68b37f73060060db2f59bd7e9ed66f09d576 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_model_construction.py @@ -0,0 +1,792 @@ +"""Private logic for creating models.""" + +from __future__ import annotations as _annotations + +import builtins +import operator +import sys +import typing +import warnings +import weakref +from abc import ABCMeta +from functools import cache, partial, wraps +from types import FunctionType +from typing import Any, Callable, Generic, Literal, NoReturn, cast + +from pydantic_core import PydanticUndefined, SchemaSerializer +from typing_extensions import TypeAliasType, dataclass_transform, deprecated, get_args, get_origin +from typing_inspection import typing_objects + +from ..errors import PydanticUndefinedAnnotation, PydanticUserError +from ..plugin._schema_validator import create_schema_validator +from ..warnings import GenericBeforeBaseModelWarning, PydanticDeprecatedSince20 +from ._config import ConfigWrapper +from ._decorators import DecoratorInfos, PydanticDescriptorProxy, get_attribute_from_bases, unwrap_wrapped_function +from ._fields import collect_model_fields, is_valid_field_name, is_valid_privateattr_name +from ._generate_schema import GenerateSchema, InvalidSchemaError +from ._generics import PydanticGenericMetadata, get_model_typevars_map +from ._import_utils import import_cached_base_model, import_cached_field_info +from ._mock_val_ser import set_model_mocks +from ._namespace_utils import NsResolver +from ._signature import generate_pydantic_signature +from ._typing_extra import ( + _make_forward_ref, + eval_type_backport, + is_classvar_annotation, + parent_frame_namespace, +) +from ._utils import LazyClassAttribute, SafeGetItemProxy + +if typing.TYPE_CHECKING: + from ..fields import Field as PydanticModelField + from ..fields import FieldInfo, ModelPrivateAttr + from ..fields import PrivateAttr as PydanticModelPrivateAttr + from ..main import BaseModel +else: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + PydanticModelField = object() + PydanticModelPrivateAttr = object() + +object_setattr = object.__setattr__ + + +class _ModelNamespaceDict(dict): + """A dictionary subclass that intercepts attribute setting on model classes and + warns about overriding of decorators. + """ + + def __setitem__(self, k: str, v: object) -> None: + existing: Any = self.get(k, None) + if existing and v is not existing and isinstance(existing, PydanticDescriptorProxy): + warnings.warn(f'`{k}` overrides an existing Pydantic `{existing.decorator_info.decorator_repr}` decorator') + + return super().__setitem__(k, v) + + +def NoInitField( + *, + init: Literal[False] = False, +) -> Any: + """Only for typing purposes. Used as default value of `__pydantic_fields_set__`, + `__pydantic_extra__`, `__pydantic_private__`, so they could be ignored when + synthesizing the `__init__` signature. + """ + + +@dataclass_transform(kw_only_default=True, field_specifiers=(PydanticModelField, PydanticModelPrivateAttr, NoInitField)) +class ModelMetaclass(ABCMeta): + def __new__( + mcs, + cls_name: str, + bases: tuple[type[Any], ...], + namespace: dict[str, Any], + __pydantic_generic_metadata__: PydanticGenericMetadata | None = None, + __pydantic_reset_parent_namespace__: bool = True, + _create_model_module: str | None = None, + **kwargs: Any, + ) -> type: + """Metaclass for creating Pydantic models. + + Args: + cls_name: The name of the class to be created. + bases: The base classes of the class to be created. + namespace: The attribute dictionary of the class to be created. + __pydantic_generic_metadata__: Metadata for generic models. + __pydantic_reset_parent_namespace__: Reset parent namespace. + _create_model_module: The module of the class to be created, if created by `create_model`. + **kwargs: Catch-all for any other keyword arguments. + + Returns: + The new class created by the metaclass. + """ + # Note `ModelMetaclass` refers to `BaseModel`, but is also used to *create* `BaseModel`, so we rely on the fact + # that `BaseModel` itself won't have any bases, but any subclass of it will, to determine whether the `__new__` + # call we're in the middle of is for the `BaseModel` class. + if bases: + base_field_names, class_vars, base_private_attributes = mcs._collect_bases_data(bases) + + config_wrapper = ConfigWrapper.for_model(bases, namespace, kwargs) + namespace['model_config'] = config_wrapper.config_dict + private_attributes = inspect_namespace( + namespace, config_wrapper.ignored_types, class_vars, base_field_names + ) + if private_attributes or base_private_attributes: + original_model_post_init = get_model_post_init(namespace, bases) + if original_model_post_init is not None: + # if there are private_attributes and a model_post_init function, we handle both + + @wraps(original_model_post_init) + def wrapped_model_post_init(self: BaseModel, context: Any, /) -> None: + """We need to both initialize private attributes and call the user-defined model_post_init + method. + """ + init_private_attributes(self, context) + original_model_post_init(self, context) + + namespace['model_post_init'] = wrapped_model_post_init + else: + namespace['model_post_init'] = init_private_attributes + + namespace['__class_vars__'] = class_vars + namespace['__private_attributes__'] = {**base_private_attributes, **private_attributes} + + cls = cast('type[BaseModel]', super().__new__(mcs, cls_name, bases, namespace, **kwargs)) + BaseModel_ = import_cached_base_model() + + mro = cls.__mro__ + if Generic in mro and mro.index(Generic) < mro.index(BaseModel_): + warnings.warn( + GenericBeforeBaseModelWarning( + 'Classes should inherit from `BaseModel` before generic classes (e.g. `typing.Generic[T]`) ' + 'for pydantic generics to work properly.' + ), + stacklevel=2, + ) + + cls.__pydantic_custom_init__ = not getattr(cls.__init__, '__pydantic_base_init__', False) + cls.__pydantic_post_init__ = ( + None if cls.model_post_init is BaseModel_.model_post_init else 'model_post_init' + ) + + cls.__pydantic_setattr_handlers__ = {} + + cls.__pydantic_decorators__ = DecoratorInfos.build(cls) + + # Use the getattr below to grab the __parameters__ from the `typing.Generic` parent class + if __pydantic_generic_metadata__: + cls.__pydantic_generic_metadata__ = __pydantic_generic_metadata__ + else: + parent_parameters = getattr(cls, '__pydantic_generic_metadata__', {}).get('parameters', ()) + parameters = getattr(cls, '__parameters__', None) or parent_parameters + if parameters and parent_parameters and not all(x in parameters for x in parent_parameters): + from ..root_model import RootModelRootType + + missing_parameters = tuple(x for x in parameters if x not in parent_parameters) + if RootModelRootType in parent_parameters and RootModelRootType not in parameters: + # This is a special case where the user has subclassed `RootModel`, but has not parametrized + # RootModel with the generic type identifiers being used. Ex: + # class MyModel(RootModel, Generic[T]): + # root: T + # Should instead just be: + # class MyModel(RootModel[T]): + # root: T + parameters_str = ', '.join([x.__name__ for x in missing_parameters]) + error_message = ( + f'{cls.__name__} is a subclass of `RootModel`, but does not include the generic type identifier(s) ' + f'{parameters_str} in its parameters. ' + f'You should parametrize RootModel directly, e.g., `class {cls.__name__}(RootModel[{parameters_str}]): ...`.' + ) + else: + combined_parameters = parent_parameters + missing_parameters + parameters_str = ', '.join([str(x) for x in combined_parameters]) + generic_type_label = f'typing.Generic[{parameters_str}]' + error_message = ( + f'All parameters must be present on typing.Generic;' + f' you should inherit from {generic_type_label}.' + ) + if Generic not in bases: # pragma: no cover + # We raise an error here not because it is desirable, but because some cases are mishandled. + # It would be nice to remove this error and still have things behave as expected, it's just + # challenging because we are using a custom `__class_getitem__` to parametrize generic models, + # and not returning a typing._GenericAlias from it. + bases_str = ', '.join([x.__name__ for x in bases] + [generic_type_label]) + error_message += ( + f' Note: `typing.Generic` must go last: `class {cls.__name__}({bases_str}): ...`)' + ) + raise TypeError(error_message) + + cls.__pydantic_generic_metadata__ = { + 'origin': None, + 'args': (), + 'parameters': parameters, + } + + cls.__pydantic_complete__ = False # Ensure this specific class gets completed + + # preserve `__set_name__` protocol defined in https://peps.python.org/pep-0487 + # for attributes not in `new_namespace` (e.g. private attributes) + for name, obj in private_attributes.items(): + obj.__set_name__(cls, name) + + if __pydantic_reset_parent_namespace__: + cls.__pydantic_parent_namespace__ = build_lenient_weakvaluedict(parent_frame_namespace()) + parent_namespace: dict[str, Any] | None = getattr(cls, '__pydantic_parent_namespace__', None) + if isinstance(parent_namespace, dict): + parent_namespace = unpack_lenient_weakvaluedict(parent_namespace) + + ns_resolver = NsResolver(parent_namespace=parent_namespace) + + set_model_fields(cls, config_wrapper=config_wrapper, ns_resolver=ns_resolver) + + # This is also set in `complete_model_class()`, after schema gen because they are recreated. + # We set them here as well for backwards compatibility: + cls.__pydantic_computed_fields__ = { + k: v.info for k, v in cls.__pydantic_decorators__.computed_fields.items() + } + + if config_wrapper.defer_build: + # TODO we can also stop there if `__pydantic_fields_complete__` is False. + # However, `set_model_fields()` is currently lenient and we don't have access to the `NameError`. + # (which is useful as we can provide the name in the error message: `set_model_mock(cls, e.name)`) + set_model_mocks(cls) + else: + # Any operation that requires accessing the field infos instances should be put inside + # `complete_model_class()`: + complete_model_class( + cls, + config_wrapper, + raise_errors=False, + ns_resolver=ns_resolver, + create_model_module=_create_model_module, + ) + + if config_wrapper.frozen and '__hash__' not in namespace: + set_default_hash_func(cls, bases) + + # using super(cls, cls) on the next line ensures we only call the parent class's __pydantic_init_subclass__ + # I believe the `type: ignore` is only necessary because mypy doesn't realize that this code branch is + # only hit for _proper_ subclasses of BaseModel + super(cls, cls).__pydantic_init_subclass__(**kwargs) # type: ignore[misc] + return cls + else: + # These are instance variables, but have been assigned to `NoInitField` to trick the type checker. + for instance_slot in '__pydantic_fields_set__', '__pydantic_extra__', '__pydantic_private__': + namespace.pop( + instance_slot, + None, # In case the metaclass is used with a class other than `BaseModel`. + ) + namespace.get('__annotations__', {}).clear() + return super().__new__(mcs, cls_name, bases, namespace, **kwargs) + + if not typing.TYPE_CHECKING: # pragma: no branch + # We put `__getattr__` in a non-TYPE_CHECKING block because otherwise, mypy allows arbitrary attribute access + + def __getattr__(self, item: str) -> Any: + """This is necessary to keep attribute access working for class attribute access.""" + private_attributes = self.__dict__.get('__private_attributes__') + if private_attributes and item in private_attributes: + return private_attributes[item] + raise AttributeError(item) + + @classmethod + def __prepare__(cls, *args: Any, **kwargs: Any) -> dict[str, object]: + return _ModelNamespaceDict() + + def __instancecheck__(self, instance: Any) -> bool: + """Avoid calling ABC _abc_instancecheck unless we're pretty sure. + + See #3829 and python/cpython#92810 + """ + return hasattr(instance, '__pydantic_decorators__') and super().__instancecheck__(instance) + + def __subclasscheck__(self, subclass: type[Any]) -> bool: + """Avoid calling ABC _abc_subclasscheck unless we're pretty sure. + + See #3829 and python/cpython#92810 + """ + return hasattr(subclass, '__pydantic_decorators__') and super().__subclasscheck__(subclass) + + @staticmethod + def _collect_bases_data(bases: tuple[type[Any], ...]) -> tuple[set[str], set[str], dict[str, ModelPrivateAttr]]: + BaseModel = import_cached_base_model() + + field_names: set[str] = set() + class_vars: set[str] = set() + private_attributes: dict[str, ModelPrivateAttr] = {} + for base in bases: + if issubclass(base, BaseModel) and base is not BaseModel: + # model_fields might not be defined yet in the case of generics, so we use getattr here: + field_names.update(getattr(base, '__pydantic_fields__', {}).keys()) + class_vars.update(base.__class_vars__) + private_attributes.update(base.__private_attributes__) + return field_names, class_vars, private_attributes + + @property + @deprecated('The `__fields__` attribute is deprecated, use `model_fields` instead.', category=None) + def __fields__(self) -> dict[str, FieldInfo]: + warnings.warn( + 'The `__fields__` attribute is deprecated, use `model_fields` instead.', + PydanticDeprecatedSince20, + stacklevel=2, + ) + return getattr(self, '__pydantic_fields__', {}) + + @property + def __pydantic_fields_complete__(self) -> bool: + """Whether the fields where successfully collected (i.e. type hints were successfully resolves). + + This is a private attribute, not meant to be used outside Pydantic. + """ + if not hasattr(self, '__pydantic_fields__'): + return False + + field_infos = cast('dict[str, FieldInfo]', self.__pydantic_fields__) # pyright: ignore[reportAttributeAccessIssue] + + return all(field_info._complete for field_info in field_infos.values()) + + def __dir__(self) -> list[str]: + attributes = list(super().__dir__()) + if '__fields__' in attributes: + attributes.remove('__fields__') + return attributes + + +def init_private_attributes(self: BaseModel, context: Any, /) -> None: + """This function is meant to behave like a BaseModel method to initialise private attributes. + + It takes context as an argument since that's what pydantic-core passes when calling it. + + Args: + self: The BaseModel instance. + context: The context. + """ + if getattr(self, '__pydantic_private__', None) is None: + pydantic_private = {} + for name, private_attr in self.__private_attributes__.items(): + default = private_attr.get_default() + if default is not PydanticUndefined: + pydantic_private[name] = default + object_setattr(self, '__pydantic_private__', pydantic_private) + + +def get_model_post_init(namespace: dict[str, Any], bases: tuple[type[Any], ...]) -> Callable[..., Any] | None: + """Get the `model_post_init` method from the namespace or the class bases, or `None` if not defined.""" + if 'model_post_init' in namespace: + return namespace['model_post_init'] + + BaseModel = import_cached_base_model() + + model_post_init = get_attribute_from_bases(bases, 'model_post_init') + if model_post_init is not BaseModel.model_post_init: + return model_post_init + + +def inspect_namespace( # noqa C901 + namespace: dict[str, Any], + ignored_types: tuple[type[Any], ...], + base_class_vars: set[str], + base_class_fields: set[str], +) -> dict[str, ModelPrivateAttr]: + """Iterate over the namespace and: + * gather private attributes + * check for items which look like fields but are not (e.g. have no annotation) and warn. + + Args: + namespace: The attribute dictionary of the class to be created. + ignored_types: A tuple of ignore types. + base_class_vars: A set of base class class variables. + base_class_fields: A set of base class fields. + + Returns: + A dict contains private attributes info. + + Raises: + TypeError: If there is a `__root__` field in model. + NameError: If private attribute name is invalid. + PydanticUserError: + - If a field does not have a type annotation. + - If a field on base class was overridden by a non-annotated attribute. + """ + from ..fields import ModelPrivateAttr, PrivateAttr + + FieldInfo = import_cached_field_info() + + all_ignored_types = ignored_types + default_ignored_types() + + private_attributes: dict[str, ModelPrivateAttr] = {} + raw_annotations = namespace.get('__annotations__', {}) + + if '__root__' in raw_annotations or '__root__' in namespace: + raise TypeError("To define root models, use `pydantic.RootModel` rather than a field called '__root__'") + + ignored_names: set[str] = set() + for var_name, value in list(namespace.items()): + if var_name == 'model_config' or var_name == '__pydantic_extra__': + continue + elif ( + isinstance(value, type) + and value.__module__ == namespace['__module__'] + and '__qualname__' in namespace + and value.__qualname__.startswith(namespace['__qualname__']) + ): + # `value` is a nested type defined in this namespace; don't error + continue + elif isinstance(value, all_ignored_types) or value.__class__.__module__ == 'functools': + ignored_names.add(var_name) + continue + elif isinstance(value, ModelPrivateAttr): + if var_name.startswith('__'): + raise NameError( + 'Private attributes must not use dunder names;' + f' use a single underscore prefix instead of {var_name!r}.' + ) + elif is_valid_field_name(var_name): + raise NameError( + 'Private attributes must not use valid field names;' + f' use sunder names, e.g. {"_" + var_name!r} instead of {var_name!r}.' + ) + private_attributes[var_name] = value + del namespace[var_name] + elif isinstance(value, FieldInfo) and not is_valid_field_name(var_name): + suggested_name = var_name.lstrip('_') or 'my_field' # don't suggest '' for all-underscore name + raise NameError( + f'Fields must not use names with leading underscores;' + f' e.g., use {suggested_name!r} instead of {var_name!r}.' + ) + + elif var_name.startswith('__'): + continue + elif is_valid_privateattr_name(var_name): + if var_name not in raw_annotations or not is_classvar_annotation(raw_annotations[var_name]): + private_attributes[var_name] = cast(ModelPrivateAttr, PrivateAttr(default=value)) + del namespace[var_name] + elif var_name in base_class_vars: + continue + elif var_name not in raw_annotations: + if var_name in base_class_fields: + raise PydanticUserError( + f'Field {var_name!r} defined on a base class was overridden by a non-annotated attribute. ' + f'All field definitions, including overrides, require a type annotation.', + code='model-field-overridden', + ) + elif isinstance(value, FieldInfo): + raise PydanticUserError( + f'Field {var_name!r} requires a type annotation', code='model-field-missing-annotation' + ) + else: + raise PydanticUserError( + f'A non-annotated attribute was detected: `{var_name} = {value!r}`. All model fields require a ' + f'type annotation; if `{var_name}` is not meant to be a field, you may be able to resolve this ' + f"error by annotating it as a `ClassVar` or updating `model_config['ignored_types']`.", + code='model-field-missing-annotation', + ) + + for ann_name, ann_type in raw_annotations.items(): + if ( + is_valid_privateattr_name(ann_name) + and ann_name not in private_attributes + and ann_name not in ignored_names + # This condition can be a false negative when `ann_type` is stringified, + # but it is handled in most cases in `set_model_fields`: + and not is_classvar_annotation(ann_type) + and ann_type not in all_ignored_types + and getattr(ann_type, '__module__', None) != 'functools' + ): + if isinstance(ann_type, str): + # Walking up the frames to get the module namespace where the model is defined + # (as the model class wasn't created yet, we unfortunately can't use `cls.__module__`): + frame = sys._getframe(2) + if frame is not None: + try: + ann_type = eval_type_backport( + _make_forward_ref(ann_type, is_argument=False, is_class=True), + globalns=frame.f_globals, + localns=frame.f_locals, + ) + except (NameError, TypeError): + pass + + if typing_objects.is_annotated(get_origin(ann_type)): + _, *metadata = get_args(ann_type) + private_attr = next((v for v in metadata if isinstance(v, ModelPrivateAttr)), None) + if private_attr is not None: + private_attributes[ann_name] = private_attr + continue + private_attributes[ann_name] = PrivateAttr() + + return private_attributes + + +def set_default_hash_func(cls: type[BaseModel], bases: tuple[type[Any], ...]) -> None: + base_hash_func = get_attribute_from_bases(bases, '__hash__') + new_hash_func = make_hash_func(cls) + if base_hash_func in {None, object.__hash__} or getattr(base_hash_func, '__code__', None) == new_hash_func.__code__: + # If `__hash__` is some default, we generate a hash function. + # It will be `None` if not overridden from BaseModel. + # It may be `object.__hash__` if there is another + # parent class earlier in the bases which doesn't override `__hash__` (e.g. `typing.Generic`). + # It may be a value set by `set_default_hash_func` if `cls` is a subclass of another frozen model. + # In the last case we still need a new hash function to account for new `model_fields`. + cls.__hash__ = new_hash_func + + +def make_hash_func(cls: type[BaseModel]) -> Any: + getter = operator.itemgetter(*cls.__pydantic_fields__.keys()) if cls.__pydantic_fields__ else lambda _: 0 + + def hash_func(self: Any) -> int: + try: + return hash(getter(self.__dict__)) + except KeyError: + # In rare cases (such as when using the deprecated copy method), the __dict__ may not contain + # all model fields, which is how we can get here. + # getter(self.__dict__) is much faster than any 'safe' method that accounts for missing keys, + # and wrapping it in a `try` doesn't slow things down much in the common case. + return hash(getter(SafeGetItemProxy(self.__dict__))) + + return hash_func + + +def set_model_fields( + cls: type[BaseModel], + config_wrapper: ConfigWrapper, + ns_resolver: NsResolver | None, +) -> None: + """Collect and set `cls.__pydantic_fields__` and `cls.__class_vars__`. + + Args: + cls: BaseModel or dataclass. + config_wrapper: The config wrapper instance. + ns_resolver: Namespace resolver to use when getting model annotations. + """ + typevars_map = get_model_typevars_map(cls) + fields, class_vars = collect_model_fields(cls, config_wrapper, ns_resolver, typevars_map=typevars_map) + + cls.__pydantic_fields__ = fields + cls.__class_vars__.update(class_vars) + + for k in class_vars: + # Class vars should not be private attributes + # We remove them _here_ and not earlier because we rely on inspecting the class to determine its classvars, + # but private attributes are determined by inspecting the namespace _prior_ to class creation. + # In the case that a classvar with a leading-'_' is defined via a ForwardRef (e.g., when using + # `__future__.annotations`), we want to remove the private attribute which was detected _before_ we knew it + # evaluated to a classvar + + value = cls.__private_attributes__.pop(k, None) + if value is not None and value.default is not PydanticUndefined: + setattr(cls, k, value.default) + + +def complete_model_class( + cls: type[BaseModel], + config_wrapper: ConfigWrapper, + *, + raise_errors: bool = True, + ns_resolver: NsResolver | None = None, + create_model_module: str | None = None, +) -> bool: + """Finish building a model class. + + This logic must be called after class has been created since validation functions must be bound + and `get_type_hints` requires a class object. + + Args: + cls: BaseModel or dataclass. + config_wrapper: The config wrapper instance. + raise_errors: Whether to raise errors. + ns_resolver: The namespace resolver instance to use during schema building. + create_model_module: The module of the class to be created, if created by `create_model`. + + Returns: + `True` if the model is successfully completed, else `False`. + + Raises: + PydanticUndefinedAnnotation: If `PydanticUndefinedAnnotation` occurs in`__get_pydantic_core_schema__` + and `raise_errors=True`. + """ + typevars_map = get_model_typevars_map(cls) + gen_schema = GenerateSchema( + config_wrapper, + ns_resolver, + typevars_map, + ) + + try: + schema = gen_schema.generate_schema(cls) + except PydanticUndefinedAnnotation as e: + if raise_errors: + raise + set_model_mocks(cls, f'`{e.name}`') + return False + + core_config = config_wrapper.core_config(title=cls.__name__) + + try: + schema = gen_schema.clean_schema(schema) + except InvalidSchemaError: + set_model_mocks(cls) + return False + + # This needs to happen *after* model schema generation, as the return type + # of the properties are evaluated and the `ComputedFieldInfo` are recreated: + cls.__pydantic_computed_fields__ = {k: v.info for k, v in cls.__pydantic_decorators__.computed_fields.items()} + + set_deprecated_descriptors(cls) + + cls.__pydantic_core_schema__ = schema + + cls.__pydantic_validator__ = create_schema_validator( + schema, + cls, + create_model_module or cls.__module__, + cls.__qualname__, + 'create_model' if create_model_module else 'BaseModel', + core_config, + config_wrapper.plugin_settings, + ) + cls.__pydantic_serializer__ = SchemaSerializer(schema, core_config) + cls.__pydantic_complete__ = True + + # set __signature__ attr only for model class, but not for its instances + # (because instances can define `__call__`, and `inspect.signature` shouldn't + # use the `__signature__` attribute and instead generate from `__call__`). + cls.__signature__ = LazyClassAttribute( + '__signature__', + partial( + generate_pydantic_signature, + init=cls.__init__, + fields=cls.__pydantic_fields__, + validate_by_name=config_wrapper.validate_by_name, + extra=config_wrapper.extra, + ), + ) + return True + + +def set_deprecated_descriptors(cls: type[BaseModel]) -> None: + """Set data descriptors on the class for deprecated fields.""" + for field, field_info in cls.__pydantic_fields__.items(): + if (msg := field_info.deprecation_message) is not None: + desc = _DeprecatedFieldDescriptor(msg) + desc.__set_name__(cls, field) + setattr(cls, field, desc) + + for field, computed_field_info in cls.__pydantic_computed_fields__.items(): + if ( + (msg := computed_field_info.deprecation_message) is not None + # Avoid having two warnings emitted: + and not hasattr(unwrap_wrapped_function(computed_field_info.wrapped_property), '__deprecated__') + ): + desc = _DeprecatedFieldDescriptor(msg, computed_field_info.wrapped_property) + desc.__set_name__(cls, field) + setattr(cls, field, desc) + + +class _DeprecatedFieldDescriptor: + """Read-only data descriptor used to emit a runtime deprecation warning before accessing a deprecated field. + + Attributes: + msg: The deprecation message to be emitted. + wrapped_property: The property instance if the deprecated field is a computed field, or `None`. + field_name: The name of the field being deprecated. + """ + + field_name: str + + def __init__(self, msg: str, wrapped_property: property | None = None) -> None: + self.msg = msg + self.wrapped_property = wrapped_property + + def __set_name__(self, cls: type[BaseModel], name: str) -> None: + self.field_name = name + + def __get__(self, obj: BaseModel | None, obj_type: type[BaseModel] | None = None) -> Any: + if obj is None: + if self.wrapped_property is not None: + return self.wrapped_property.__get__(None, obj_type) + raise AttributeError(self.field_name) + + warnings.warn(self.msg, builtins.DeprecationWarning, stacklevel=2) + + if self.wrapped_property is not None: + return self.wrapped_property.__get__(obj, obj_type) + return obj.__dict__[self.field_name] + + # Defined to make it a data descriptor and take precedence over the instance's dictionary. + # Note that it will not be called when setting a value on a model instance + # as `BaseModel.__setattr__` is defined and takes priority. + def __set__(self, obj: Any, value: Any) -> NoReturn: + raise AttributeError(self.field_name) + + +class _PydanticWeakRef: + """Wrapper for `weakref.ref` that enables `pickle` serialization. + + Cloudpickle fails to serialize `weakref.ref` objects due to an arcane error related + to abstract base classes (`abc.ABC`). This class works around the issue by wrapping + `weakref.ref` instead of subclassing it. + + See https://github.com/pydantic/pydantic/issues/6763 for context. + + Semantics: + - If not pickled, behaves the same as a `weakref.ref`. + - If pickled along with the referenced object, the same `weakref.ref` behavior + will be maintained between them after unpickling. + - If pickled without the referenced object, after unpickling the underlying + reference will be cleared (`__call__` will always return `None`). + """ + + def __init__(self, obj: Any): + if obj is None: + # The object will be `None` upon deserialization if the serialized weakref + # had lost its underlying object. + self._wr = None + else: + self._wr = weakref.ref(obj) + + def __call__(self) -> Any: + if self._wr is None: + return None + else: + return self._wr() + + def __reduce__(self) -> tuple[Callable, tuple[weakref.ReferenceType | None]]: + return _PydanticWeakRef, (self(),) + + +def build_lenient_weakvaluedict(d: dict[str, Any] | None) -> dict[str, Any] | None: + """Takes an input dictionary, and produces a new value that (invertibly) replaces the values with weakrefs. + + We can't just use a WeakValueDictionary because many types (including int, str, etc.) can't be stored as values + in a WeakValueDictionary. + + The `unpack_lenient_weakvaluedict` function can be used to reverse this operation. + """ + if d is None: + return None + result = {} + for k, v in d.items(): + try: + proxy = _PydanticWeakRef(v) + except TypeError: + proxy = v + result[k] = proxy + return result + + +def unpack_lenient_weakvaluedict(d: dict[str, Any] | None) -> dict[str, Any] | None: + """Inverts the transform performed by `build_lenient_weakvaluedict`.""" + if d is None: + return None + + result = {} + for k, v in d.items(): + if isinstance(v, _PydanticWeakRef): + v = v() + if v is not None: + result[k] = v + else: + result[k] = v + return result + + +@cache +def default_ignored_types() -> tuple[type[Any], ...]: + from ..fields import ComputedFieldInfo + + ignored_types = [ + FunctionType, + property, + classmethod, + staticmethod, + PydanticDescriptorProxy, + ComputedFieldInfo, + TypeAliasType, # from `typing_extensions` + ] + + if sys.version_info >= (3, 12): + ignored_types.append(typing.TypeAliasType) + + return tuple(ignored_types) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_namespace_utils.py b/lib/python3.12/site-packages/pydantic/_internal/_namespace_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..781dfa2d7c489d3a5ec207a63522e771d3cef22a --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_namespace_utils.py @@ -0,0 +1,293 @@ +from __future__ import annotations + +import sys +from collections.abc import Generator, Iterator, Mapping +from contextlib import contextmanager +from functools import cached_property +from typing import Any, Callable, NamedTuple, TypeVar + +from typing_extensions import ParamSpec, TypeAlias, TypeAliasType, TypeVarTuple + +GlobalsNamespace: TypeAlias = 'dict[str, Any]' +"""A global namespace. + +In most cases, this is a reference to the `__dict__` attribute of a module. +This namespace type is expected as the `globals` argument during annotations evaluation. +""" + +MappingNamespace: TypeAlias = Mapping[str, Any] +"""Any kind of namespace. + +In most cases, this is a local namespace (e.g. the `__dict__` attribute of a class, +the [`f_locals`][frame.f_locals] attribute of a frame object, when dealing with types +defined inside functions). +This namespace type is expected as the `locals` argument during annotations evaluation. +""" + +_TypeVarLike: TypeAlias = 'TypeVar | ParamSpec | TypeVarTuple' + + +class NamespacesTuple(NamedTuple): + """A tuple of globals and locals to be used during annotations evaluation. + + This datastructure is defined as a named tuple so that it can easily be unpacked: + + ```python {lint="skip" test="skip"} + def eval_type(typ: type[Any], ns: NamespacesTuple) -> None: + return eval(typ, *ns) + ``` + """ + + globals: GlobalsNamespace + """The namespace to be used as the `globals` argument during annotations evaluation.""" + + locals: MappingNamespace + """The namespace to be used as the `locals` argument during annotations evaluation.""" + + +def get_module_ns_of(obj: Any) -> dict[str, Any]: + """Get the namespace of the module where the object is defined. + + Caution: this function does not return a copy of the module namespace, so the result + should not be mutated. The burden of enforcing this is on the caller. + """ + module_name = getattr(obj, '__module__', None) + if module_name: + try: + return sys.modules[module_name].__dict__ + except KeyError: + # happens occasionally, see https://github.com/pydantic/pydantic/issues/2363 + return {} + return {} + + +# Note that this class is almost identical to `collections.ChainMap`, but need to enforce +# immutable mappings here: +class LazyLocalNamespace(Mapping[str, Any]): + """A lazily evaluated mapping, to be used as the `locals` argument during annotations evaluation. + + While the [`eval`][eval] function expects a mapping as the `locals` argument, it only + performs `__getitem__` calls. The [`Mapping`][collections.abc.Mapping] abstract base class + is fully implemented only for type checking purposes. + + Args: + *namespaces: The namespaces to consider, in ascending order of priority. + + Example: + ```python {lint="skip" test="skip"} + ns = LazyLocalNamespace({'a': 1, 'b': 2}, {'a': 3}) + ns['a'] + #> 3 + ns['b'] + #> 2 + ``` + """ + + def __init__(self, *namespaces: MappingNamespace) -> None: + self._namespaces = namespaces + + @cached_property + def data(self) -> dict[str, Any]: + return {k: v for ns in self._namespaces for k, v in ns.items()} + + def __len__(self) -> int: + return len(self.data) + + def __getitem__(self, key: str) -> Any: + return self.data[key] + + def __contains__(self, key: object) -> bool: + return key in self.data + + def __iter__(self) -> Iterator[str]: + return iter(self.data) + + +def ns_for_function(obj: Callable[..., Any], parent_namespace: MappingNamespace | None = None) -> NamespacesTuple: + """Return the global and local namespaces to be used when evaluating annotations for the provided function. + + The global namespace will be the `__dict__` attribute of the module the function was defined in. + The local namespace will contain the `__type_params__` introduced by PEP 695. + + Args: + obj: The object to use when building namespaces. + parent_namespace: Optional namespace to be added with the lowest priority in the local namespace. + If the passed function is a method, the `parent_namespace` will be the namespace of the class + the method is defined in. Thus, we also fetch type `__type_params__` from there (i.e. the + class-scoped type variables). + """ + locals_list: list[MappingNamespace] = [] + if parent_namespace is not None: + locals_list.append(parent_namespace) + + # Get the `__type_params__` attribute introduced by PEP 695. + # Note that the `typing._eval_type` function expects type params to be + # passed as a separate argument. However, internally, `_eval_type` calls + # `ForwardRef._evaluate` which will merge type params with the localns, + # essentially mimicking what we do here. + type_params: tuple[_TypeVarLike, ...] = getattr(obj, '__type_params__', ()) + if parent_namespace is not None: + # We also fetch type params from the parent namespace. If present, it probably + # means the function was defined in a class. This is to support the following: + # https://github.com/python/cpython/issues/124089. + type_params += parent_namespace.get('__type_params__', ()) + + locals_list.append({t.__name__: t for t in type_params}) + + # What about short-cirtuiting to `obj.__globals__`? + globalns = get_module_ns_of(obj) + + return NamespacesTuple(globalns, LazyLocalNamespace(*locals_list)) + + +class NsResolver: + """A class responsible for the namespaces resolving logic for annotations evaluation. + + This class handles the namespace logic when evaluating annotations mainly for class objects. + + It holds a stack of classes that are being inspected during the core schema building, + and the `types_namespace` property exposes the globals and locals to be used for + type annotation evaluation. Additionally -- if no class is present in the stack -- a + fallback globals and locals can be provided using the `namespaces_tuple` argument + (this is useful when generating a schema for a simple annotation, e.g. when using + `TypeAdapter`). + + The namespace creation logic is unfortunately flawed in some cases, for backwards + compatibility reasons and to better support valid edge cases. See the description + for the `parent_namespace` argument and the example for more details. + + Args: + namespaces_tuple: The default globals and locals to use if no class is present + on the stack. This can be useful when using the `GenerateSchema` class + with `TypeAdapter`, where the "type" being analyzed is a simple annotation. + parent_namespace: An optional parent namespace that will be added to the locals + with the lowest priority. For a given class defined in a function, the locals + of this function are usually used as the parent namespace: + + ```python {lint="skip" test="skip"} + from pydantic import BaseModel + + def func() -> None: + SomeType = int + + class Model(BaseModel): + f: 'SomeType' + + # when collecting fields, an namespace resolver instance will be created + # this way: + # ns_resolver = NsResolver(parent_namespace={'SomeType': SomeType}) + ``` + + For backwards compatibility reasons and to support valid edge cases, this parent + namespace will be used for *every* type being pushed to the stack. In the future, + we might want to be smarter by only doing so when the type being pushed is defined + in the same module as the parent namespace. + + Example: + ```python {lint="skip" test="skip"} + ns_resolver = NsResolver( + parent_namespace={'fallback': 1}, + ) + + class Sub: + m: 'Model' + + class Model: + some_local = 1 + sub: Sub + + ns_resolver = NsResolver() + + # This is roughly what happens when we build a core schema for `Model`: + with ns_resolver.push(Model): + ns_resolver.types_namespace + #> NamespacesTuple({'Sub': Sub}, {'Model': Model, 'some_local': 1}) + # First thing to notice here, the model being pushed is added to the locals. + # Because `NsResolver` is being used during the model definition, it is not + # yet added to the globals. This is useful when resolving self-referencing annotations. + + with ns_resolver.push(Sub): + ns_resolver.types_namespace + #> NamespacesTuple({'Sub': Sub}, {'Sub': Sub, 'Model': Model}) + # Second thing to notice: `Sub` is present in both the globals and locals. + # This is not an issue, just that as described above, the model being pushed + # is added to the locals, but it happens to be present in the globals as well + # because it is already defined. + # Third thing to notice: `Model` is also added in locals. This is a backwards + # compatibility workaround that allows for `Sub` to be able to resolve `'Model'` + # correctly (as otherwise models would have to be rebuilt even though this + # doesn't look necessary). + ``` + """ + + def __init__( + self, + namespaces_tuple: NamespacesTuple | None = None, + parent_namespace: MappingNamespace | None = None, + ) -> None: + self._base_ns_tuple = namespaces_tuple or NamespacesTuple({}, {}) + self._parent_ns = parent_namespace + self._types_stack: list[type[Any] | TypeAliasType] = [] + + @cached_property + def types_namespace(self) -> NamespacesTuple: + """The current global and local namespaces to be used for annotations evaluation.""" + if not self._types_stack: + # TODO: should we merge the parent namespace here? + # This is relevant for TypeAdapter, where there are no types on the stack, and we might + # need access to the parent_ns. Right now, we sidestep this in `type_adapter.py` by passing + # locals to both parent_ns and the base_ns_tuple, but this is a bit hacky. + # we might consider something like: + # if self._parent_ns is not None: + # # Hacky workarounds, see class docstring: + # # An optional parent namespace that will be added to the locals with the lowest priority + # locals_list: list[MappingNamespace] = [self._parent_ns, self._base_ns_tuple.locals] + # return NamespacesTuple(self._base_ns_tuple.globals, LazyLocalNamespace(*locals_list)) + return self._base_ns_tuple + + typ = self._types_stack[-1] + + globalns = get_module_ns_of(typ) + + locals_list: list[MappingNamespace] = [] + # Hacky workarounds, see class docstring: + # An optional parent namespace that will be added to the locals with the lowest priority + if self._parent_ns is not None: + locals_list.append(self._parent_ns) + if len(self._types_stack) > 1: + first_type = self._types_stack[0] + locals_list.append({first_type.__name__: first_type}) + + # Adding `__type_params__` *before* `vars(typ)`, as the latter takes priority + # (see https://github.com/python/cpython/pull/120272). + # TODO `typ.__type_params__` when we drop support for Python 3.11: + type_params: tuple[_TypeVarLike, ...] = getattr(typ, '__type_params__', ()) + if type_params: + # Adding `__type_params__` is mostly useful for generic classes defined using + # PEP 695 syntax *and* using forward annotations (see the example in + # https://github.com/python/cpython/issues/114053). For TypeAliasType instances, + # it is way less common, but still required if using a string annotation in the alias + # value, e.g. `type A[T] = 'T'` (which is not necessary in most cases). + locals_list.append({t.__name__: t for t in type_params}) + + # TypeAliasType instances don't have a `__dict__` attribute, so the check + # is necessary: + if hasattr(typ, '__dict__'): + locals_list.append(vars(typ)) + + # The `len(self._types_stack) > 1` check above prevents this from being added twice: + locals_list.append({typ.__name__: typ}) + + return NamespacesTuple(globalns, LazyLocalNamespace(*locals_list)) + + @contextmanager + def push(self, typ: type[Any] | TypeAliasType, /) -> Generator[None]: + """Push a type to the stack.""" + self._types_stack.append(typ) + # Reset the cached property: + self.__dict__.pop('types_namespace', None) + try: + yield + finally: + self._types_stack.pop() + self.__dict__.pop('types_namespace', None) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_repr.py b/lib/python3.12/site-packages/pydantic/_internal/_repr.py new file mode 100644 index 0000000000000000000000000000000000000000..bf3cae53e3dc2083daa85288984b42137fe26848 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_repr.py @@ -0,0 +1,125 @@ +"""Tools to provide pretty/human-readable display of objects.""" + +from __future__ import annotations as _annotations + +import types +import typing +from typing import Any + +import typing_extensions +from typing_inspection import typing_objects +from typing_inspection.introspection import is_union_origin + +from . import _typing_extra + +if typing.TYPE_CHECKING: + ReprArgs: typing_extensions.TypeAlias = 'typing.Iterable[tuple[str | None, Any]]' + RichReprResult: typing_extensions.TypeAlias = ( + 'typing.Iterable[Any | tuple[Any] | tuple[str, Any] | tuple[str, Any, Any]]' + ) + + +class PlainRepr(str): + """String class where repr doesn't include quotes. Useful with Representation when you want to return a string + representation of something that is valid (or pseudo-valid) python. + """ + + def __repr__(self) -> str: + return str(self) + + +class Representation: + # Mixin to provide `__str__`, `__repr__`, and `__pretty__` and `__rich_repr__` methods. + # `__pretty__` is used by [devtools](https://python-devtools.helpmanual.io/). + # `__rich_repr__` is used by [rich](https://rich.readthedocs.io/en/stable/pretty.html). + # (this is not a docstring to avoid adding a docstring to classes which inherit from Representation) + + # we don't want to use a type annotation here as it can break get_type_hints + __slots__ = () # type: typing.Collection[str] + + def __repr_args__(self) -> ReprArgs: + """Returns the attributes to show in __str__, __repr__, and __pretty__ this is generally overridden. + + Can either return: + * name - value pairs, e.g.: `[('foo_name', 'foo'), ('bar_name', ['b', 'a', 'r'])]` + * or, just values, e.g.: `[(None, 'foo'), (None, ['b', 'a', 'r'])]` + """ + attrs_names = self.__slots__ + if not attrs_names and hasattr(self, '__dict__'): + attrs_names = self.__dict__.keys() + attrs = ((s, getattr(self, s)) for s in attrs_names) + return [(a, v if v is not self else self.__repr_recursion__(v)) for a, v in attrs if v is not None] + + def __repr_name__(self) -> str: + """Name of the instance's class, used in __repr__.""" + return self.__class__.__name__ + + def __repr_recursion__(self, object: Any) -> str: + """Returns the string representation of a recursive object.""" + # This is copied over from the stdlib `pprint` module: + return f'' + + def __repr_str__(self, join_str: str) -> str: + return join_str.join(repr(v) if a is None else f'{a}={v!r}' for a, v in self.__repr_args__()) + + def __pretty__(self, fmt: typing.Callable[[Any], Any], **kwargs: Any) -> typing.Generator[Any, None, None]: + """Used by devtools (https://python-devtools.helpmanual.io/) to pretty print objects.""" + yield self.__repr_name__() + '(' + yield 1 + for name, value in self.__repr_args__(): + if name is not None: + yield name + '=' + yield fmt(value) + yield ',' + yield 0 + yield -1 + yield ')' + + def __rich_repr__(self) -> RichReprResult: + """Used by Rich (https://rich.readthedocs.io/en/stable/pretty.html) to pretty print objects.""" + for name, field_repr in self.__repr_args__(): + if name is None: + yield field_repr + else: + yield name, field_repr + + def __str__(self) -> str: + return self.__repr_str__(' ') + + def __repr__(self) -> str: + return f'{self.__repr_name__()}({self.__repr_str__(", ")})' + + +def display_as_type(obj: Any) -> str: + """Pretty representation of a type, should be as close as possible to the original type definition string. + + Takes some logic from `typing._type_repr`. + """ + if isinstance(obj, (types.FunctionType, types.BuiltinFunctionType)): + return obj.__name__ + elif obj is ...: + return '...' + elif isinstance(obj, Representation): + return repr(obj) + elif isinstance(obj, typing.ForwardRef) or typing_objects.is_typealiastype(obj): + return str(obj) + + if not isinstance(obj, (_typing_extra.typing_base, _typing_extra.WithArgsTypes, type)): + obj = obj.__class__ + + if is_union_origin(typing_extensions.get_origin(obj)): + args = ', '.join(map(display_as_type, typing_extensions.get_args(obj))) + return f'Union[{args}]' + elif isinstance(obj, _typing_extra.WithArgsTypes): + if typing_objects.is_literal(typing_extensions.get_origin(obj)): + args = ', '.join(map(repr, typing_extensions.get_args(obj))) + else: + args = ', '.join(map(display_as_type, typing_extensions.get_args(obj))) + try: + return f'{obj.__qualname__}[{args}]' + except AttributeError: + return str(obj).replace('typing.', '').replace('typing_extensions.', '') # handles TypeAliasType in 3.12 + elif isinstance(obj, type): + return obj.__qualname__ + else: + return repr(obj).replace('typing.', '').replace('typing_extensions.', '') diff --git a/lib/python3.12/site-packages/pydantic/_internal/_schema_gather.py b/lib/python3.12/site-packages/pydantic/_internal/_schema_gather.py new file mode 100644 index 0000000000000000000000000000000000000000..fc2d806ead6554f22919f0827e681b43889378ad --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_schema_gather.py @@ -0,0 +1,209 @@ +# pyright: reportTypedDictNotRequiredAccess=false, reportGeneralTypeIssues=false, reportArgumentType=false, reportAttributeAccessIssue=false +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TypedDict + +from pydantic_core.core_schema import ComputedField, CoreSchema, DefinitionReferenceSchema, SerSchema +from typing_extensions import TypeAlias + +AllSchemas: TypeAlias = 'CoreSchema | SerSchema | ComputedField' + + +class GatherResult(TypedDict): + """Schema traversing result.""" + + collected_references: dict[str, DefinitionReferenceSchema | None] + """The collected definition references. + + If a definition reference schema can be inlined, it means that there is + only one in the whole core schema. As such, it is stored as the value. + Otherwise, the value is set to `None`. + """ + + deferred_discriminator_schemas: list[CoreSchema] + """The list of core schemas having the discriminator application deferred.""" + + +class MissingDefinitionError(LookupError): + """A reference was pointing to a non-existing core schema.""" + + def __init__(self, schema_reference: str, /) -> None: + self.schema_reference = schema_reference + + +@dataclass +class GatherContext: + """The current context used during core schema traversing. + + Context instances should only be used during schema traversing. + """ + + definitions: dict[str, CoreSchema] + """The available definitions.""" + + deferred_discriminator_schemas: list[CoreSchema] = field(init=False, default_factory=list) + """The list of core schemas having the discriminator application deferred. + + Internally, these core schemas have a specific key set in the core metadata dict. + """ + + collected_references: dict[str, DefinitionReferenceSchema | None] = field(init=False, default_factory=dict) + """The collected definition references. + + If a definition reference schema can be inlined, it means that there is + only one in the whole core schema. As such, it is stored as the value. + Otherwise, the value is set to `None`. + + During schema traversing, definition reference schemas can be added as candidates, or removed + (by setting the value to `None`). + """ + + +def traverse_metadata(schema: AllSchemas, ctx: GatherContext) -> None: + meta = schema.get('metadata') + if meta is not None and 'pydantic_internal_union_discriminator' in meta: + ctx.deferred_discriminator_schemas.append(schema) # pyright: ignore[reportArgumentType] + + +def traverse_definition_ref(def_ref_schema: DefinitionReferenceSchema, ctx: GatherContext) -> None: + schema_ref = def_ref_schema['schema_ref'] + + if schema_ref not in ctx.collected_references: + definition = ctx.definitions.get(schema_ref) + if definition is None: + raise MissingDefinitionError(schema_ref) + + # The `'definition-ref'` schema was only encountered once, make it + # a candidate to be inlined: + ctx.collected_references[schema_ref] = def_ref_schema + traverse_schema(definition, ctx) + if 'serialization' in def_ref_schema: + traverse_schema(def_ref_schema['serialization'], ctx) + traverse_metadata(def_ref_schema, ctx) + else: + # The `'definition-ref'` schema was already encountered, meaning + # the previously encountered schema (and this one) can't be inlined: + ctx.collected_references[schema_ref] = None + + +def traverse_schema(schema: AllSchemas, context: GatherContext) -> None: + # TODO When we drop 3.9, use a match statement to get better type checking and remove + # file-level type ignore. + # (the `'type'` could also be fetched in every `if/elif` statement, but this alters performance). + schema_type = schema['type'] + + if schema_type == 'definition-ref': + traverse_definition_ref(schema, context) + # `traverse_definition_ref` handles the possible serialization and metadata schemas: + return + elif schema_type == 'definitions': + traverse_schema(schema['schema'], context) + for definition in schema['definitions']: + traverse_schema(definition, context) + elif schema_type in {'list', 'set', 'frozenset', 'generator'}: + if 'items_schema' in schema: + traverse_schema(schema['items_schema'], context) + elif schema_type == 'tuple': + if 'items_schema' in schema: + for s in schema['items_schema']: + traverse_schema(s, context) + elif schema_type == 'dict': + if 'keys_schema' in schema: + traverse_schema(schema['keys_schema'], context) + if 'values_schema' in schema: + traverse_schema(schema['values_schema'], context) + elif schema_type == 'union': + for choice in schema['choices']: + if isinstance(choice, tuple): + traverse_schema(choice[0], context) + else: + traverse_schema(choice, context) + elif schema_type == 'tagged-union': + for v in schema['choices'].values(): + traverse_schema(v, context) + elif schema_type == 'chain': + for step in schema['steps']: + traverse_schema(step, context) + elif schema_type == 'lax-or-strict': + traverse_schema(schema['lax_schema'], context) + traverse_schema(schema['strict_schema'], context) + elif schema_type == 'json-or-python': + traverse_schema(schema['json_schema'], context) + traverse_schema(schema['python_schema'], context) + elif schema_type in {'model-fields', 'typed-dict'}: + if 'extras_schema' in schema: + traverse_schema(schema['extras_schema'], context) + if 'computed_fields' in schema: + for s in schema['computed_fields']: + traverse_schema(s, context) + for s in schema['fields'].values(): + traverse_schema(s, context) + elif schema_type == 'dataclass-args': + if 'computed_fields' in schema: + for s in schema['computed_fields']: + traverse_schema(s, context) + for s in schema['fields']: + traverse_schema(s, context) + elif schema_type == 'arguments': + for s in schema['arguments_schema']: + traverse_schema(s['schema'], context) + if 'var_args_schema' in schema: + traverse_schema(schema['var_args_schema'], context) + if 'var_kwargs_schema' in schema: + traverse_schema(schema['var_kwargs_schema'], context) + elif schema_type == 'arguments-v3': + for s in schema['arguments_schema']: + traverse_schema(s['schema'], context) + elif schema_type == 'call': + traverse_schema(schema['arguments_schema'], context) + if 'return_schema' in schema: + traverse_schema(schema['return_schema'], context) + elif schema_type == 'computed-field': + traverse_schema(schema['return_schema'], context) + elif schema_type == 'function-before': + if 'schema' in schema: + traverse_schema(schema['schema'], context) + if 'json_schema_input_schema' in schema: + traverse_schema(schema['json_schema_input_schema'], context) + elif schema_type == 'function-plain': + # TODO duplicate schema types for serializers and validators, needs to be deduplicated. + if 'return_schema' in schema: + traverse_schema(schema['return_schema'], context) + if 'json_schema_input_schema' in schema: + traverse_schema(schema['json_schema_input_schema'], context) + elif schema_type == 'function-wrap': + # TODO duplicate schema types for serializers and validators, needs to be deduplicated. + if 'return_schema' in schema: + traverse_schema(schema['return_schema'], context) + if 'schema' in schema: + traverse_schema(schema['schema'], context) + if 'json_schema_input_schema' in schema: + traverse_schema(schema['json_schema_input_schema'], context) + else: + if 'schema' in schema: + traverse_schema(schema['schema'], context) + + if 'serialization' in schema: + traverse_schema(schema['serialization'], context) + traverse_metadata(schema, context) + + +def gather_schemas_for_cleaning(schema: CoreSchema, definitions: dict[str, CoreSchema]) -> GatherResult: + """Traverse the core schema and definitions and return the necessary information for schema cleaning. + + During the core schema traversing, any `'definition-ref'` schema is: + + - Validated: the reference must point to an existing definition. If this is not the case, a + `MissingDefinitionError` exception is raised. + - Stored in the context: the actual reference is stored in the context. Depending on whether + the `'definition-ref'` schema is encountered more that once, the schema itself is also + saved in the context to be inlined (i.e. replaced by the definition it points to). + """ + context = GatherContext(definitions) + traverse_schema(schema, context) + + return { + 'collected_references': context.collected_references, + 'deferred_discriminator_schemas': context.deferred_discriminator_schemas, + } diff --git a/lib/python3.12/site-packages/pydantic/_internal/_schema_generation_shared.py b/lib/python3.12/site-packages/pydantic/_internal/_schema_generation_shared.py new file mode 100644 index 0000000000000000000000000000000000000000..b231a82e1969822695bb3680217b66e73c81d017 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_schema_generation_shared.py @@ -0,0 +1,125 @@ +"""Types and utility functions used by various other internal tools.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, Literal + +from pydantic_core import core_schema + +from ..annotated_handlers import GetCoreSchemaHandler, GetJsonSchemaHandler + +if TYPE_CHECKING: + from ..json_schema import GenerateJsonSchema, JsonSchemaValue + from ._core_utils import CoreSchemaOrField + from ._generate_schema import GenerateSchema + from ._namespace_utils import NamespacesTuple + + GetJsonSchemaFunction = Callable[[CoreSchemaOrField, GetJsonSchemaHandler], JsonSchemaValue] + HandlerOverride = Callable[[CoreSchemaOrField], JsonSchemaValue] + + +class GenerateJsonSchemaHandler(GetJsonSchemaHandler): + """JsonSchemaHandler implementation that doesn't do ref unwrapping by default. + + This is used for any Annotated metadata so that we don't end up with conflicting + modifications to the definition schema. + + Used internally by Pydantic, please do not rely on this implementation. + See `GetJsonSchemaHandler` for the handler API. + """ + + def __init__(self, generate_json_schema: GenerateJsonSchema, handler_override: HandlerOverride | None) -> None: + self.generate_json_schema = generate_json_schema + self.handler = handler_override or generate_json_schema.generate_inner + self.mode = generate_json_schema.mode + + def __call__(self, core_schema: CoreSchemaOrField, /) -> JsonSchemaValue: + return self.handler(core_schema) + + def resolve_ref_schema(self, maybe_ref_json_schema: JsonSchemaValue) -> JsonSchemaValue: + """Resolves `$ref` in the json schema. + + This returns the input json schema if there is no `$ref` in json schema. + + Args: + maybe_ref_json_schema: The input json schema that may contains `$ref`. + + Returns: + Resolved json schema. + + Raises: + LookupError: If it can't find the definition for `$ref`. + """ + if '$ref' not in maybe_ref_json_schema: + return maybe_ref_json_schema + ref = maybe_ref_json_schema['$ref'] + json_schema = self.generate_json_schema.get_schema_from_definitions(ref) + if json_schema is None: + raise LookupError( + f'Could not find a ref for {ref}.' + ' Maybe you tried to call resolve_ref_schema from within a recursive model?' + ) + return json_schema + + +class CallbackGetCoreSchemaHandler(GetCoreSchemaHandler): + """Wrapper to use an arbitrary function as a `GetCoreSchemaHandler`. + + Used internally by Pydantic, please do not rely on this implementation. + See `GetCoreSchemaHandler` for the handler API. + """ + + def __init__( + self, + handler: Callable[[Any], core_schema.CoreSchema], + generate_schema: GenerateSchema, + ref_mode: Literal['to-def', 'unpack'] = 'to-def', + ) -> None: + self._handler = handler + self._generate_schema = generate_schema + self._ref_mode = ref_mode + + def __call__(self, source_type: Any, /) -> core_schema.CoreSchema: + schema = self._handler(source_type) + if self._ref_mode == 'to-def': + ref = schema.get('ref') + if ref is not None: + return self._generate_schema.defs.create_definition_reference_schema(schema) + return schema + else: # ref_mode = 'unpack' + return self.resolve_ref_schema(schema) + + def _get_types_namespace(self) -> NamespacesTuple: + return self._generate_schema._types_namespace + + def generate_schema(self, source_type: Any, /) -> core_schema.CoreSchema: + return self._generate_schema.generate_schema(source_type) + + @property + def field_name(self) -> str | None: + return self._generate_schema.field_name_stack.get() + + def resolve_ref_schema(self, maybe_ref_schema: core_schema.CoreSchema) -> core_schema.CoreSchema: + """Resolves reference in the core schema. + + Args: + maybe_ref_schema: The input core schema that may contains reference. + + Returns: + Resolved core schema. + + Raises: + LookupError: If it can't find the definition for reference. + """ + if maybe_ref_schema['type'] == 'definition-ref': + ref = maybe_ref_schema['schema_ref'] + definition = self._generate_schema.defs.get_schema_from_ref(ref) + if definition is None: + raise LookupError( + f'Could not find a ref for {ref}.' + ' Maybe you tried to call resolve_ref_schema from within a recursive model?' + ) + return definition + elif maybe_ref_schema['type'] == 'definitions': + return self.resolve_ref_schema(maybe_ref_schema['schema']) + return maybe_ref_schema diff --git a/lib/python3.12/site-packages/pydantic/_internal/_serializers.py b/lib/python3.12/site-packages/pydantic/_internal/_serializers.py new file mode 100644 index 0000000000000000000000000000000000000000..d0593218fd002acae235f244ef1c61b9d99ea1b8 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_serializers.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import collections +import collections.abc +import typing +from typing import Any + +from pydantic_core import PydanticOmit, core_schema + +SEQUENCE_ORIGIN_MAP: dict[Any, Any] = { + typing.Deque: collections.deque, # noqa: UP006 + collections.deque: collections.deque, + list: list, + typing.List: list, # noqa: UP006 + tuple: tuple, + typing.Tuple: tuple, # noqa: UP006 + set: set, + typing.AbstractSet: set, + typing.Set: set, # noqa: UP006 + frozenset: frozenset, + typing.FrozenSet: frozenset, # noqa: UP006 + typing.Sequence: list, + typing.MutableSequence: list, + typing.MutableSet: set, + # this doesn't handle subclasses of these + # parametrized typing.Set creates one of these + collections.abc.MutableSet: set, + collections.abc.Set: frozenset, +} + + +def serialize_sequence_via_list( + v: Any, handler: core_schema.SerializerFunctionWrapHandler, info: core_schema.SerializationInfo +) -> Any: + items: list[Any] = [] + + mapped_origin = SEQUENCE_ORIGIN_MAP.get(type(v), None) + if mapped_origin is None: + # we shouldn't hit this branch, should probably add a serialization error or something + return v + + for index, item in enumerate(v): + try: + v = handler(item, index) + except PydanticOmit: + pass + else: + items.append(v) + + if info.mode_is_json(): + return items + else: + return mapped_origin(items) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_signature.py b/lib/python3.12/site-packages/pydantic/_internal/_signature.py new file mode 100644 index 0000000000000000000000000000000000000000..977e5d2934016630108499c24c1f6fd5bf30179a --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_signature.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import dataclasses +from inspect import Parameter, Signature, signature +from typing import TYPE_CHECKING, Any, Callable + +from pydantic_core import PydanticUndefined + +from ._utils import is_valid_identifier + +if TYPE_CHECKING: + from ..config import ExtraValues + from ..fields import FieldInfo + + +# Copied over from stdlib dataclasses +class _HAS_DEFAULT_FACTORY_CLASS: + def __repr__(self): + return '' + + +_HAS_DEFAULT_FACTORY = _HAS_DEFAULT_FACTORY_CLASS() + + +def _field_name_for_signature(field_name: str, field_info: FieldInfo) -> str: + """Extract the correct name to use for the field when generating a signature. + + Assuming the field has a valid alias, this will return the alias. Otherwise, it will return the field name. + First priority is given to the alias, then the validation_alias, then the field name. + + Args: + field_name: The name of the field + field_info: The corresponding FieldInfo object. + + Returns: + The correct name to use when generating a signature. + """ + if isinstance(field_info.alias, str) and is_valid_identifier(field_info.alias): + return field_info.alias + if isinstance(field_info.validation_alias, str) and is_valid_identifier(field_info.validation_alias): + return field_info.validation_alias + + return field_name + + +def _process_param_defaults(param: Parameter) -> Parameter: + """Modify the signature for a parameter in a dataclass where the default value is a FieldInfo instance. + + Args: + param (Parameter): The parameter + + Returns: + Parameter: The custom processed parameter + """ + from ..fields import FieldInfo + + param_default = param.default + if isinstance(param_default, FieldInfo): + annotation = param.annotation + # Replace the annotation if appropriate + # inspect does "clever" things to show annotations as strings because we have + # `from __future__ import annotations` in main, we don't want that + if annotation == 'Any': + annotation = Any + + # Replace the field default + default = param_default.default + if default is PydanticUndefined: + if param_default.default_factory is PydanticUndefined: + default = Signature.empty + else: + # this is used by dataclasses to indicate a factory exists: + default = dataclasses._HAS_DEFAULT_FACTORY # type: ignore + return param.replace( + annotation=annotation, name=_field_name_for_signature(param.name, param_default), default=default + ) + return param + + +def _generate_signature_parameters( # noqa: C901 (ignore complexity, could use a refactor) + init: Callable[..., None], + fields: dict[str, FieldInfo], + validate_by_name: bool, + extra: ExtraValues | None, +) -> dict[str, Parameter]: + """Generate a mapping of parameter names to Parameter objects for a pydantic BaseModel or dataclass.""" + from itertools import islice + + present_params = signature(init).parameters.values() + merged_params: dict[str, Parameter] = {} + var_kw = None + use_var_kw = False + + for param in islice(present_params, 1, None): # skip self arg + # inspect does "clever" things to show annotations as strings because we have + # `from __future__ import annotations` in main, we don't want that + if fields.get(param.name): + # exclude params with init=False + if getattr(fields[param.name], 'init', True) is False: + continue + param = param.replace(name=_field_name_for_signature(param.name, fields[param.name])) + if param.annotation == 'Any': + param = param.replace(annotation=Any) + if param.kind is param.VAR_KEYWORD: + var_kw = param + continue + merged_params[param.name] = param + + if var_kw: # if custom init has no var_kw, fields which are not declared in it cannot be passed through + allow_names = validate_by_name + for field_name, field in fields.items(): + # when alias is a str it should be used for signature generation + param_name = _field_name_for_signature(field_name, field) + + if field_name in merged_params or param_name in merged_params: + continue + + if not is_valid_identifier(param_name): + if allow_names: + param_name = field_name + else: + use_var_kw = True + continue + + if field.is_required(): + default = Parameter.empty + elif field.default_factory is not None: + # Mimics stdlib dataclasses: + default = _HAS_DEFAULT_FACTORY + else: + default = field.default + merged_params[param_name] = Parameter( + param_name, + Parameter.KEYWORD_ONLY, + annotation=field.rebuild_annotation(), + default=default, + ) + + if extra == 'allow': + use_var_kw = True + + if var_kw and use_var_kw: + # Make sure the parameter for extra kwargs + # does not have the same name as a field + default_model_signature = [ + ('self', Parameter.POSITIONAL_ONLY), + ('data', Parameter.VAR_KEYWORD), + ] + if [(p.name, p.kind) for p in present_params] == default_model_signature: + # if this is the standard model signature, use extra_data as the extra args name + var_kw_name = 'extra_data' + else: + # else start from var_kw + var_kw_name = var_kw.name + + # generate a name that's definitely unique + while var_kw_name in fields: + var_kw_name += '_' + merged_params[var_kw_name] = var_kw.replace(name=var_kw_name) + + return merged_params + + +def generate_pydantic_signature( + init: Callable[..., None], + fields: dict[str, FieldInfo], + validate_by_name: bool, + extra: ExtraValues | None, + is_dataclass: bool = False, +) -> Signature: + """Generate signature for a pydantic BaseModel or dataclass. + + Args: + init: The class init. + fields: The model fields. + validate_by_name: The `validate_by_name` value of the config. + extra: The `extra` value of the config. + is_dataclass: Whether the model is a dataclass. + + Returns: + The dataclass/BaseModel subclass signature. + """ + merged_params = _generate_signature_parameters(init, fields, validate_by_name, extra) + + if is_dataclass: + merged_params = {k: _process_param_defaults(v) for k, v in merged_params.items()} + + return Signature(parameters=list(merged_params.values()), return_annotation=None) diff --git a/lib/python3.12/site-packages/pydantic/_internal/_typing_extra.py b/lib/python3.12/site-packages/pydantic/_internal/_typing_extra.py new file mode 100644 index 0000000000000000000000000000000000000000..4be1a09cb3214fdacae75ddc01d05746fc8ebc23 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_typing_extra.py @@ -0,0 +1,714 @@ +"""Logic for interacting with type annotations, mostly extensions, shims and hacks to wrap Python's typing module.""" + +from __future__ import annotations + +import collections.abc +import re +import sys +import types +import typing +from functools import partial +from typing import TYPE_CHECKING, Any, Callable, cast + +import typing_extensions +from typing_extensions import deprecated, get_args, get_origin +from typing_inspection import typing_objects +from typing_inspection.introspection import is_union_origin + +from pydantic.version import version_short + +from ._namespace_utils import GlobalsNamespace, MappingNamespace, NsResolver, get_module_ns_of + +if sys.version_info < (3, 10): + NoneType = type(None) + EllipsisType = type(Ellipsis) +else: + from types import EllipsisType as EllipsisType + from types import NoneType as NoneType + +if TYPE_CHECKING: + from pydantic import BaseModel + +# As per https://typing-extensions.readthedocs.io/en/latest/#runtime-use-of-types, +# always check for both `typing` and `typing_extensions` variants of a typing construct. +# (this is implemented differently than the suggested approach in the `typing_extensions` +# docs for performance). + + +_t_annotated = typing.Annotated +_te_annotated = typing_extensions.Annotated + + +def is_annotated(tp: Any, /) -> bool: + """Return whether the provided argument is a `Annotated` special form. + + ```python {test="skip" lint="skip"} + is_annotated(Annotated[int, ...]) + #> True + ``` + """ + origin = get_origin(tp) + return origin is _t_annotated or origin is _te_annotated + + +def annotated_type(tp: Any, /) -> Any | None: + """Return the type of the `Annotated` special form, or `None`.""" + return tp.__origin__ if typing_objects.is_annotated(get_origin(tp)) else None + + +def unpack_type(tp: Any, /) -> Any | None: + """Return the type wrapped by the `Unpack` special form, or `None`.""" + return get_args(tp)[0] if typing_objects.is_unpack(get_origin(tp)) else None + + +def is_hashable(tp: Any, /) -> bool: + """Return whether the provided argument is the `Hashable` class. + + ```python {test="skip" lint="skip"} + is_hashable(Hashable) + #> True + ``` + """ + # `get_origin` is documented as normalizing any typing-module aliases to `collections` classes, + # hence the second check: + return tp is collections.abc.Hashable or get_origin(tp) is collections.abc.Hashable + + +def is_callable(tp: Any, /) -> bool: + """Return whether the provided argument is a `Callable`, parametrized or not. + + ```python {test="skip" lint="skip"} + is_callable(Callable[[int], str]) + #> True + is_callable(typing.Callable) + #> True + is_callable(collections.abc.Callable) + #> True + ``` + """ + # `get_origin` is documented as normalizing any typing-module aliases to `collections` classes, + # hence the second check: + return tp is collections.abc.Callable or get_origin(tp) is collections.abc.Callable + + +_classvar_re = re.compile(r'((\w+\.)?Annotated\[)?(\w+\.)?ClassVar\[') + + +def is_classvar_annotation(tp: Any, /) -> bool: + """Return whether the provided argument represents a class variable annotation. + + Although not explicitly stated by the typing specification, `ClassVar` can be used + inside `Annotated` and as such, this function checks for this specific scenario. + + Because this function is used to detect class variables before evaluating forward references + (or because evaluation failed), we also implement a naive regex match implementation. This is + required because class variables are inspected before fields are collected, so we try to be + as accurate as possible. + """ + if typing_objects.is_classvar(tp): + return True + + origin = get_origin(tp) + + if typing_objects.is_classvar(origin): + return True + + if typing_objects.is_annotated(origin): + annotated_type = tp.__origin__ + if typing_objects.is_classvar(annotated_type) or typing_objects.is_classvar(get_origin(annotated_type)): + return True + + str_ann: str | None = None + if isinstance(tp, typing.ForwardRef): + str_ann = tp.__forward_arg__ + if isinstance(tp, str): + str_ann = tp + + if str_ann is not None and _classvar_re.match(str_ann): + # stdlib dataclasses do something similar, although a bit more advanced + # (see `dataclass._is_type`). + return True + + return False + + +_t_final = typing.Final +_te_final = typing_extensions.Final + + +# TODO implement `is_finalvar_annotation` as Final can be wrapped with other special forms: +def is_finalvar(tp: Any, /) -> bool: + """Return whether the provided argument is a `Final` special form, parametrized or not. + + ```python {test="skip" lint="skip"} + is_finalvar(Final[int]) + #> True + is_finalvar(Final) + #> True + """ + # Final is not necessarily parametrized: + if tp is _t_final or tp is _te_final: + return True + origin = get_origin(tp) + return origin is _t_final or origin is _te_final + + +_NONE_TYPES: tuple[Any, ...] = (None, NoneType, typing.Literal[None], typing_extensions.Literal[None]) + + +def is_none_type(tp: Any, /) -> bool: + """Return whether the argument represents the `None` type as part of an annotation. + + ```python {test="skip" lint="skip"} + is_none_type(None) + #> True + is_none_type(NoneType) + #> True + is_none_type(Literal[None]) + #> True + is_none_type(type[None]) + #> False + """ + return tp in _NONE_TYPES + + +def is_namedtuple(tp: Any, /) -> bool: + """Return whether the provided argument is a named tuple class. + + The class can be created using `typing.NamedTuple` or `collections.namedtuple`. + Parametrized generic classes are *not* assumed to be named tuples. + """ + from ._utils import lenient_issubclass # circ. import + + return lenient_issubclass(tp, tuple) and hasattr(tp, '_fields') + + +# TODO In 2.12, delete this export. It is currently defined only to not break +# pydantic-settings which relies on it: +origin_is_union = is_union_origin + + +def is_generic_alias(tp: Any, /) -> bool: + return isinstance(tp, (types.GenericAlias, typing._GenericAlias)) # pyright: ignore[reportAttributeAccessIssue] + + +# TODO: Ideally, we should avoid relying on the private `typing` constructs: + +if sys.version_info < (3, 10): + WithArgsTypes: tuple[Any, ...] = (typing._GenericAlias, types.GenericAlias) # pyright: ignore[reportAttributeAccessIssue] +else: + WithArgsTypes: tuple[Any, ...] = (typing._GenericAlias, types.GenericAlias, types.UnionType) # pyright: ignore[reportAttributeAccessIssue] + + +# Similarly, we shouldn't rely on this `_Final` class, which is even more private than `_GenericAlias`: +typing_base: Any = typing._Final # pyright: ignore[reportAttributeAccessIssue] + + +### Annotation evaluations functions: + + +def parent_frame_namespace(*, parent_depth: int = 2, force: bool = False) -> dict[str, Any] | None: + """Fetch the local namespace of the parent frame where this function is called. + + Using this function is mostly useful to resolve forward annotations pointing to members defined in a local namespace, + such as assignments inside a function. Using the standard library tools, it is currently not possible to resolve + such annotations: + + ```python {lint="skip" test="skip"} + from typing import get_type_hints + + def func() -> None: + Alias = int + + class C: + a: 'Alias' + + # Raises a `NameError: 'Alias' is not defined` + get_type_hints(C) + ``` + + Pydantic uses this function when a Pydantic model is being defined to fetch the parent frame locals. However, + this only allows us to fetch the parent frame namespace and not other parents (e.g. a model defined in a function, + itself defined in another function). Inspecting the next outer frames (using `f_back`) is not reliable enough + (see https://discuss.python.org/t/20659). + + Because this function is mostly used to better resolve forward annotations, nothing is returned if the parent frame's + code object is defined at the module level. In this case, the locals of the frame will be the same as the module + globals where the class is defined (see `_namespace_utils.get_module_ns_of`). However, if you still want to fetch + the module globals (e.g. when rebuilding a model, where the frame where the rebuild call is performed might contain + members that you want to use for forward annotations evaluation), you can use the `force` parameter. + + Args: + parent_depth: The depth at which to get the frame. Defaults to 2, meaning the parent frame where this function + is called will be used. + force: Whether to always return the frame locals, even if the frame's code object is defined at the module level. + + Returns: + The locals of the namespace, or `None` if it was skipped as per the described logic. + """ + frame = sys._getframe(parent_depth) + + if frame.f_code.co_name.startswith('`, + # and we need to skip this frame as it is irrelevant. + frame = cast(types.FrameType, frame.f_back) # guaranteed to not be `None` + + # note, we don't copy frame.f_locals here (or during the last return call), because we don't expect the namespace to be + # modified down the line if this becomes a problem, we could implement some sort of frozen mapping structure to enforce this. + if force: + return frame.f_locals + + # If either of the following conditions are true, the class is defined at the top module level. + # To better understand why we need both of these checks, see + # https://github.com/pydantic/pydantic/pull/10113#discussion_r1714981531. + if frame.f_back is None or frame.f_code.co_name == '': + return None + + return frame.f_locals + + +def _type_convert(arg: Any) -> Any: + """Convert `None` to `NoneType` and strings to `ForwardRef` instances. + + This is a backport of the private `typing._type_convert` function. When + evaluating a type, `ForwardRef._evaluate` ends up being called, and is + responsible for making this conversion. However, we still have to apply + it for the first argument passed to our type evaluation functions, similarly + to the `typing.get_type_hints` function. + """ + if arg is None: + return NoneType + if isinstance(arg, str): + # Like `typing.get_type_hints`, assume the arg can be in any context, + # hence the proper `is_argument` and `is_class` args: + return _make_forward_ref(arg, is_argument=False, is_class=True) + return arg + + +def get_model_type_hints( + obj: type[BaseModel], + *, + ns_resolver: NsResolver | None = None, +) -> dict[str, tuple[Any, bool]]: + """Collect annotations from a Pydantic model class, including those from parent classes. + + Args: + obj: The Pydantic model to inspect. + ns_resolver: A namespace resolver instance to use. Defaults to an empty instance. + + Returns: + A dictionary mapping annotation names to a two-tuple: the first element is the evaluated + type or the original annotation if a `NameError` occurred, the second element is a boolean + indicating if whether the evaluation succeeded. + """ + hints: dict[str, Any] | dict[str, tuple[Any, bool]] = {} + ns_resolver = ns_resolver or NsResolver() + + for base in reversed(obj.__mro__): + ann: dict[str, Any] | None = base.__dict__.get('__annotations__') + if not ann or isinstance(ann, types.GetSetDescriptorType): + continue + with ns_resolver.push(base): + globalns, localns = ns_resolver.types_namespace + for name, value in ann.items(): + if name.startswith('_'): + # For private attributes, we only need the annotation to detect the `ClassVar` special form. + # For this reason, we still try to evaluate it, but we also catch any possible exception (on + # top of the `NameError`s caught in `try_eval_type`) that could happen so that users are free + # to use any kind of forward annotation for private fields (e.g. circular imports, new typing + # syntax, etc). + try: + hints[name] = try_eval_type(value, globalns, localns) + except Exception: + hints[name] = (value, False) + else: + hints[name] = try_eval_type(value, globalns, localns) + return hints + + +def get_cls_type_hints( + obj: type[Any], + *, + ns_resolver: NsResolver | None = None, +) -> dict[str, Any]: + """Collect annotations from a class, including those from parent classes. + + Args: + obj: The class to inspect. + ns_resolver: A namespace resolver instance to use. Defaults to an empty instance. + """ + hints: dict[str, Any] | dict[str, tuple[Any, bool]] = {} + ns_resolver = ns_resolver or NsResolver() + + for base in reversed(obj.__mro__): + ann: dict[str, Any] | None = base.__dict__.get('__annotations__') + if not ann or isinstance(ann, types.GetSetDescriptorType): + continue + with ns_resolver.push(base): + globalns, localns = ns_resolver.types_namespace + for name, value in ann.items(): + hints[name] = eval_type(value, globalns, localns) + return hints + + +def try_eval_type( + value: Any, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, +) -> tuple[Any, bool]: + """Try evaluating the annotation using the provided namespaces. + + Args: + value: The value to evaluate. If `None`, it will be replaced by `type[None]`. If an instance + of `str`, it will be converted to a `ForwardRef`. + localns: The global namespace to use during annotation evaluation. + globalns: The local namespace to use during annotation evaluation. + + Returns: + A two-tuple containing the possibly evaluated type and a boolean indicating + whether the evaluation succeeded or not. + """ + value = _type_convert(value) + + try: + return eval_type_backport(value, globalns, localns), True + except NameError: + return value, False + + +def eval_type( + value: Any, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, +) -> Any: + """Evaluate the annotation using the provided namespaces. + + Args: + value: The value to evaluate. If `None`, it will be replaced by `type[None]`. If an instance + of `str`, it will be converted to a `ForwardRef`. + localns: The global namespace to use during annotation evaluation. + globalns: The local namespace to use during annotation evaluation. + """ + value = _type_convert(value) + return eval_type_backport(value, globalns, localns) + + +@deprecated( + '`eval_type_lenient` is deprecated, use `try_eval_type` instead.', + category=None, +) +def eval_type_lenient( + value: Any, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, +) -> Any: + ev, _ = try_eval_type(value, globalns, localns) + return ev + + +def eval_type_backport( + value: Any, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, + type_params: tuple[Any, ...] | None = None, +) -> Any: + """An enhanced version of `typing._eval_type` which will fall back to using the `eval_type_backport` + package if it's installed to let older Python versions use newer typing constructs. + + Specifically, this transforms `X | Y` into `typing.Union[X, Y]` and `list[X]` into `typing.List[X]` + (as well as all the types made generic in PEP 585) if the original syntax is not supported in the + current Python version. + + This function will also display a helpful error if the value passed fails to evaluate. + """ + try: + return _eval_type_backport(value, globalns, localns, type_params) + except TypeError as e: + if 'Unable to evaluate type annotation' in str(e): + raise + + # If it is a `TypeError` and value isn't a `ForwardRef`, it would have failed during annotation definition. + # Thus we assert here for type checking purposes: + assert isinstance(value, typing.ForwardRef) + + message = f'Unable to evaluate type annotation {value.__forward_arg__!r}.' + if sys.version_info >= (3, 11): + e.add_note(message) + raise + else: + raise TypeError(message) from e + except RecursionError as e: + # TODO ideally recursion errors should be checked in `eval_type` above, but `eval_type_backport` + # is used directly in some places. + message = ( + "If you made use of an implicit recursive type alias (e.g. `MyType = list['MyType']), " + 'consider using PEP 695 type aliases instead. For more details, refer to the documentation: ' + f'https://docs.pydantic.dev/{version_short()}/concepts/types/#named-recursive-types' + ) + if sys.version_info >= (3, 11): + e.add_note(message) + raise + else: + raise RecursionError(f'{e.args[0]}\n{message}') + + +def _eval_type_backport( + value: Any, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, + type_params: tuple[Any, ...] | None = None, +) -> Any: + try: + return _eval_type(value, globalns, localns, type_params) + except TypeError as e: + if not (isinstance(value, typing.ForwardRef) and is_backport_fixable_error(e)): + raise + + try: + from eval_type_backport import eval_type_backport + except ImportError: + raise TypeError( + f'Unable to evaluate type annotation {value.__forward_arg__!r}. If you are making use ' + 'of the new typing syntax (unions using `|` since Python 3.10 or builtins subscripting ' + 'since Python 3.9), you should either replace the use of new syntax with the existing ' + '`typing` constructs or install the `eval_type_backport` package.' + ) from e + + return eval_type_backport( + value, + globalns, + localns, # pyright: ignore[reportArgumentType], waiting on a new `eval_type_backport` release. + try_default=False, + ) + + +def _eval_type( + value: Any, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, + type_params: tuple[Any, ...] | None = None, +) -> Any: + if sys.version_info >= (3, 13): + return typing._eval_type( # type: ignore + value, globalns, localns, type_params=type_params + ) + else: + return typing._eval_type( # type: ignore + value, globalns, localns + ) + + +def is_backport_fixable_error(e: TypeError) -> bool: + msg = str(e) + + return sys.version_info < (3, 10) and msg.startswith('unsupported operand type(s) for |: ') + + +def get_function_type_hints( + function: Callable[..., Any], + *, + include_keys: set[str] | None = None, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, +) -> dict[str, Any]: + """Return type hints for a function. + + This is similar to the `typing.get_type_hints` function, with a few differences: + - Support `functools.partial` by using the underlying `func` attribute. + - Do not wrap type annotation of a parameter with `Optional` if it has a default value of `None` + (related bug: https://github.com/python/cpython/issues/90353, only fixed in 3.11+). + """ + try: + if isinstance(function, partial): + annotations = function.func.__annotations__ + else: + annotations = function.__annotations__ + except AttributeError: + # Some functions (e.g. builtins) don't have annotations: + return {} + + if globalns is None: + globalns = get_module_ns_of(function) + type_params: tuple[Any, ...] | None = None + if localns is None: + # If localns was specified, it is assumed to already contain type params. This is because + # Pydantic has more advanced logic to do so (see `_namespace_utils.ns_for_function`). + type_params = getattr(function, '__type_params__', ()) + + type_hints = {} + for name, value in annotations.items(): + if include_keys is not None and name not in include_keys: + continue + if value is None: + value = NoneType + elif isinstance(value, str): + value = _make_forward_ref(value) + + type_hints[name] = eval_type_backport(value, globalns, localns, type_params) + + return type_hints + + +if sys.version_info < (3, 9, 8) or (3, 10) <= sys.version_info < (3, 10, 1): + + def _make_forward_ref( + arg: Any, + is_argument: bool = True, + *, + is_class: bool = False, + ) -> typing.ForwardRef: + """Wrapper for ForwardRef that accounts for the `is_class` argument missing in older versions. + The `module` argument is omitted as it breaks <3.9.8, =3.10.0 and isn't used in the calls below. + + See https://github.com/python/cpython/pull/28560 for some background. + The backport happened on 3.9.8, see: + https://github.com/pydantic/pydantic/discussions/6244#discussioncomment-6275458, + and on 3.10.1 for the 3.10 branch, see: + https://github.com/pydantic/pydantic/issues/6912 + + Implemented as EAFP with memory. + """ + return typing.ForwardRef(arg, is_argument) + +else: + _make_forward_ref = typing.ForwardRef + + +if sys.version_info >= (3, 10): + get_type_hints = typing.get_type_hints + +else: + """ + For older versions of python, we have a custom implementation of `get_type_hints` which is a close as possible to + the implementation in CPython 3.10.8. + """ + + @typing.no_type_check + def get_type_hints( # noqa: C901 + obj: Any, + globalns: dict[str, Any] | None = None, + localns: dict[str, Any] | None = None, + include_extras: bool = False, + ) -> dict[str, Any]: # pragma: no cover + """Taken verbatim from python 3.10.8 unchanged, except: + * type annotations of the function definition above. + * prefixing `typing.` where appropriate + * Use `_make_forward_ref` instead of `typing.ForwardRef` to handle the `is_class` argument. + + https://github.com/python/cpython/blob/aaaf5174241496afca7ce4d4584570190ff972fe/Lib/typing.py#L1773-L1875 + + DO NOT CHANGE THIS METHOD UNLESS ABSOLUTELY NECESSARY. + ====================================================== + + Return type hints for an object. + + This is often the same as obj.__annotations__, but it handles + forward references encoded as string literals, adds Optional[t] if a + default value equal to None is set and recursively replaces all + 'Annotated[T, ...]' with 'T' (unless 'include_extras=True'). + + The argument may be a module, class, method, or function. The annotations + are returned as a dictionary. For classes, annotations include also + inherited members. + + TypeError is raised if the argument is not of a type that can contain + annotations, and an empty dictionary is returned if no annotations are + present. + + BEWARE -- the behavior of globalns and localns is counterintuitive + (unless you are familiar with how eval() and exec() work). The + search order is locals first, then globals. + + - If no dict arguments are passed, an attempt is made to use the + globals from obj (or the respective module's globals for classes), + and these are also used as the locals. If the object does not appear + to have globals, an empty dictionary is used. For classes, the search + order is globals first then locals. + + - If one dict argument is passed, it is used for both globals and + locals. + + - If two dict arguments are passed, they specify globals and + locals, respectively. + """ + if getattr(obj, '__no_type_check__', None): + return {} + # Classes require a special treatment. + if isinstance(obj, type): + hints = {} + for base in reversed(obj.__mro__): + if globalns is None: + base_globals = getattr(sys.modules.get(base.__module__, None), '__dict__', {}) + else: + base_globals = globalns + ann = base.__dict__.get('__annotations__', {}) + if isinstance(ann, types.GetSetDescriptorType): + ann = {} + base_locals = dict(vars(base)) if localns is None else localns + if localns is None and globalns is None: + # This is surprising, but required. Before Python 3.10, + # get_type_hints only evaluated the globalns of + # a class. To maintain backwards compatibility, we reverse + # the globalns and localns order so that eval() looks into + # *base_globals* first rather than *base_locals*. + # This only affects ForwardRefs. + base_globals, base_locals = base_locals, base_globals + for name, value in ann.items(): + if value is None: + value = type(None) + if isinstance(value, str): + value = _make_forward_ref(value, is_argument=False, is_class=True) + + value = eval_type_backport(value, base_globals, base_locals) + hints[name] = value + if not include_extras and hasattr(typing, '_strip_annotations'): + return { + k: typing._strip_annotations(t) # type: ignore + for k, t in hints.items() + } + else: + return hints + + if globalns is None: + if isinstance(obj, types.ModuleType): + globalns = obj.__dict__ + else: + nsobj = obj + # Find globalns for the unwrapped object. + while hasattr(nsobj, '__wrapped__'): + nsobj = nsobj.__wrapped__ + globalns = getattr(nsobj, '__globals__', {}) + if localns is None: + localns = globalns + elif localns is None: + localns = globalns + hints = getattr(obj, '__annotations__', None) + if hints is None: + # Return empty annotations for something that _could_ have them. + if isinstance(obj, typing._allowed_types): # type: ignore + return {} + else: + raise TypeError(f'{obj!r} is not a module, class, method, or function.') + defaults = typing._get_defaults(obj) # type: ignore + hints = dict(hints) + for name, value in hints.items(): + if value is None: + value = type(None) + if isinstance(value, str): + # class-level forward refs were handled above, this must be either + # a module-level annotation or a function argument annotation + + value = _make_forward_ref( + value, + is_argument=not isinstance(obj, types.ModuleType), + is_class=False, + ) + value = eval_type_backport(value, globalns, localns) + if name in defaults and defaults[name] is None: + value = typing.Optional[value] + hints[name] = value + return hints if include_extras else {k: typing._strip_annotations(t) for k, t in hints.items()} # type: ignore diff --git a/lib/python3.12/site-packages/pydantic/_internal/_utils.py b/lib/python3.12/site-packages/pydantic/_internal/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f33464987bc985f1a0deea4ed4653b99f7a8d731 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_utils.py @@ -0,0 +1,431 @@ +"""Bucket of reusable internal utilities. + +This should be reduced as much as possible with functions only used in one place, moved to that place. +""" + +from __future__ import annotations as _annotations + +import dataclasses +import keyword +import sys +import typing +import warnings +import weakref +from collections import OrderedDict, defaultdict, deque +from collections.abc import Mapping +from copy import deepcopy +from functools import cached_property +from inspect import Parameter +from itertools import zip_longest +from types import BuiltinFunctionType, CodeType, FunctionType, GeneratorType, LambdaType, ModuleType +from typing import Any, Callable, Generic, TypeVar, overload + +from typing_extensions import TypeAlias, TypeGuard, deprecated + +from pydantic import PydanticDeprecatedSince211 + +from . import _repr, _typing_extra +from ._import_utils import import_cached_base_model + +if typing.TYPE_CHECKING: + MappingIntStrAny: TypeAlias = 'typing.Mapping[int, Any] | typing.Mapping[str, Any]' + AbstractSetIntStr: TypeAlias = 'typing.AbstractSet[int] | typing.AbstractSet[str]' + from ..main import BaseModel + + +# these are types that are returned unchanged by deepcopy +IMMUTABLE_NON_COLLECTIONS_TYPES: set[type[Any]] = { + int, + float, + complex, + str, + bool, + bytes, + type, + _typing_extra.NoneType, + FunctionType, + BuiltinFunctionType, + LambdaType, + weakref.ref, + CodeType, + # note: including ModuleType will differ from behaviour of deepcopy by not producing error. + # It might be not a good idea in general, but considering that this function used only internally + # against default values of fields, this will allow to actually have a field with module as default value + ModuleType, + NotImplemented.__class__, + Ellipsis.__class__, +} + +# these are types that if empty, might be copied with simple copy() instead of deepcopy() +BUILTIN_COLLECTIONS: set[type[Any]] = { + list, + set, + tuple, + frozenset, + dict, + OrderedDict, + defaultdict, + deque, +} + + +def can_be_positional(param: Parameter) -> bool: + """Return whether the parameter accepts a positional argument. + + ```python {test="skip" lint="skip"} + def func(a, /, b, *, c): + pass + + params = inspect.signature(func).parameters + can_be_positional(params['a']) + #> True + can_be_positional(params['b']) + #> True + can_be_positional(params['c']) + #> False + ``` + """ + return param.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD) + + +def sequence_like(v: Any) -> bool: + return isinstance(v, (list, tuple, set, frozenset, GeneratorType, deque)) + + +def lenient_isinstance(o: Any, class_or_tuple: type[Any] | tuple[type[Any], ...] | None) -> bool: # pragma: no cover + try: + return isinstance(o, class_or_tuple) # type: ignore[arg-type] + except TypeError: + return False + + +def lenient_issubclass(cls: Any, class_or_tuple: Any) -> bool: # pragma: no cover + try: + return isinstance(cls, type) and issubclass(cls, class_or_tuple) + except TypeError: + if isinstance(cls, _typing_extra.WithArgsTypes): + return False + raise # pragma: no cover + + +def is_model_class(cls: Any) -> TypeGuard[type[BaseModel]]: + """Returns true if cls is a _proper_ subclass of BaseModel, and provides proper type-checking, + unlike raw calls to lenient_issubclass. + """ + BaseModel = import_cached_base_model() + + return lenient_issubclass(cls, BaseModel) and cls is not BaseModel + + +def is_valid_identifier(identifier: str) -> bool: + """Checks that a string is a valid identifier and not a Python keyword. + :param identifier: The identifier to test. + :return: True if the identifier is valid. + """ + return identifier.isidentifier() and not keyword.iskeyword(identifier) + + +KeyType = TypeVar('KeyType') + + +def deep_update(mapping: dict[KeyType, Any], *updating_mappings: dict[KeyType, Any]) -> dict[KeyType, Any]: + updated_mapping = mapping.copy() + for updating_mapping in updating_mappings: + for k, v in updating_mapping.items(): + if k in updated_mapping and isinstance(updated_mapping[k], dict) and isinstance(v, dict): + updated_mapping[k] = deep_update(updated_mapping[k], v) + else: + updated_mapping[k] = v + return updated_mapping + + +def update_not_none(mapping: dict[Any, Any], **update: Any) -> None: + mapping.update({k: v for k, v in update.items() if v is not None}) + + +T = TypeVar('T') + + +def unique_list( + input_list: list[T] | tuple[T, ...], + *, + name_factory: typing.Callable[[T], str] = str, +) -> list[T]: + """Make a list unique while maintaining order. + We update the list if another one with the same name is set + (e.g. model validator overridden in subclass). + """ + result: list[T] = [] + result_names: list[str] = [] + for v in input_list: + v_name = name_factory(v) + if v_name not in result_names: + result_names.append(v_name) + result.append(v) + else: + result[result_names.index(v_name)] = v + + return result + + +class ValueItems(_repr.Representation): + """Class for more convenient calculation of excluded or included fields on values.""" + + __slots__ = ('_items', '_type') + + def __init__(self, value: Any, items: AbstractSetIntStr | MappingIntStrAny) -> None: + items = self._coerce_items(items) + + if isinstance(value, (list, tuple)): + items = self._normalize_indexes(items, len(value)) # type: ignore + + self._items: MappingIntStrAny = items # type: ignore + + def is_excluded(self, item: Any) -> bool: + """Check if item is fully excluded. + + :param item: key or index of a value + """ + return self.is_true(self._items.get(item)) + + def is_included(self, item: Any) -> bool: + """Check if value is contained in self._items. + + :param item: key or index of value + """ + return item in self._items + + def for_element(self, e: int | str) -> AbstractSetIntStr | MappingIntStrAny | None: + """:param e: key or index of element on value + :return: raw values for element if self._items is dict and contain needed element + """ + item = self._items.get(e) # type: ignore + return item if not self.is_true(item) else None + + def _normalize_indexes(self, items: MappingIntStrAny, v_length: int) -> dict[int | str, Any]: + """:param items: dict or set of indexes which will be normalized + :param v_length: length of sequence indexes of which will be + + >>> self._normalize_indexes({0: True, -2: True, -1: True}, 4) + {0: True, 2: True, 3: True} + >>> self._normalize_indexes({'__all__': True}, 4) + {0: True, 1: True, 2: True, 3: True} + """ + normalized_items: dict[int | str, Any] = {} + all_items = None + for i, v in items.items(): + if not (isinstance(v, typing.Mapping) or isinstance(v, typing.AbstractSet) or self.is_true(v)): + raise TypeError(f'Unexpected type of exclude value for index "{i}" {v.__class__}') + if i == '__all__': + all_items = self._coerce_value(v) + continue + if not isinstance(i, int): + raise TypeError( + 'Excluding fields from a sequence of sub-models or dicts must be performed index-wise: ' + 'expected integer keys or keyword "__all__"' + ) + normalized_i = v_length + i if i < 0 else i + normalized_items[normalized_i] = self.merge(v, normalized_items.get(normalized_i)) + + if not all_items: + return normalized_items + if self.is_true(all_items): + for i in range(v_length): + normalized_items.setdefault(i, ...) + return normalized_items + for i in range(v_length): + normalized_item = normalized_items.setdefault(i, {}) + if not self.is_true(normalized_item): + normalized_items[i] = self.merge(all_items, normalized_item) + return normalized_items + + @classmethod + def merge(cls, base: Any, override: Any, intersect: bool = False) -> Any: + """Merge a `base` item with an `override` item. + + Both `base` and `override` are converted to dictionaries if possible. + Sets are converted to dictionaries with the sets entries as keys and + Ellipsis as values. + + Each key-value pair existing in `base` is merged with `override`, + while the rest of the key-value pairs are updated recursively with this function. + + Merging takes place based on the "union" of keys if `intersect` is + set to `False` (default) and on the intersection of keys if + `intersect` is set to `True`. + """ + override = cls._coerce_value(override) + base = cls._coerce_value(base) + if override is None: + return base + if cls.is_true(base) or base is None: + return override + if cls.is_true(override): + return base if intersect else override + + # intersection or union of keys while preserving ordering: + if intersect: + merge_keys = [k for k in base if k in override] + [k for k in override if k in base] + else: + merge_keys = list(base) + [k for k in override if k not in base] + + merged: dict[int | str, Any] = {} + for k in merge_keys: + merged_item = cls.merge(base.get(k), override.get(k), intersect=intersect) + if merged_item is not None: + merged[k] = merged_item + + return merged + + @staticmethod + def _coerce_items(items: AbstractSetIntStr | MappingIntStrAny) -> MappingIntStrAny: + if isinstance(items, typing.Mapping): + pass + elif isinstance(items, typing.AbstractSet): + items = dict.fromkeys(items, ...) # type: ignore + else: + class_name = getattr(items, '__class__', '???') + raise TypeError(f'Unexpected type of exclude value {class_name}') + return items # type: ignore + + @classmethod + def _coerce_value(cls, value: Any) -> Any: + if value is None or cls.is_true(value): + return value + return cls._coerce_items(value) + + @staticmethod + def is_true(v: Any) -> bool: + return v is True or v is ... + + def __repr_args__(self) -> _repr.ReprArgs: + return [(None, self._items)] + + +if typing.TYPE_CHECKING: + + def LazyClassAttribute(name: str, get_value: Callable[[], T]) -> T: ... + +else: + + class LazyClassAttribute: + """A descriptor exposing an attribute only accessible on a class (hidden from instances). + + The attribute is lazily computed and cached during the first access. + """ + + def __init__(self, name: str, get_value: Callable[[], Any]) -> None: + self.name = name + self.get_value = get_value + + @cached_property + def value(self) -> Any: + return self.get_value() + + def __get__(self, instance: Any, owner: type[Any]) -> None: + if instance is None: + return self.value + raise AttributeError(f'{self.name!r} attribute of {owner.__name__!r} is class-only') + + +Obj = TypeVar('Obj') + + +def smart_deepcopy(obj: Obj) -> Obj: + """Return type as is for immutable built-in types + Use obj.copy() for built-in empty collections + Use copy.deepcopy() for non-empty collections and unknown objects. + """ + obj_type = obj.__class__ + if obj_type in IMMUTABLE_NON_COLLECTIONS_TYPES: + return obj # fastest case: obj is immutable and not collection therefore will not be copied anyway + try: + if not obj and obj_type in BUILTIN_COLLECTIONS: + # faster way for empty collections, no need to copy its members + return obj if obj_type is tuple else obj.copy() # tuple doesn't have copy method # type: ignore + except (TypeError, ValueError, RuntimeError): + # do we really dare to catch ALL errors? Seems a bit risky + pass + + return deepcopy(obj) # slowest way when we actually might need a deepcopy + + +_SENTINEL = object() + + +def all_identical(left: typing.Iterable[Any], right: typing.Iterable[Any]) -> bool: + """Check that the items of `left` are the same objects as those in `right`. + + >>> a, b = object(), object() + >>> all_identical([a, b, a], [a, b, a]) + True + >>> all_identical([a, b, [a]], [a, b, [a]]) # new list object, while "equal" is not "identical" + False + """ + for left_item, right_item in zip_longest(left, right, fillvalue=_SENTINEL): + if left_item is not right_item: + return False + return True + + +@dataclasses.dataclass(frozen=True) +class SafeGetItemProxy: + """Wrapper redirecting `__getitem__` to `get` with a sentinel value as default + + This makes is safe to use in `operator.itemgetter` when some keys may be missing + """ + + # Define __slots__manually for performances + # @dataclasses.dataclass() only support slots=True in python>=3.10 + __slots__ = ('wrapped',) + + wrapped: Mapping[str, Any] + + def __getitem__(self, key: str, /) -> Any: + return self.wrapped.get(key, _SENTINEL) + + # required to pass the object to operator.itemgetter() instances due to a quirk of typeshed + # https://github.com/python/mypy/issues/13713 + # https://github.com/python/typeshed/pull/8785 + # Since this is typing-only, hide it in a typing.TYPE_CHECKING block + if typing.TYPE_CHECKING: + + def __contains__(self, key: str, /) -> bool: + return self.wrapped.__contains__(key) + + +_ModelT = TypeVar('_ModelT', bound='BaseModel') +_RT = TypeVar('_RT') + + +class deprecated_instance_property(Generic[_ModelT, _RT]): + """A decorator exposing the decorated class method as a property, with a warning on instance access. + + This decorator takes a class method defined on the `BaseModel` class and transforms it into + an attribute. The attribute can be accessed on both the class and instances of the class. If accessed + via an instance, a deprecation warning is emitted stating that instance access will be removed in V3. + """ + + def __init__(self, fget: Callable[[type[_ModelT]], _RT], /) -> None: + # Note: fget should be a classmethod: + self.fget = fget + + @overload + def __get__(self, instance: None, objtype: type[_ModelT]) -> _RT: ... + @overload + @deprecated( + 'Accessing this attribute on the instance is deprecated, and will be removed in Pydantic V3. ' + 'Instead, you should access this attribute from the model class.', + category=None, + ) + def __get__(self, instance: _ModelT, objtype: type[_ModelT]) -> _RT: ... + def __get__(self, instance: _ModelT | None, objtype: type[_ModelT]) -> _RT: + if instance is not None: + attr_name = self.fget.__name__ if sys.version_info >= (3, 10) else self.fget.__func__.__name__ + warnings.warn( + f'Accessing the {attr_name!r} attribute on the instance is deprecated. ' + 'Instead, you should access this attribute from the model class.', + category=PydanticDeprecatedSince211, + stacklevel=2, + ) + return self.fget.__get__(instance, objtype)() diff --git a/lib/python3.12/site-packages/pydantic/_internal/_validate_call.py b/lib/python3.12/site-packages/pydantic/_internal/_validate_call.py new file mode 100644 index 0000000000000000000000000000000000000000..ab82832fe4dfb246ee8d06afcbd5b70652f88e27 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_validate_call.py @@ -0,0 +1,140 @@ +from __future__ import annotations as _annotations + +import functools +import inspect +from collections.abc import Awaitable +from functools import partial +from typing import Any, Callable + +import pydantic_core + +from ..config import ConfigDict +from ..plugin._schema_validator import create_schema_validator +from ._config import ConfigWrapper +from ._generate_schema import GenerateSchema, ValidateCallSupportedTypes +from ._namespace_utils import MappingNamespace, NsResolver, ns_for_function + + +def extract_function_name(func: ValidateCallSupportedTypes) -> str: + """Extract the name of a `ValidateCallSupportedTypes` object.""" + return f'partial({func.func.__name__})' if isinstance(func, functools.partial) else func.__name__ + + +def extract_function_qualname(func: ValidateCallSupportedTypes) -> str: + """Extract the qualname of a `ValidateCallSupportedTypes` object.""" + return f'partial({func.func.__qualname__})' if isinstance(func, functools.partial) else func.__qualname__ + + +def update_wrapper_attributes(wrapped: ValidateCallSupportedTypes, wrapper: Callable[..., Any]): + """Update the `wrapper` function with the attributes of the `wrapped` function. Return the updated function.""" + if inspect.iscoroutinefunction(wrapped): + + @functools.wraps(wrapped) + async def wrapper_function(*args, **kwargs): # type: ignore + return await wrapper(*args, **kwargs) + else: + + @functools.wraps(wrapped) + def wrapper_function(*args, **kwargs): + return wrapper(*args, **kwargs) + + # We need to manually update this because `partial` object has no `__name__` and `__qualname__`. + wrapper_function.__name__ = extract_function_name(wrapped) + wrapper_function.__qualname__ = extract_function_qualname(wrapped) + wrapper_function.raw_function = wrapped # type: ignore + + return wrapper_function + + +class ValidateCallWrapper: + """This is a wrapper around a function that validates the arguments passed to it, and optionally the return value.""" + + __slots__ = ( + 'function', + 'validate_return', + 'schema_type', + 'module', + 'qualname', + 'ns_resolver', + 'config_wrapper', + '__pydantic_complete__', + '__pydantic_validator__', + '__return_pydantic_validator__', + ) + + def __init__( + self, + function: ValidateCallSupportedTypes, + config: ConfigDict | None, + validate_return: bool, + parent_namespace: MappingNamespace | None, + ) -> None: + self.function = function + self.validate_return = validate_return + if isinstance(function, partial): + self.schema_type = function.func + self.module = function.func.__module__ + else: + self.schema_type = function + self.module = function.__module__ + self.qualname = extract_function_qualname(function) + + self.ns_resolver = NsResolver( + namespaces_tuple=ns_for_function(self.schema_type, parent_namespace=parent_namespace) + ) + self.config_wrapper = ConfigWrapper(config) + if not self.config_wrapper.defer_build: + self._create_validators() + else: + self.__pydantic_complete__ = False + + def _create_validators(self) -> None: + gen_schema = GenerateSchema(self.config_wrapper, self.ns_resolver) + schema = gen_schema.clean_schema(gen_schema.generate_schema(self.function)) + core_config = self.config_wrapper.core_config(title=self.qualname) + + self.__pydantic_validator__ = create_schema_validator( + schema, + self.schema_type, + self.module, + self.qualname, + 'validate_call', + core_config, + self.config_wrapper.plugin_settings, + ) + if self.validate_return: + signature = inspect.signature(self.function) + return_type = signature.return_annotation if signature.return_annotation is not signature.empty else Any + gen_schema = GenerateSchema(self.config_wrapper, self.ns_resolver) + schema = gen_schema.clean_schema(gen_schema.generate_schema(return_type)) + validator = create_schema_validator( + schema, + self.schema_type, + self.module, + self.qualname, + 'validate_call', + core_config, + self.config_wrapper.plugin_settings, + ) + if inspect.iscoroutinefunction(self.function): + + async def return_val_wrapper(aw: Awaitable[Any]) -> None: + return validator.validate_python(await aw) + + self.__return_pydantic_validator__ = return_val_wrapper + else: + self.__return_pydantic_validator__ = validator.validate_python + else: + self.__return_pydantic_validator__ = None + + self.__pydantic_complete__ = True + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + if not self.__pydantic_complete__: + self._create_validators() + + res = self.__pydantic_validator__.validate_python(pydantic_core.ArgsKwargs(args, kwargs)) + if self.__return_pydantic_validator__: + return self.__return_pydantic_validator__(res) + else: + return res diff --git a/lib/python3.12/site-packages/pydantic/_internal/_validators.py b/lib/python3.12/site-packages/pydantic/_internal/_validators.py new file mode 100644 index 0000000000000000000000000000000000000000..803363ca2f34eb9b57702d3f7ed02258fd0c1610 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_internal/_validators.py @@ -0,0 +1,532 @@ +"""Validator functions for standard library types. + +Import of this module is deferred since it contains imports of many standard library modules. +""" + +from __future__ import annotations as _annotations + +import collections.abc +import math +import re +import typing +from decimal import Decimal +from fractions import Fraction +from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network +from typing import Any, Callable, Union, cast, get_origin +from zoneinfo import ZoneInfo, ZoneInfoNotFoundError + +import typing_extensions +from pydantic_core import PydanticCustomError, core_schema +from pydantic_core._pydantic_core import PydanticKnownError +from typing_inspection import typing_objects + +from pydantic._internal._import_utils import import_cached_field_info +from pydantic.errors import PydanticSchemaGenerationError + + +def sequence_validator( + input_value: typing.Sequence[Any], + /, + validator: core_schema.ValidatorFunctionWrapHandler, +) -> typing.Sequence[Any]: + """Validator for `Sequence` types, isinstance(v, Sequence) has already been called.""" + value_type = type(input_value) + + # We don't accept any plain string as a sequence + # Relevant issue: https://github.com/pydantic/pydantic/issues/5595 + if issubclass(value_type, (str, bytes)): + raise PydanticCustomError( + 'sequence_str', + "'{type_name}' instances are not allowed as a Sequence value", + {'type_name': value_type.__name__}, + ) + + # TODO: refactor sequence validation to validate with either a list or a tuple + # schema, depending on the type of the value. + # Additionally, we should be able to remove one of either this validator or the + # SequenceValidator in _std_types_schema.py (preferably this one, while porting over some logic). + # Effectively, a refactor for sequence validation is needed. + if value_type is tuple: + input_value = list(input_value) + + v_list = validator(input_value) + + # the rest of the logic is just re-creating the original type from `v_list` + if value_type is list: + return v_list + elif issubclass(value_type, range): + # return the list as we probably can't re-create the range + return v_list + elif value_type is tuple: + return tuple(v_list) + else: + # best guess at how to re-create the original type, more custom construction logic might be required + return value_type(v_list) # type: ignore[call-arg] + + +def import_string(value: Any) -> Any: + if isinstance(value, str): + try: + return _import_string_logic(value) + except ImportError as e: + raise PydanticCustomError('import_error', 'Invalid python path: {error}', {'error': str(e)}) from e + else: + # otherwise we just return the value and let the next validator do the rest of the work + return value + + +def _import_string_logic(dotted_path: str) -> Any: + """Inspired by uvicorn — dotted paths should include a colon before the final item if that item is not a module. + (This is necessary to distinguish between a submodule and an attribute when there is a conflict.). + + If the dotted path does not include a colon and the final item is not a valid module, importing as an attribute + rather than a submodule will be attempted automatically. + + So, for example, the following values of `dotted_path` result in the following returned values: + * 'collections': + * 'collections.abc': + * 'collections.abc:Mapping': + * `collections.abc.Mapping`: (though this is a bit slower than the previous line) + + An error will be raised under any of the following scenarios: + * `dotted_path` contains more than one colon (e.g., 'collections:abc:Mapping') + * the substring of `dotted_path` before the colon is not a valid module in the environment (e.g., '123:Mapping') + * the substring of `dotted_path` after the colon is not an attribute of the module (e.g., 'collections:abc123') + """ + from importlib import import_module + + components = dotted_path.strip().split(':') + if len(components) > 2: + raise ImportError(f"Import strings should have at most one ':'; received {dotted_path!r}") + + module_path = components[0] + if not module_path: + raise ImportError(f'Import strings should have a nonempty module name; received {dotted_path!r}') + + try: + module = import_module(module_path) + except ModuleNotFoundError as e: + if '.' in module_path: + # Check if it would be valid if the final item was separated from its module with a `:` + maybe_module_path, maybe_attribute = dotted_path.strip().rsplit('.', 1) + try: + return _import_string_logic(f'{maybe_module_path}:{maybe_attribute}') + except ImportError: + pass + raise ImportError(f'No module named {module_path!r}') from e + raise e + + if len(components) > 1: + attribute = components[1] + try: + return getattr(module, attribute) + except AttributeError as e: + raise ImportError(f'cannot import name {attribute!r} from {module_path!r}') from e + else: + return module + + +def pattern_either_validator(input_value: Any, /) -> typing.Pattern[Any]: + if isinstance(input_value, typing.Pattern): + return input_value + elif isinstance(input_value, (str, bytes)): + # todo strict mode + return compile_pattern(input_value) # type: ignore + else: + raise PydanticCustomError('pattern_type', 'Input should be a valid pattern') + + +def pattern_str_validator(input_value: Any, /) -> typing.Pattern[str]: + if isinstance(input_value, typing.Pattern): + if isinstance(input_value.pattern, str): + return input_value + else: + raise PydanticCustomError('pattern_str_type', 'Input should be a string pattern') + elif isinstance(input_value, str): + return compile_pattern(input_value) + elif isinstance(input_value, bytes): + raise PydanticCustomError('pattern_str_type', 'Input should be a string pattern') + else: + raise PydanticCustomError('pattern_type', 'Input should be a valid pattern') + + +def pattern_bytes_validator(input_value: Any, /) -> typing.Pattern[bytes]: + if isinstance(input_value, typing.Pattern): + if isinstance(input_value.pattern, bytes): + return input_value + else: + raise PydanticCustomError('pattern_bytes_type', 'Input should be a bytes pattern') + elif isinstance(input_value, bytes): + return compile_pattern(input_value) + elif isinstance(input_value, str): + raise PydanticCustomError('pattern_bytes_type', 'Input should be a bytes pattern') + else: + raise PydanticCustomError('pattern_type', 'Input should be a valid pattern') + + +PatternType = typing.TypeVar('PatternType', str, bytes) + + +def compile_pattern(pattern: PatternType) -> typing.Pattern[PatternType]: + try: + return re.compile(pattern) + except re.error: + raise PydanticCustomError('pattern_regex', 'Input should be a valid regular expression') + + +def ip_v4_address_validator(input_value: Any, /) -> IPv4Address: + if isinstance(input_value, IPv4Address): + return input_value + + try: + return IPv4Address(input_value) + except ValueError: + raise PydanticCustomError('ip_v4_address', 'Input is not a valid IPv4 address') + + +def ip_v6_address_validator(input_value: Any, /) -> IPv6Address: + if isinstance(input_value, IPv6Address): + return input_value + + try: + return IPv6Address(input_value) + except ValueError: + raise PydanticCustomError('ip_v6_address', 'Input is not a valid IPv6 address') + + +def ip_v4_network_validator(input_value: Any, /) -> IPv4Network: + """Assume IPv4Network initialised with a default `strict` argument. + + See more: + https://docs.python.org/library/ipaddress.html#ipaddress.IPv4Network + """ + if isinstance(input_value, IPv4Network): + return input_value + + try: + return IPv4Network(input_value) + except ValueError: + raise PydanticCustomError('ip_v4_network', 'Input is not a valid IPv4 network') + + +def ip_v6_network_validator(input_value: Any, /) -> IPv6Network: + """Assume IPv6Network initialised with a default `strict` argument. + + See more: + https://docs.python.org/library/ipaddress.html#ipaddress.IPv6Network + """ + if isinstance(input_value, IPv6Network): + return input_value + + try: + return IPv6Network(input_value) + except ValueError: + raise PydanticCustomError('ip_v6_network', 'Input is not a valid IPv6 network') + + +def ip_v4_interface_validator(input_value: Any, /) -> IPv4Interface: + if isinstance(input_value, IPv4Interface): + return input_value + + try: + return IPv4Interface(input_value) + except ValueError: + raise PydanticCustomError('ip_v4_interface', 'Input is not a valid IPv4 interface') + + +def ip_v6_interface_validator(input_value: Any, /) -> IPv6Interface: + if isinstance(input_value, IPv6Interface): + return input_value + + try: + return IPv6Interface(input_value) + except ValueError: + raise PydanticCustomError('ip_v6_interface', 'Input is not a valid IPv6 interface') + + +def fraction_validator(input_value: Any, /) -> Fraction: + if isinstance(input_value, Fraction): + return input_value + + try: + return Fraction(input_value) + except ValueError: + raise PydanticCustomError('fraction_parsing', 'Input is not a valid fraction') + + +def forbid_inf_nan_check(x: Any) -> Any: + if not math.isfinite(x): + raise PydanticKnownError('finite_number') + return x + + +def _safe_repr(v: Any) -> int | float | str: + """The context argument for `PydanticKnownError` requires a number or str type, so we do a simple repr() coercion for types like timedelta. + + See tests/test_types.py::test_annotated_metadata_any_order for some context. + """ + if isinstance(v, (int, float, str)): + return v + return repr(v) + + +def greater_than_validator(x: Any, gt: Any) -> Any: + try: + if not (x > gt): + raise PydanticKnownError('greater_than', {'gt': _safe_repr(gt)}) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'gt' to supplied value {x}") + + +def greater_than_or_equal_validator(x: Any, ge: Any) -> Any: + try: + if not (x >= ge): + raise PydanticKnownError('greater_than_equal', {'ge': _safe_repr(ge)}) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'ge' to supplied value {x}") + + +def less_than_validator(x: Any, lt: Any) -> Any: + try: + if not (x < lt): + raise PydanticKnownError('less_than', {'lt': _safe_repr(lt)}) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'lt' to supplied value {x}") + + +def less_than_or_equal_validator(x: Any, le: Any) -> Any: + try: + if not (x <= le): + raise PydanticKnownError('less_than_equal', {'le': _safe_repr(le)}) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'le' to supplied value {x}") + + +def multiple_of_validator(x: Any, multiple_of: Any) -> Any: + try: + if x % multiple_of: + raise PydanticKnownError('multiple_of', {'multiple_of': _safe_repr(multiple_of)}) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'multiple_of' to supplied value {x}") + + +def min_length_validator(x: Any, min_length: Any) -> Any: + try: + if not (len(x) >= min_length): + raise PydanticKnownError( + 'too_short', {'field_type': 'Value', 'min_length': min_length, 'actual_length': len(x)} + ) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'min_length' to supplied value {x}") + + +def max_length_validator(x: Any, max_length: Any) -> Any: + try: + if len(x) > max_length: + raise PydanticKnownError( + 'too_long', + {'field_type': 'Value', 'max_length': max_length, 'actual_length': len(x)}, + ) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'max_length' to supplied value {x}") + + +def _extract_decimal_digits_info(decimal: Decimal) -> tuple[int, int]: + """Compute the total number of digits and decimal places for a given [`Decimal`][decimal.Decimal] instance. + + This function handles both normalized and non-normalized Decimal instances. + Example: Decimal('1.230') -> 4 digits, 3 decimal places + + Args: + decimal (Decimal): The decimal number to analyze. + + Returns: + tuple[int, int]: A tuple containing the number of decimal places and total digits. + + Though this could be divided into two separate functions, the logic is easier to follow if we couple the computation + of the number of decimals and digits together. + """ + try: + decimal_tuple = decimal.as_tuple() + + assert isinstance(decimal_tuple.exponent, int) + + exponent = decimal_tuple.exponent + num_digits = len(decimal_tuple.digits) + + if exponent >= 0: + # A positive exponent adds that many trailing zeros + # Ex: digit_tuple=(1, 2, 3), exponent=2 -> 12300 -> 0 decimal places, 5 digits + num_digits += exponent + decimal_places = 0 + else: + # If the absolute value of the negative exponent is larger than the + # number of digits, then it's the same as the number of digits, + # because it'll consume all the digits in digit_tuple and then + # add abs(exponent) - len(digit_tuple) leading zeros after the decimal point. + # Ex: digit_tuple=(1, 2, 3), exponent=-2 -> 1.23 -> 2 decimal places, 3 digits + # Ex: digit_tuple=(1, 2, 3), exponent=-4 -> 0.0123 -> 4 decimal places, 4 digits + decimal_places = abs(exponent) + num_digits = max(num_digits, decimal_places) + + return decimal_places, num_digits + except (AssertionError, AttributeError): + raise TypeError(f'Unable to extract decimal digits info from supplied value {decimal}') + + +def max_digits_validator(x: Any, max_digits: Any) -> Any: + try: + _, num_digits = _extract_decimal_digits_info(x) + _, normalized_num_digits = _extract_decimal_digits_info(x.normalize()) + if (num_digits > max_digits) and (normalized_num_digits > max_digits): + raise PydanticKnownError( + 'decimal_max_digits', + {'max_digits': max_digits}, + ) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'max_digits' to supplied value {x}") + + +def decimal_places_validator(x: Any, decimal_places: Any) -> Any: + try: + decimal_places_, _ = _extract_decimal_digits_info(x) + if decimal_places_ > decimal_places: + normalized_decimal_places, _ = _extract_decimal_digits_info(x.normalize()) + if normalized_decimal_places > decimal_places: + raise PydanticKnownError( + 'decimal_max_places', + {'decimal_places': decimal_places}, + ) + return x + except TypeError: + raise TypeError(f"Unable to apply constraint 'decimal_places' to supplied value {x}") + + +def deque_validator(input_value: Any, handler: core_schema.ValidatorFunctionWrapHandler) -> collections.deque[Any]: + return collections.deque(handler(input_value), maxlen=getattr(input_value, 'maxlen', None)) + + +def defaultdict_validator( + input_value: Any, handler: core_schema.ValidatorFunctionWrapHandler, default_default_factory: Callable[[], Any] +) -> collections.defaultdict[Any, Any]: + if isinstance(input_value, collections.defaultdict): + default_factory = input_value.default_factory + return collections.defaultdict(default_factory, handler(input_value)) + else: + return collections.defaultdict(default_default_factory, handler(input_value)) + + +def get_defaultdict_default_default_factory(values_source_type: Any) -> Callable[[], Any]: + FieldInfo = import_cached_field_info() + + values_type_origin = get_origin(values_source_type) + + def infer_default() -> Callable[[], Any]: + allowed_default_types: dict[Any, Any] = { + tuple: tuple, + collections.abc.Sequence: tuple, + collections.abc.MutableSequence: list, + list: list, + typing.Sequence: list, + set: set, + typing.MutableSet: set, + collections.abc.MutableSet: set, + collections.abc.Set: frozenset, + typing.MutableMapping: dict, + typing.Mapping: dict, + collections.abc.Mapping: dict, + collections.abc.MutableMapping: dict, + float: float, + int: int, + str: str, + bool: bool, + } + values_type = values_type_origin or values_source_type + instructions = 'set using `DefaultDict[..., Annotated[..., Field(default_factory=...)]]`' + if typing_objects.is_typevar(values_type): + + def type_var_default_factory() -> None: + raise RuntimeError( + 'Generic defaultdict cannot be used without a concrete value type or an' + ' explicit default factory, ' + instructions + ) + + return type_var_default_factory + elif values_type not in allowed_default_types: + # a somewhat subjective set of types that have reasonable default values + allowed_msg = ', '.join([t.__name__ for t in set(allowed_default_types.values())]) + raise PydanticSchemaGenerationError( + f'Unable to infer a default factory for keys of type {values_source_type}.' + f' Only {allowed_msg} are supported, other types require an explicit default factory' + ' ' + instructions + ) + return allowed_default_types[values_type] + + # Assume Annotated[..., Field(...)] + if typing_objects.is_annotated(values_type_origin): + field_info = next((v for v in typing_extensions.get_args(values_source_type) if isinstance(v, FieldInfo)), None) + else: + field_info = None + if field_info and field_info.default_factory: + # Assume the default factory does not take any argument: + default_default_factory = cast(Callable[[], Any], field_info.default_factory) + else: + default_default_factory = infer_default() + return default_default_factory + + +def validate_str_is_valid_iana_tz(value: Any, /) -> ZoneInfo: + if isinstance(value, ZoneInfo): + return value + try: + return ZoneInfo(value) + except (ZoneInfoNotFoundError, ValueError, TypeError): + raise PydanticCustomError('zoneinfo_str', 'invalid timezone: {value}', {'value': value}) + + +NUMERIC_VALIDATOR_LOOKUP: dict[str, Callable] = { + 'gt': greater_than_validator, + 'ge': greater_than_or_equal_validator, + 'lt': less_than_validator, + 'le': less_than_or_equal_validator, + 'multiple_of': multiple_of_validator, + 'min_length': min_length_validator, + 'max_length': max_length_validator, + 'max_digits': max_digits_validator, + 'decimal_places': decimal_places_validator, +} + +IpType = Union[IPv4Address, IPv6Address, IPv4Network, IPv6Network, IPv4Interface, IPv6Interface] + +IP_VALIDATOR_LOOKUP: dict[type[IpType], Callable] = { + IPv4Address: ip_v4_address_validator, + IPv6Address: ip_v6_address_validator, + IPv4Network: ip_v4_network_validator, + IPv6Network: ip_v6_network_validator, + IPv4Interface: ip_v4_interface_validator, + IPv6Interface: ip_v6_interface_validator, +} + +MAPPING_ORIGIN_MAP: dict[Any, Any] = { + typing.DefaultDict: collections.defaultdict, # noqa: UP006 + collections.defaultdict: collections.defaultdict, + typing.OrderedDict: collections.OrderedDict, # noqa: UP006 + collections.OrderedDict: collections.OrderedDict, + typing_extensions.OrderedDict: collections.OrderedDict, + typing.Counter: collections.Counter, + collections.Counter: collections.Counter, + # this doesn't handle subclasses of these + typing.Mapping: dict, + typing.MutableMapping: dict, + # parametrized typing.{Mutable}Mapping creates one of these + collections.abc.Mapping: dict, + collections.abc.MutableMapping: dict, +} diff --git a/lib/python3.12/site-packages/pydantic/_migration.py b/lib/python3.12/site-packages/pydantic/_migration.py new file mode 100644 index 0000000000000000000000000000000000000000..980dfd23eedec7aaa03f5d98e4ee9ae82ad80baf --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/_migration.py @@ -0,0 +1,308 @@ +import sys +from typing import Any, Callable + +from .version import version_short + +MOVED_IN_V2 = { + 'pydantic.utils:version_info': 'pydantic.version:version_info', + 'pydantic.error_wrappers:ValidationError': 'pydantic:ValidationError', + 'pydantic.utils:to_camel': 'pydantic.alias_generators:to_pascal', + 'pydantic.utils:to_lower_camel': 'pydantic.alias_generators:to_camel', + 'pydantic:PyObject': 'pydantic.types:ImportString', + 'pydantic.types:PyObject': 'pydantic.types:ImportString', + 'pydantic.generics:GenericModel': 'pydantic.BaseModel', +} + +DEPRECATED_MOVED_IN_V2 = { + 'pydantic.tools:schema_of': 'pydantic.deprecated.tools:schema_of', + 'pydantic.tools:parse_obj_as': 'pydantic.deprecated.tools:parse_obj_as', + 'pydantic.tools:schema_json_of': 'pydantic.deprecated.tools:schema_json_of', + 'pydantic.json:pydantic_encoder': 'pydantic.deprecated.json:pydantic_encoder', + 'pydantic:validate_arguments': 'pydantic.deprecated.decorator:validate_arguments', + 'pydantic.json:custom_pydantic_encoder': 'pydantic.deprecated.json:custom_pydantic_encoder', + 'pydantic.json:timedelta_isoformat': 'pydantic.deprecated.json:timedelta_isoformat', + 'pydantic.decorator:validate_arguments': 'pydantic.deprecated.decorator:validate_arguments', + 'pydantic.class_validators:validator': 'pydantic.deprecated.class_validators:validator', + 'pydantic.class_validators:root_validator': 'pydantic.deprecated.class_validators:root_validator', + 'pydantic.config:BaseConfig': 'pydantic.deprecated.config:BaseConfig', + 'pydantic.config:Extra': 'pydantic.deprecated.config:Extra', +} + +REDIRECT_TO_V1 = { + f'pydantic.utils:{obj}': f'pydantic.v1.utils:{obj}' + for obj in ( + 'deep_update', + 'GetterDict', + 'lenient_issubclass', + 'lenient_isinstance', + 'is_valid_field', + 'update_not_none', + 'import_string', + 'Representation', + 'ROOT_KEY', + 'smart_deepcopy', + 'sequence_like', + ) +} + + +REMOVED_IN_V2 = { + 'pydantic:ConstrainedBytes', + 'pydantic:ConstrainedDate', + 'pydantic:ConstrainedDecimal', + 'pydantic:ConstrainedFloat', + 'pydantic:ConstrainedFrozenSet', + 'pydantic:ConstrainedInt', + 'pydantic:ConstrainedList', + 'pydantic:ConstrainedSet', + 'pydantic:ConstrainedStr', + 'pydantic:JsonWrapper', + 'pydantic:NoneBytes', + 'pydantic:NoneStr', + 'pydantic:NoneStrBytes', + 'pydantic:Protocol', + 'pydantic:Required', + 'pydantic:StrBytes', + 'pydantic:compiled', + 'pydantic.config:get_config', + 'pydantic.config:inherit_config', + 'pydantic.config:prepare_config', + 'pydantic:create_model_from_namedtuple', + 'pydantic:create_model_from_typeddict', + 'pydantic.dataclasses:create_pydantic_model_from_dataclass', + 'pydantic.dataclasses:make_dataclass_validator', + 'pydantic.dataclasses:set_validation', + 'pydantic.datetime_parse:parse_date', + 'pydantic.datetime_parse:parse_time', + 'pydantic.datetime_parse:parse_datetime', + 'pydantic.datetime_parse:parse_duration', + 'pydantic.error_wrappers:ErrorWrapper', + 'pydantic.errors:AnyStrMaxLengthError', + 'pydantic.errors:AnyStrMinLengthError', + 'pydantic.errors:ArbitraryTypeError', + 'pydantic.errors:BoolError', + 'pydantic.errors:BytesError', + 'pydantic.errors:CallableError', + 'pydantic.errors:ClassError', + 'pydantic.errors:ColorError', + 'pydantic.errors:ConfigError', + 'pydantic.errors:DataclassTypeError', + 'pydantic.errors:DateError', + 'pydantic.errors:DateNotInTheFutureError', + 'pydantic.errors:DateNotInThePastError', + 'pydantic.errors:DateTimeError', + 'pydantic.errors:DecimalError', + 'pydantic.errors:DecimalIsNotFiniteError', + 'pydantic.errors:DecimalMaxDigitsError', + 'pydantic.errors:DecimalMaxPlacesError', + 'pydantic.errors:DecimalWholeDigitsError', + 'pydantic.errors:DictError', + 'pydantic.errors:DurationError', + 'pydantic.errors:EmailError', + 'pydantic.errors:EnumError', + 'pydantic.errors:EnumMemberError', + 'pydantic.errors:ExtraError', + 'pydantic.errors:FloatError', + 'pydantic.errors:FrozenSetError', + 'pydantic.errors:FrozenSetMaxLengthError', + 'pydantic.errors:FrozenSetMinLengthError', + 'pydantic.errors:HashableError', + 'pydantic.errors:IPv4AddressError', + 'pydantic.errors:IPv4InterfaceError', + 'pydantic.errors:IPv4NetworkError', + 'pydantic.errors:IPv6AddressError', + 'pydantic.errors:IPv6InterfaceError', + 'pydantic.errors:IPv6NetworkError', + 'pydantic.errors:IPvAnyAddressError', + 'pydantic.errors:IPvAnyInterfaceError', + 'pydantic.errors:IPvAnyNetworkError', + 'pydantic.errors:IntEnumError', + 'pydantic.errors:IntegerError', + 'pydantic.errors:InvalidByteSize', + 'pydantic.errors:InvalidByteSizeUnit', + 'pydantic.errors:InvalidDiscriminator', + 'pydantic.errors:InvalidLengthForBrand', + 'pydantic.errors:JsonError', + 'pydantic.errors:JsonTypeError', + 'pydantic.errors:ListError', + 'pydantic.errors:ListMaxLengthError', + 'pydantic.errors:ListMinLengthError', + 'pydantic.errors:ListUniqueItemsError', + 'pydantic.errors:LuhnValidationError', + 'pydantic.errors:MissingDiscriminator', + 'pydantic.errors:MissingError', + 'pydantic.errors:NoneIsAllowedError', + 'pydantic.errors:NoneIsNotAllowedError', + 'pydantic.errors:NotDigitError', + 'pydantic.errors:NotNoneError', + 'pydantic.errors:NumberNotGeError', + 'pydantic.errors:NumberNotGtError', + 'pydantic.errors:NumberNotLeError', + 'pydantic.errors:NumberNotLtError', + 'pydantic.errors:NumberNotMultipleError', + 'pydantic.errors:PathError', + 'pydantic.errors:PathNotADirectoryError', + 'pydantic.errors:PathNotAFileError', + 'pydantic.errors:PathNotExistsError', + 'pydantic.errors:PatternError', + 'pydantic.errors:PyObjectError', + 'pydantic.errors:PydanticTypeError', + 'pydantic.errors:PydanticValueError', + 'pydantic.errors:SequenceError', + 'pydantic.errors:SetError', + 'pydantic.errors:SetMaxLengthError', + 'pydantic.errors:SetMinLengthError', + 'pydantic.errors:StrError', + 'pydantic.errors:StrRegexError', + 'pydantic.errors:StrictBoolError', + 'pydantic.errors:SubclassError', + 'pydantic.errors:TimeError', + 'pydantic.errors:TupleError', + 'pydantic.errors:TupleLengthError', + 'pydantic.errors:UUIDError', + 'pydantic.errors:UUIDVersionError', + 'pydantic.errors:UrlError', + 'pydantic.errors:UrlExtraError', + 'pydantic.errors:UrlHostError', + 'pydantic.errors:UrlHostTldError', + 'pydantic.errors:UrlPortError', + 'pydantic.errors:UrlSchemeError', + 'pydantic.errors:UrlSchemePermittedError', + 'pydantic.errors:UrlUserInfoError', + 'pydantic.errors:WrongConstantError', + 'pydantic.main:validate_model', + 'pydantic.networks:stricturl', + 'pydantic:parse_file_as', + 'pydantic:parse_raw_as', + 'pydantic:stricturl', + 'pydantic.tools:parse_file_as', + 'pydantic.tools:parse_raw_as', + 'pydantic.types:ConstrainedBytes', + 'pydantic.types:ConstrainedDate', + 'pydantic.types:ConstrainedDecimal', + 'pydantic.types:ConstrainedFloat', + 'pydantic.types:ConstrainedFrozenSet', + 'pydantic.types:ConstrainedInt', + 'pydantic.types:ConstrainedList', + 'pydantic.types:ConstrainedSet', + 'pydantic.types:ConstrainedStr', + 'pydantic.types:JsonWrapper', + 'pydantic.types:NoneBytes', + 'pydantic.types:NoneStr', + 'pydantic.types:NoneStrBytes', + 'pydantic.types:StrBytes', + 'pydantic.typing:evaluate_forwardref', + 'pydantic.typing:AbstractSetIntStr', + 'pydantic.typing:AnyCallable', + 'pydantic.typing:AnyClassMethod', + 'pydantic.typing:CallableGenerator', + 'pydantic.typing:DictAny', + 'pydantic.typing:DictIntStrAny', + 'pydantic.typing:DictStrAny', + 'pydantic.typing:IntStr', + 'pydantic.typing:ListStr', + 'pydantic.typing:MappingIntStrAny', + 'pydantic.typing:NoArgAnyCallable', + 'pydantic.typing:NoneType', + 'pydantic.typing:ReprArgs', + 'pydantic.typing:SetStr', + 'pydantic.typing:StrPath', + 'pydantic.typing:TupleGenerator', + 'pydantic.typing:WithArgsTypes', + 'pydantic.typing:all_literal_values', + 'pydantic.typing:display_as_type', + 'pydantic.typing:get_all_type_hints', + 'pydantic.typing:get_args', + 'pydantic.typing:get_origin', + 'pydantic.typing:get_sub_types', + 'pydantic.typing:is_callable_type', + 'pydantic.typing:is_classvar', + 'pydantic.typing:is_finalvar', + 'pydantic.typing:is_literal_type', + 'pydantic.typing:is_namedtuple', + 'pydantic.typing:is_new_type', + 'pydantic.typing:is_none_type', + 'pydantic.typing:is_typeddict', + 'pydantic.typing:is_typeddict_special', + 'pydantic.typing:is_union', + 'pydantic.typing:new_type_supertype', + 'pydantic.typing:resolve_annotations', + 'pydantic.typing:typing_base', + 'pydantic.typing:update_field_forward_refs', + 'pydantic.typing:update_model_forward_refs', + 'pydantic.utils:ClassAttribute', + 'pydantic.utils:DUNDER_ATTRIBUTES', + 'pydantic.utils:PyObjectStr', + 'pydantic.utils:ValueItems', + 'pydantic.utils:almost_equal_floats', + 'pydantic.utils:get_discriminator_alias_and_values', + 'pydantic.utils:get_model', + 'pydantic.utils:get_unique_discriminator_alias', + 'pydantic.utils:in_ipython', + 'pydantic.utils:is_valid_identifier', + 'pydantic.utils:path_type', + 'pydantic.utils:validate_field_name', + 'pydantic:validate_model', +} + + +def getattr_migration(module: str) -> Callable[[str], Any]: + """Implement PEP 562 for objects that were either moved or removed on the migration + to V2. + + Args: + module: The module name. + + Returns: + A callable that will raise an error if the object is not found. + """ + # This avoids circular import with errors.py. + from .errors import PydanticImportError + + def wrapper(name: str) -> object: + """Raise an error if the object is not found, or warn if it was moved. + + In case it was moved, it still returns the object. + + Args: + name: The object name. + + Returns: + The object. + """ + if name == '__path__': + raise AttributeError(f'module {module!r} has no attribute {name!r}') + + import warnings + + from ._internal._validators import import_string + + import_path = f'{module}:{name}' + if import_path in MOVED_IN_V2.keys(): + new_location = MOVED_IN_V2[import_path] + warnings.warn(f'`{import_path}` has been moved to `{new_location}`.') + return import_string(MOVED_IN_V2[import_path]) + if import_path in DEPRECATED_MOVED_IN_V2: + # skip the warning here because a deprecation warning will be raised elsewhere + return import_string(DEPRECATED_MOVED_IN_V2[import_path]) + if import_path in REDIRECT_TO_V1: + new_location = REDIRECT_TO_V1[import_path] + warnings.warn( + f'`{import_path}` has been removed. We are importing from `{new_location}` instead.' + 'See the migration guide for more details: https://docs.pydantic.dev/latest/migration/' + ) + return import_string(REDIRECT_TO_V1[import_path]) + if import_path == 'pydantic:BaseSettings': + raise PydanticImportError( + '`BaseSettings` has been moved to the `pydantic-settings` package. ' + f'See https://docs.pydantic.dev/{version_short()}/migration/#basesettings-has-moved-to-pydantic-settings ' + 'for more details.' + ) + if import_path in REMOVED_IN_V2: + raise PydanticImportError(f'`{import_path}` has been removed in V2.') + globals: dict[str, Any] = sys.modules[module].__dict__ + if name in globals: + return globals[name] + raise AttributeError(f'module {module!r} has no attribute {name!r}') + + return wrapper diff --git a/lib/python3.12/site-packages/pydantic/alias_generators.py b/lib/python3.12/site-packages/pydantic/alias_generators.py new file mode 100644 index 0000000000000000000000000000000000000000..0b7653f5825979edcd446d698d320bdb980b8e9e --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/alias_generators.py @@ -0,0 +1,62 @@ +"""Alias generators for converting between different capitalization conventions.""" + +import re + +__all__ = ('to_pascal', 'to_camel', 'to_snake') + +# TODO: in V3, change the argument names to be more descriptive +# Generally, don't only convert from snake_case, or name the functions +# more specifically like snake_to_camel. + + +def to_pascal(snake: str) -> str: + """Convert a snake_case string to PascalCase. + + Args: + snake: The string to convert. + + Returns: + The PascalCase string. + """ + camel = snake.title() + return re.sub('([0-9A-Za-z])_(?=[0-9A-Z])', lambda m: m.group(1), camel) + + +def to_camel(snake: str) -> str: + """Convert a snake_case string to camelCase. + + Args: + snake: The string to convert. + + Returns: + The converted camelCase string. + """ + # If the string is already in camelCase and does not contain a digit followed + # by a lowercase letter, return it as it is + if re.match('^[a-z]+[A-Za-z0-9]*$', snake) and not re.search(r'\d[a-z]', snake): + return snake + + camel = to_pascal(snake) + return re.sub('(^_*[A-Z])', lambda m: m.group(1).lower(), camel) + + +def to_snake(camel: str) -> str: + """Convert a PascalCase, camelCase, or kebab-case string to snake_case. + + Args: + camel: The string to convert. + + Returns: + The converted string in snake_case. + """ + # Handle the sequence of uppercase letters followed by a lowercase letter + snake = re.sub(r'([A-Z]+)([A-Z][a-z])', lambda m: f'{m.group(1)}_{m.group(2)}', camel) + # Insert an underscore between a lowercase letter and an uppercase letter + snake = re.sub(r'([a-z])([A-Z])', lambda m: f'{m.group(1)}_{m.group(2)}', snake) + # Insert an underscore between a digit and an uppercase letter + snake = re.sub(r'([0-9])([A-Z])', lambda m: f'{m.group(1)}_{m.group(2)}', snake) + # Insert an underscore between a lowercase letter and a digit + snake = re.sub(r'([a-z])([0-9])', lambda m: f'{m.group(1)}_{m.group(2)}', snake) + # Replace hyphens with underscores to handle kebab-case + snake = snake.replace('-', '_') + return snake.lower() diff --git a/lib/python3.12/site-packages/pydantic/aliases.py b/lib/python3.12/site-packages/pydantic/aliases.py new file mode 100644 index 0000000000000000000000000000000000000000..ac2273701b5a0ae7cb0d0171cd08a7bbf77a9178 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/aliases.py @@ -0,0 +1,135 @@ +"""Support for alias configurations.""" + +from __future__ import annotations + +import dataclasses +from typing import Any, Callable, Literal + +from pydantic_core import PydanticUndefined + +from ._internal import _internal_dataclass + +__all__ = ('AliasGenerator', 'AliasPath', 'AliasChoices') + + +@dataclasses.dataclass(**_internal_dataclass.slots_true) +class AliasPath: + """!!! abstract "Usage Documentation" + [`AliasPath` and `AliasChoices`](../concepts/alias.md#aliaspath-and-aliaschoices) + + A data class used by `validation_alias` as a convenience to create aliases. + + Attributes: + path: A list of string or integer aliases. + """ + + path: list[int | str] + + def __init__(self, first_arg: str, *args: str | int) -> None: + self.path = [first_arg] + list(args) + + def convert_to_aliases(self) -> list[str | int]: + """Converts arguments to a list of string or integer aliases. + + Returns: + The list of aliases. + """ + return self.path + + def search_dict_for_path(self, d: dict) -> Any: + """Searches a dictionary for the path specified by the alias. + + Returns: + The value at the specified path, or `PydanticUndefined` if the path is not found. + """ + v = d + for k in self.path: + if isinstance(v, str): + # disallow indexing into a str, like for AliasPath('x', 0) and x='abc' + return PydanticUndefined + try: + v = v[k] + except (KeyError, IndexError, TypeError): + return PydanticUndefined + return v + + +@dataclasses.dataclass(**_internal_dataclass.slots_true) +class AliasChoices: + """!!! abstract "Usage Documentation" + [`AliasPath` and `AliasChoices`](../concepts/alias.md#aliaspath-and-aliaschoices) + + A data class used by `validation_alias` as a convenience to create aliases. + + Attributes: + choices: A list containing a string or `AliasPath`. + """ + + choices: list[str | AliasPath] + + def __init__(self, first_choice: str | AliasPath, *choices: str | AliasPath) -> None: + self.choices = [first_choice] + list(choices) + + def convert_to_aliases(self) -> list[list[str | int]]: + """Converts arguments to a list of lists containing string or integer aliases. + + Returns: + The list of aliases. + """ + aliases: list[list[str | int]] = [] + for c in self.choices: + if isinstance(c, AliasPath): + aliases.append(c.convert_to_aliases()) + else: + aliases.append([c]) + return aliases + + +@dataclasses.dataclass(**_internal_dataclass.slots_true) +class AliasGenerator: + """!!! abstract "Usage Documentation" + [Using an `AliasGenerator`](../concepts/alias.md#using-an-aliasgenerator) + + A data class used by `alias_generator` as a convenience to create various aliases. + + Attributes: + alias: A callable that takes a field name and returns an alias for it. + validation_alias: A callable that takes a field name and returns a validation alias for it. + serialization_alias: A callable that takes a field name and returns a serialization alias for it. + """ + + alias: Callable[[str], str] | None = None + validation_alias: Callable[[str], str | AliasPath | AliasChoices] | None = None + serialization_alias: Callable[[str], str] | None = None + + def _generate_alias( + self, + alias_kind: Literal['alias', 'validation_alias', 'serialization_alias'], + allowed_types: tuple[type[str] | type[AliasPath] | type[AliasChoices], ...], + field_name: str, + ) -> str | AliasPath | AliasChoices | None: + """Generate an alias of the specified kind. Returns None if the alias generator is None. + + Raises: + TypeError: If the alias generator produces an invalid type. + """ + alias = None + if alias_generator := getattr(self, alias_kind): + alias = alias_generator(field_name) + if alias and not isinstance(alias, allowed_types): + raise TypeError( + f'Invalid `{alias_kind}` type. `{alias_kind}` generator must produce one of `{allowed_types}`' + ) + return alias + + def generate_aliases(self, field_name: str) -> tuple[str | None, str | AliasPath | AliasChoices | None, str | None]: + """Generate `alias`, `validation_alias`, and `serialization_alias` for a field. + + Returns: + A tuple of three aliases - validation, alias, and serialization. + """ + alias = self._generate_alias('alias', (str,), field_name) + validation_alias = self._generate_alias('validation_alias', (str, AliasChoices, AliasPath), field_name) + serialization_alias = self._generate_alias('serialization_alias', (str,), field_name) + + return alias, validation_alias, serialization_alias # type: ignore diff --git a/lib/python3.12/site-packages/pydantic/annotated_handlers.py b/lib/python3.12/site-packages/pydantic/annotated_handlers.py new file mode 100644 index 0000000000000000000000000000000000000000..d0cb5d3d4616f938aea2aae42985b4c166b500cc --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/annotated_handlers.py @@ -0,0 +1,122 @@ +"""Type annotations to use with `__get_pydantic_core_schema__` and `__get_pydantic_json_schema__`.""" + +from __future__ import annotations as _annotations + +from typing import TYPE_CHECKING, Any, Union + +from pydantic_core import core_schema + +if TYPE_CHECKING: + from ._internal._namespace_utils import NamespacesTuple + from .json_schema import JsonSchemaMode, JsonSchemaValue + + CoreSchemaOrField = Union[ + core_schema.CoreSchema, + core_schema.ModelField, + core_schema.DataclassField, + core_schema.TypedDictField, + core_schema.ComputedField, + ] + +__all__ = 'GetJsonSchemaHandler', 'GetCoreSchemaHandler' + + +class GetJsonSchemaHandler: + """Handler to call into the next JSON schema generation function. + + Attributes: + mode: Json schema mode, can be `validation` or `serialization`. + """ + + mode: JsonSchemaMode + + def __call__(self, core_schema: CoreSchemaOrField, /) -> JsonSchemaValue: + """Call the inner handler and get the JsonSchemaValue it returns. + This will call the next JSON schema modifying function up until it calls + into `pydantic.json_schema.GenerateJsonSchema`, which will raise a + `pydantic.errors.PydanticInvalidForJsonSchema` error if it cannot generate + a JSON schema. + + Args: + core_schema: A `pydantic_core.core_schema.CoreSchema`. + + Returns: + JsonSchemaValue: The JSON schema generated by the inner JSON schema modify + functions. + """ + raise NotImplementedError + + def resolve_ref_schema(self, maybe_ref_json_schema: JsonSchemaValue, /) -> JsonSchemaValue: + """Get the real schema for a `{"$ref": ...}` schema. + If the schema given is not a `$ref` schema, it will be returned as is. + This means you don't have to check before calling this function. + + Args: + maybe_ref_json_schema: A JsonSchemaValue which may be a `$ref` schema. + + Raises: + LookupError: If the ref is not found. + + Returns: + JsonSchemaValue: A JsonSchemaValue that has no `$ref`. + """ + raise NotImplementedError + + +class GetCoreSchemaHandler: + """Handler to call into the next CoreSchema schema generation function.""" + + def __call__(self, source_type: Any, /) -> core_schema.CoreSchema: + """Call the inner handler and get the CoreSchema it returns. + This will call the next CoreSchema modifying function up until it calls + into Pydantic's internal schema generation machinery, which will raise a + `pydantic.errors.PydanticSchemaGenerationError` error if it cannot generate + a CoreSchema for the given source type. + + Args: + source_type: The input type. + + Returns: + CoreSchema: The `pydantic-core` CoreSchema generated. + """ + raise NotImplementedError + + def generate_schema(self, source_type: Any, /) -> core_schema.CoreSchema: + """Generate a schema unrelated to the current context. + Use this function if e.g. you are handling schema generation for a sequence + and want to generate a schema for its items. + Otherwise, you may end up doing something like applying a `min_length` constraint + that was intended for the sequence itself to its items! + + Args: + source_type: The input type. + + Returns: + CoreSchema: The `pydantic-core` CoreSchema generated. + """ + raise NotImplementedError + + def resolve_ref_schema(self, maybe_ref_schema: core_schema.CoreSchema, /) -> core_schema.CoreSchema: + """Get the real schema for a `definition-ref` schema. + If the schema given is not a `definition-ref` schema, it will be returned as is. + This means you don't have to check before calling this function. + + Args: + maybe_ref_schema: A `CoreSchema`, `ref`-based or not. + + Raises: + LookupError: If the `ref` is not found. + + Returns: + A concrete `CoreSchema`. + """ + raise NotImplementedError + + @property + def field_name(self) -> str | None: + """Get the name of the closest field to this validator.""" + raise NotImplementedError + + def _get_types_namespace(self) -> NamespacesTuple: + """Internal method used during type resolution for serializer annotations.""" + raise NotImplementedError diff --git a/lib/python3.12/site-packages/pydantic/class_validators.py b/lib/python3.12/site-packages/pydantic/class_validators.py new file mode 100644 index 0000000000000000000000000000000000000000..9977150c92fcb083fcdfa632c9de3b5fa92470cb --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/class_validators.py @@ -0,0 +1,5 @@ +"""`class_validators` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/color.py b/lib/python3.12/site-packages/pydantic/color.py new file mode 100644 index 0000000000000000000000000000000000000000..9a42d586fd71a786ba63a4d5ce55a0ed93d866b0 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/color.py @@ -0,0 +1,604 @@ +"""Color definitions are used as per the CSS3 +[CSS Color Module Level 3](http://www.w3.org/TR/css3-color/#svg-color) specification. + +A few colors have multiple names referring to the sames colors, eg. `grey` and `gray` or `aqua` and `cyan`. + +In these cases the _last_ color when sorted alphabetically takes preferences, +eg. `Color((0, 255, 255)).as_named() == 'cyan'` because "cyan" comes after "aqua". + +Warning: Deprecated + The `Color` class is deprecated, use `pydantic_extra_types` instead. + See [`pydantic-extra-types.Color`](../usage/types/extra_types/color_types.md) + for more information. +""" + +import math +import re +from colorsys import hls_to_rgb, rgb_to_hls +from typing import Any, Callable, Optional, Union, cast + +from pydantic_core import CoreSchema, PydanticCustomError, core_schema +from typing_extensions import deprecated + +from ._internal import _repr +from ._internal._schema_generation_shared import GetJsonSchemaHandler as _GetJsonSchemaHandler +from .json_schema import JsonSchemaValue +from .warnings import PydanticDeprecatedSince20 + +ColorTuple = Union[tuple[int, int, int], tuple[int, int, int, float]] +ColorType = Union[ColorTuple, str] +HslColorTuple = Union[tuple[float, float, float], tuple[float, float, float, float]] + + +class RGBA: + """Internal use only as a representation of a color.""" + + __slots__ = 'r', 'g', 'b', 'alpha', '_tuple' + + def __init__(self, r: float, g: float, b: float, alpha: Optional[float]): + self.r = r + self.g = g + self.b = b + self.alpha = alpha + + self._tuple: tuple[float, float, float, Optional[float]] = (r, g, b, alpha) + + def __getitem__(self, item: Any) -> Any: + return self._tuple[item] + + +# these are not compiled here to avoid import slowdown, they'll be compiled the first time they're used, then cached +_r_255 = r'(\d{1,3}(?:\.\d+)?)' +_r_comma = r'\s*,\s*' +_r_alpha = r'(\d(?:\.\d+)?|\.\d+|\d{1,2}%)' +_r_h = r'(-?\d+(?:\.\d+)?|-?\.\d+)(deg|rad|turn)?' +_r_sl = r'(\d{1,3}(?:\.\d+)?)%' +r_hex_short = r'\s*(?:#|0x)?([0-9a-f])([0-9a-f])([0-9a-f])([0-9a-f])?\s*' +r_hex_long = r'\s*(?:#|0x)?([0-9a-f]{2})([0-9a-f]{2})([0-9a-f]{2})([0-9a-f]{2})?\s*' +# CSS3 RGB examples: rgb(0, 0, 0), rgba(0, 0, 0, 0.5), rgba(0, 0, 0, 50%) +r_rgb = rf'\s*rgba?\(\s*{_r_255}{_r_comma}{_r_255}{_r_comma}{_r_255}(?:{_r_comma}{_r_alpha})?\s*\)\s*' +# CSS3 HSL examples: hsl(270, 60%, 50%), hsla(270, 60%, 50%, 0.5), hsla(270, 60%, 50%, 50%) +r_hsl = rf'\s*hsla?\(\s*{_r_h}{_r_comma}{_r_sl}{_r_comma}{_r_sl}(?:{_r_comma}{_r_alpha})?\s*\)\s*' +# CSS4 RGB examples: rgb(0 0 0), rgb(0 0 0 / 0.5), rgb(0 0 0 / 50%), rgba(0 0 0 / 50%) +r_rgb_v4_style = rf'\s*rgba?\(\s*{_r_255}\s+{_r_255}\s+{_r_255}(?:\s*/\s*{_r_alpha})?\s*\)\s*' +# CSS4 HSL examples: hsl(270 60% 50%), hsl(270 60% 50% / 0.5), hsl(270 60% 50% / 50%), hsla(270 60% 50% / 50%) +r_hsl_v4_style = rf'\s*hsla?\(\s*{_r_h}\s+{_r_sl}\s+{_r_sl}(?:\s*/\s*{_r_alpha})?\s*\)\s*' + +# colors where the two hex characters are the same, if all colors match this the short version of hex colors can be used +repeat_colors = {int(c * 2, 16) for c in '0123456789abcdef'} +rads = 2 * math.pi + + +@deprecated( + 'The `Color` class is deprecated, use `pydantic_extra_types` instead. ' + 'See https://docs.pydantic.dev/latest/api/pydantic_extra_types_color/.', + category=PydanticDeprecatedSince20, +) +class Color(_repr.Representation): + """Represents a color.""" + + __slots__ = '_original', '_rgba' + + def __init__(self, value: ColorType) -> None: + self._rgba: RGBA + self._original: ColorType + if isinstance(value, (tuple, list)): + self._rgba = parse_tuple(value) + elif isinstance(value, str): + self._rgba = parse_str(value) + elif isinstance(value, Color): + self._rgba = value._rgba + value = value._original + else: + raise PydanticCustomError( + 'color_error', 'value is not a valid color: value must be a tuple, list or string' + ) + + # if we've got here value must be a valid color + self._original = value + + @classmethod + def __get_pydantic_json_schema__( + cls, core_schema: core_schema.CoreSchema, handler: _GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = {} + field_schema.update(type='string', format='color') + return field_schema + + def original(self) -> ColorType: + """Original value passed to `Color`.""" + return self._original + + def as_named(self, *, fallback: bool = False) -> str: + """Returns the name of the color if it can be found in `COLORS_BY_VALUE` dictionary, + otherwise returns the hexadecimal representation of the color or raises `ValueError`. + + Args: + fallback: If True, falls back to returning the hexadecimal representation of + the color instead of raising a ValueError when no named color is found. + + Returns: + The name of the color, or the hexadecimal representation of the color. + + Raises: + ValueError: When no named color is found and fallback is `False`. + """ + if self._rgba.alpha is None: + rgb = cast(tuple[int, int, int], self.as_rgb_tuple()) + try: + return COLORS_BY_VALUE[rgb] + except KeyError as e: + if fallback: + return self.as_hex() + else: + raise ValueError('no named color found, use fallback=True, as_hex() or as_rgb()') from e + else: + return self.as_hex() + + def as_hex(self) -> str: + """Returns the hexadecimal representation of the color. + + Hex string representing the color can be 3, 4, 6, or 8 characters depending on whether the string + a "short" representation of the color is possible and whether there's an alpha channel. + + Returns: + The hexadecimal representation of the color. + """ + values = [float_to_255(c) for c in self._rgba[:3]] + if self._rgba.alpha is not None: + values.append(float_to_255(self._rgba.alpha)) + + as_hex = ''.join(f'{v:02x}' for v in values) + if all(c in repeat_colors for c in values): + as_hex = ''.join(as_hex[c] for c in range(0, len(as_hex), 2)) + return '#' + as_hex + + def as_rgb(self) -> str: + """Color as an `rgb(, , )` or `rgba(, , , )` string.""" + if self._rgba.alpha is None: + return f'rgb({float_to_255(self._rgba.r)}, {float_to_255(self._rgba.g)}, {float_to_255(self._rgba.b)})' + else: + return ( + f'rgba({float_to_255(self._rgba.r)}, {float_to_255(self._rgba.g)}, {float_to_255(self._rgba.b)}, ' + f'{round(self._alpha_float(), 2)})' + ) + + def as_rgb_tuple(self, *, alpha: Optional[bool] = None) -> ColorTuple: + """Returns the color as an RGB or RGBA tuple. + + Args: + alpha: Whether to include the alpha channel. There are three options for this input: + + - `None` (default): Include alpha only if it's set. (e.g. not `None`) + - `True`: Always include alpha. + - `False`: Always omit alpha. + + Returns: + A tuple that contains the values of the red, green, and blue channels in the range 0 to 255. + If alpha is included, it is in the range 0 to 1. + """ + r, g, b = (float_to_255(c) for c in self._rgba[:3]) + if alpha is None: + if self._rgba.alpha is None: + return r, g, b + else: + return r, g, b, self._alpha_float() + elif alpha: + return r, g, b, self._alpha_float() + else: + # alpha is False + return r, g, b + + def as_hsl(self) -> str: + """Color as an `hsl(, , )` or `hsl(, , , )` string.""" + if self._rgba.alpha is None: + h, s, li = self.as_hsl_tuple(alpha=False) # type: ignore + return f'hsl({h * 360:0.0f}, {s:0.0%}, {li:0.0%})' + else: + h, s, li, a = self.as_hsl_tuple(alpha=True) # type: ignore + return f'hsl({h * 360:0.0f}, {s:0.0%}, {li:0.0%}, {round(a, 2)})' + + def as_hsl_tuple(self, *, alpha: Optional[bool] = None) -> HslColorTuple: + """Returns the color as an HSL or HSLA tuple. + + Args: + alpha: Whether to include the alpha channel. + + - `None` (default): Include the alpha channel only if it's set (e.g. not `None`). + - `True`: Always include alpha. + - `False`: Always omit alpha. + + Returns: + The color as a tuple of hue, saturation, lightness, and alpha (if included). + All elements are in the range 0 to 1. + + Note: + This is HSL as used in HTML and most other places, not HLS as used in Python's `colorsys`. + """ + h, l, s = rgb_to_hls(self._rgba.r, self._rgba.g, self._rgba.b) # noqa: E741 + if alpha is None: + if self._rgba.alpha is None: + return h, s, l + else: + return h, s, l, self._alpha_float() + if alpha: + return h, s, l, self._alpha_float() + else: + # alpha is False + return h, s, l + + def _alpha_float(self) -> float: + return 1 if self._rgba.alpha is None else self._rgba.alpha + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[Any], handler: Callable[[Any], CoreSchema] + ) -> core_schema.CoreSchema: + return core_schema.with_info_plain_validator_function( + cls._validate, serialization=core_schema.to_string_ser_schema() + ) + + @classmethod + def _validate(cls, __input_value: Any, _: Any) -> 'Color': + return cls(__input_value) + + def __str__(self) -> str: + return self.as_named(fallback=True) + + def __repr_args__(self) -> '_repr.ReprArgs': + return [(None, self.as_named(fallback=True))] + [('rgb', self.as_rgb_tuple())] + + def __eq__(self, other: Any) -> bool: + return isinstance(other, Color) and self.as_rgb_tuple() == other.as_rgb_tuple() + + def __hash__(self) -> int: + return hash(self.as_rgb_tuple()) + + +def parse_tuple(value: tuple[Any, ...]) -> RGBA: + """Parse a tuple or list to get RGBA values. + + Args: + value: A tuple or list. + + Returns: + An `RGBA` tuple parsed from the input tuple. + + Raises: + PydanticCustomError: If tuple is not valid. + """ + if len(value) == 3: + r, g, b = (parse_color_value(v) for v in value) + return RGBA(r, g, b, None) + elif len(value) == 4: + r, g, b = (parse_color_value(v) for v in value[:3]) + return RGBA(r, g, b, parse_float_alpha(value[3])) + else: + raise PydanticCustomError('color_error', 'value is not a valid color: tuples must have length 3 or 4') + + +def parse_str(value: str) -> RGBA: + """Parse a string representing a color to an RGBA tuple. + + Possible formats for the input string include: + + * named color, see `COLORS_BY_NAME` + * hex short eg. `fff` (prefix can be `#`, `0x` or nothing) + * hex long eg. `ffffff` (prefix can be `#`, `0x` or nothing) + * `rgb(, , )` + * `rgba(, , , )` + + Args: + value: A string representing a color. + + Returns: + An `RGBA` tuple parsed from the input string. + + Raises: + ValueError: If the input string cannot be parsed to an RGBA tuple. + """ + value_lower = value.lower() + try: + r, g, b = COLORS_BY_NAME[value_lower] + except KeyError: + pass + else: + return ints_to_rgba(r, g, b, None) + + m = re.fullmatch(r_hex_short, value_lower) + if m: + *rgb, a = m.groups() + r, g, b = (int(v * 2, 16) for v in rgb) + if a: + alpha: Optional[float] = int(a * 2, 16) / 255 + else: + alpha = None + return ints_to_rgba(r, g, b, alpha) + + m = re.fullmatch(r_hex_long, value_lower) + if m: + *rgb, a = m.groups() + r, g, b = (int(v, 16) for v in rgb) + if a: + alpha = int(a, 16) / 255 + else: + alpha = None + return ints_to_rgba(r, g, b, alpha) + + m = re.fullmatch(r_rgb, value_lower) or re.fullmatch(r_rgb_v4_style, value_lower) + if m: + return ints_to_rgba(*m.groups()) # type: ignore + + m = re.fullmatch(r_hsl, value_lower) or re.fullmatch(r_hsl_v4_style, value_lower) + if m: + return parse_hsl(*m.groups()) # type: ignore + + raise PydanticCustomError('color_error', 'value is not a valid color: string not recognised as a valid color') + + +def ints_to_rgba(r: Union[int, str], g: Union[int, str], b: Union[int, str], alpha: Optional[float] = None) -> RGBA: + """Converts integer or string values for RGB color and an optional alpha value to an `RGBA` object. + + Args: + r: An integer or string representing the red color value. + g: An integer or string representing the green color value. + b: An integer or string representing the blue color value. + alpha: A float representing the alpha value. Defaults to None. + + Returns: + An instance of the `RGBA` class with the corresponding color and alpha values. + """ + return RGBA(parse_color_value(r), parse_color_value(g), parse_color_value(b), parse_float_alpha(alpha)) + + +def parse_color_value(value: Union[int, str], max_val: int = 255) -> float: + """Parse the color value provided and return a number between 0 and 1. + + Args: + value: An integer or string color value. + max_val: Maximum range value. Defaults to 255. + + Raises: + PydanticCustomError: If the value is not a valid color. + + Returns: + A number between 0 and 1. + """ + try: + color = float(value) + except ValueError: + raise PydanticCustomError('color_error', 'value is not a valid color: color values must be a valid number') + if 0 <= color <= max_val: + return color / max_val + else: + raise PydanticCustomError( + 'color_error', + 'value is not a valid color: color values must be in the range 0 to {max_val}', + {'max_val': max_val}, + ) + + +def parse_float_alpha(value: Union[None, str, float, int]) -> Optional[float]: + """Parse an alpha value checking it's a valid float in the range 0 to 1. + + Args: + value: The input value to parse. + + Returns: + The parsed value as a float, or `None` if the value was None or equal 1. + + Raises: + PydanticCustomError: If the input value cannot be successfully parsed as a float in the expected range. + """ + if value is None: + return None + try: + if isinstance(value, str) and value.endswith('%'): + alpha = float(value[:-1]) / 100 + else: + alpha = float(value) + except ValueError: + raise PydanticCustomError('color_error', 'value is not a valid color: alpha values must be a valid float') + + if math.isclose(alpha, 1): + return None + elif 0 <= alpha <= 1: + return alpha + else: + raise PydanticCustomError('color_error', 'value is not a valid color: alpha values must be in the range 0 to 1') + + +def parse_hsl(h: str, h_units: str, sat: str, light: str, alpha: Optional[float] = None) -> RGBA: + """Parse raw hue, saturation, lightness, and alpha values and convert to RGBA. + + Args: + h: The hue value. + h_units: The unit for hue value. + sat: The saturation value. + light: The lightness value. + alpha: Alpha value. + + Returns: + An instance of `RGBA`. + """ + s_value, l_value = parse_color_value(sat, 100), parse_color_value(light, 100) + + h_value = float(h) + if h_units in {None, 'deg'}: + h_value = h_value % 360 / 360 + elif h_units == 'rad': + h_value = h_value % rads / rads + else: + # turns + h_value = h_value % 1 + + r, g, b = hls_to_rgb(h_value, l_value, s_value) + return RGBA(r, g, b, parse_float_alpha(alpha)) + + +def float_to_255(c: float) -> int: + """Converts a float value between 0 and 1 (inclusive) to an integer between 0 and 255 (inclusive). + + Args: + c: The float value to be converted. Must be between 0 and 1 (inclusive). + + Returns: + The integer equivalent of the given float value rounded to the nearest whole number. + + Raises: + ValueError: If the given float value is outside the acceptable range of 0 to 1 (inclusive). + """ + return int(round(c * 255)) + + +COLORS_BY_NAME = { + 'aliceblue': (240, 248, 255), + 'antiquewhite': (250, 235, 215), + 'aqua': (0, 255, 255), + 'aquamarine': (127, 255, 212), + 'azure': (240, 255, 255), + 'beige': (245, 245, 220), + 'bisque': (255, 228, 196), + 'black': (0, 0, 0), + 'blanchedalmond': (255, 235, 205), + 'blue': (0, 0, 255), + 'blueviolet': (138, 43, 226), + 'brown': (165, 42, 42), + 'burlywood': (222, 184, 135), + 'cadetblue': (95, 158, 160), + 'chartreuse': (127, 255, 0), + 'chocolate': (210, 105, 30), + 'coral': (255, 127, 80), + 'cornflowerblue': (100, 149, 237), + 'cornsilk': (255, 248, 220), + 'crimson': (220, 20, 60), + 'cyan': (0, 255, 255), + 'darkblue': (0, 0, 139), + 'darkcyan': (0, 139, 139), + 'darkgoldenrod': (184, 134, 11), + 'darkgray': (169, 169, 169), + 'darkgreen': (0, 100, 0), + 'darkgrey': (169, 169, 169), + 'darkkhaki': (189, 183, 107), + 'darkmagenta': (139, 0, 139), + 'darkolivegreen': (85, 107, 47), + 'darkorange': (255, 140, 0), + 'darkorchid': (153, 50, 204), + 'darkred': (139, 0, 0), + 'darksalmon': (233, 150, 122), + 'darkseagreen': (143, 188, 143), + 'darkslateblue': (72, 61, 139), + 'darkslategray': (47, 79, 79), + 'darkslategrey': (47, 79, 79), + 'darkturquoise': (0, 206, 209), + 'darkviolet': (148, 0, 211), + 'deeppink': (255, 20, 147), + 'deepskyblue': (0, 191, 255), + 'dimgray': (105, 105, 105), + 'dimgrey': (105, 105, 105), + 'dodgerblue': (30, 144, 255), + 'firebrick': (178, 34, 34), + 'floralwhite': (255, 250, 240), + 'forestgreen': (34, 139, 34), + 'fuchsia': (255, 0, 255), + 'gainsboro': (220, 220, 220), + 'ghostwhite': (248, 248, 255), + 'gold': (255, 215, 0), + 'goldenrod': (218, 165, 32), + 'gray': (128, 128, 128), + 'green': (0, 128, 0), + 'greenyellow': (173, 255, 47), + 'grey': (128, 128, 128), + 'honeydew': (240, 255, 240), + 'hotpink': (255, 105, 180), + 'indianred': (205, 92, 92), + 'indigo': (75, 0, 130), + 'ivory': (255, 255, 240), + 'khaki': (240, 230, 140), + 'lavender': (230, 230, 250), + 'lavenderblush': (255, 240, 245), + 'lawngreen': (124, 252, 0), + 'lemonchiffon': (255, 250, 205), + 'lightblue': (173, 216, 230), + 'lightcoral': (240, 128, 128), + 'lightcyan': (224, 255, 255), + 'lightgoldenrodyellow': (250, 250, 210), + 'lightgray': (211, 211, 211), + 'lightgreen': (144, 238, 144), + 'lightgrey': (211, 211, 211), + 'lightpink': (255, 182, 193), + 'lightsalmon': (255, 160, 122), + 'lightseagreen': (32, 178, 170), + 'lightskyblue': (135, 206, 250), + 'lightslategray': (119, 136, 153), + 'lightslategrey': (119, 136, 153), + 'lightsteelblue': (176, 196, 222), + 'lightyellow': (255, 255, 224), + 'lime': (0, 255, 0), + 'limegreen': (50, 205, 50), + 'linen': (250, 240, 230), + 'magenta': (255, 0, 255), + 'maroon': (128, 0, 0), + 'mediumaquamarine': (102, 205, 170), + 'mediumblue': (0, 0, 205), + 'mediumorchid': (186, 85, 211), + 'mediumpurple': (147, 112, 219), + 'mediumseagreen': (60, 179, 113), + 'mediumslateblue': (123, 104, 238), + 'mediumspringgreen': (0, 250, 154), + 'mediumturquoise': (72, 209, 204), + 'mediumvioletred': (199, 21, 133), + 'midnightblue': (25, 25, 112), + 'mintcream': (245, 255, 250), + 'mistyrose': (255, 228, 225), + 'moccasin': (255, 228, 181), + 'navajowhite': (255, 222, 173), + 'navy': (0, 0, 128), + 'oldlace': (253, 245, 230), + 'olive': (128, 128, 0), + 'olivedrab': (107, 142, 35), + 'orange': (255, 165, 0), + 'orangered': (255, 69, 0), + 'orchid': (218, 112, 214), + 'palegoldenrod': (238, 232, 170), + 'palegreen': (152, 251, 152), + 'paleturquoise': (175, 238, 238), + 'palevioletred': (219, 112, 147), + 'papayawhip': (255, 239, 213), + 'peachpuff': (255, 218, 185), + 'peru': (205, 133, 63), + 'pink': (255, 192, 203), + 'plum': (221, 160, 221), + 'powderblue': (176, 224, 230), + 'purple': (128, 0, 128), + 'red': (255, 0, 0), + 'rosybrown': (188, 143, 143), + 'royalblue': (65, 105, 225), + 'saddlebrown': (139, 69, 19), + 'salmon': (250, 128, 114), + 'sandybrown': (244, 164, 96), + 'seagreen': (46, 139, 87), + 'seashell': (255, 245, 238), + 'sienna': (160, 82, 45), + 'silver': (192, 192, 192), + 'skyblue': (135, 206, 235), + 'slateblue': (106, 90, 205), + 'slategray': (112, 128, 144), + 'slategrey': (112, 128, 144), + 'snow': (255, 250, 250), + 'springgreen': (0, 255, 127), + 'steelblue': (70, 130, 180), + 'tan': (210, 180, 140), + 'teal': (0, 128, 128), + 'thistle': (216, 191, 216), + 'tomato': (255, 99, 71), + 'turquoise': (64, 224, 208), + 'violet': (238, 130, 238), + 'wheat': (245, 222, 179), + 'white': (255, 255, 255), + 'whitesmoke': (245, 245, 245), + 'yellow': (255, 255, 0), + 'yellowgreen': (154, 205, 50), +} + +COLORS_BY_VALUE = {v: k for k, v in COLORS_BY_NAME.items()} diff --git a/lib/python3.12/site-packages/pydantic/config.py b/lib/python3.12/site-packages/pydantic/config.py new file mode 100644 index 0000000000000000000000000000000000000000..12fef106e53a8e7147b18cb2bb2545de7f646a00 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/config.py @@ -0,0 +1,1213 @@ +"""Configuration for Pydantic models.""" + +from __future__ import annotations as _annotations + +import warnings +from re import Pattern +from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar, Union, cast, overload + +from typing_extensions import TypeAlias, TypedDict, Unpack, deprecated + +from ._migration import getattr_migration +from .aliases import AliasGenerator +from .errors import PydanticUserError +from .warnings import PydanticDeprecatedSince211 + +if TYPE_CHECKING: + from ._internal._generate_schema import GenerateSchema as _GenerateSchema + from .fields import ComputedFieldInfo, FieldInfo + +__all__ = ('ConfigDict', 'with_config') + + +JsonValue: TypeAlias = Union[int, float, str, bool, None, list['JsonValue'], 'JsonDict'] +JsonDict: TypeAlias = dict[str, JsonValue] + +JsonEncoder = Callable[[Any], Any] + +JsonSchemaExtraCallable: TypeAlias = Union[ + Callable[[JsonDict], None], + Callable[[JsonDict, type[Any]], None], +] + +ExtraValues = Literal['allow', 'ignore', 'forbid'] + + +class ConfigDict(TypedDict, total=False): + """A TypedDict for configuring Pydantic behaviour.""" + + title: str | None + """The title for the generated JSON schema, defaults to the model's name""" + + model_title_generator: Callable[[type], str] | None + """A callable that takes a model class and returns the title for it. Defaults to `None`.""" + + field_title_generator: Callable[[str, FieldInfo | ComputedFieldInfo], str] | None + """A callable that takes a field's name and info and returns title for it. Defaults to `None`.""" + + str_to_lower: bool + """Whether to convert all characters to lowercase for str types. Defaults to `False`.""" + + str_to_upper: bool + """Whether to convert all characters to uppercase for str types. Defaults to `False`.""" + + str_strip_whitespace: bool + """Whether to strip leading and trailing whitespace for str types.""" + + str_min_length: int + """The minimum length for str types. Defaults to `None`.""" + + str_max_length: int | None + """The maximum length for str types. Defaults to `None`.""" + + extra: ExtraValues | None + ''' + Whether to ignore, allow, or forbid extra data during model initialization. Defaults to `'ignore'`. + + Three configuration values are available: + + - `'ignore'`: Providing extra data is ignored (the default): + ```python + from pydantic import BaseModel, ConfigDict + + class User(BaseModel): + model_config = ConfigDict(extra='ignore') # (1)! + + name: str + + user = User(name='John Doe', age=20) # (2)! + print(user) + #> name='John Doe' + ``` + + 1. This is the default behaviour. + 2. The `age` argument is ignored. + + - `'forbid'`: Providing extra data is not permitted, and a [`ValidationError`][pydantic_core.ValidationError] + will be raised if this is the case: + ```python + from pydantic import BaseModel, ConfigDict, ValidationError + + + class Model(BaseModel): + x: int + + model_config = ConfigDict(extra='forbid') + + + try: + Model(x=1, y='a') + except ValidationError as exc: + print(exc) + """ + 1 validation error for Model + y + Extra inputs are not permitted [type=extra_forbidden, input_value='a', input_type=str] + """ + ``` + + - `'allow'`: Providing extra data is allowed and stored in the `__pydantic_extra__` dictionary attribute: + ```python + from pydantic import BaseModel, ConfigDict + + + class Model(BaseModel): + x: int + + model_config = ConfigDict(extra='allow') + + + m = Model(x=1, y='a') + assert m.__pydantic_extra__ == {'y': 'a'} + ``` + By default, no validation will be applied to these extra items, but you can set a type for the values by overriding + the type annotation for `__pydantic_extra__`: + ```python + from pydantic import BaseModel, ConfigDict, Field, ValidationError + + + class Model(BaseModel): + __pydantic_extra__: dict[str, int] = Field(init=False) # (1)! + + x: int + + model_config = ConfigDict(extra='allow') + + + try: + Model(x=1, y='a') + except ValidationError as exc: + print(exc) + """ + 1 validation error for Model + y + Input should be a valid integer, unable to parse string as an integer [type=int_parsing, input_value='a', input_type=str] + """ + + m = Model(x=1, y='2') + assert m.x == 1 + assert m.y == 2 + assert m.model_dump() == {'x': 1, 'y': 2} + assert m.__pydantic_extra__ == {'y': 2} + ``` + + 1. The `= Field(init=False)` does not have any effect at runtime, but prevents the `__pydantic_extra__` field from + being included as a parameter to the model's `__init__` method by type checkers. + ''' + + frozen: bool + """ + Whether models are faux-immutable, i.e. whether `__setattr__` is allowed, and also generates + a `__hash__()` method for the model. This makes instances of the model potentially hashable if all the + attributes are hashable. Defaults to `False`. + + Note: + On V1, the inverse of this setting was called `allow_mutation`, and was `True` by default. + """ + + populate_by_name: bool + """ + Whether an aliased field may be populated by its name as given by the model + attribute, as well as the alias. Defaults to `False`. + + !!! warning + `populate_by_name` usage is not recommended in v2.11+ and will be deprecated in v3. + Instead, you should use the [`validate_by_name`][pydantic.config.ConfigDict.validate_by_name] configuration setting. + + When `validate_by_name=True` and `validate_by_alias=True`, this is strictly equivalent to the + previous behavior of `populate_by_name=True`. + + In v2.11, we also introduced a [`validate_by_alias`][pydantic.config.ConfigDict.validate_by_alias] setting that introduces more fine grained + control for validation behavior. + + Here's how you might go about using the new settings to achieve the same behavior: + + ```python + from pydantic import BaseModel, ConfigDict, Field + + class Model(BaseModel): + model_config = ConfigDict(validate_by_name=True, validate_by_alias=True) + + my_field: str = Field(alias='my_alias') # (1)! + + m = Model(my_alias='foo') # (2)! + print(m) + #> my_field='foo' + + m = Model(my_alias='foo') # (3)! + print(m) + #> my_field='foo' + ``` + + 1. The field `'my_field'` has an alias `'my_alias'`. + 2. The model is populated by the alias `'my_alias'`. + 3. The model is populated by the attribute name `'my_field'`. + """ + + use_enum_values: bool + """ + Whether to populate models with the `value` property of enums, rather than the raw enum. + This may be useful if you want to serialize `model.model_dump()` later. Defaults to `False`. + + !!! note + If you have an `Optional[Enum]` value that you set a default for, you need to use `validate_default=True` + for said Field to ensure that the `use_enum_values` flag takes effect on the default, as extracting an + enum's value occurs during validation, not serialization. + + ```python + from enum import Enum + from typing import Optional + + from pydantic import BaseModel, ConfigDict, Field + + class SomeEnum(Enum): + FOO = 'foo' + BAR = 'bar' + BAZ = 'baz' + + class SomeModel(BaseModel): + model_config = ConfigDict(use_enum_values=True) + + some_enum: SomeEnum + another_enum: Optional[SomeEnum] = Field( + default=SomeEnum.FOO, validate_default=True + ) + + model1 = SomeModel(some_enum=SomeEnum.BAR) + print(model1.model_dump()) + #> {'some_enum': 'bar', 'another_enum': 'foo'} + + model2 = SomeModel(some_enum=SomeEnum.BAR, another_enum=SomeEnum.BAZ) + print(model2.model_dump()) + #> {'some_enum': 'bar', 'another_enum': 'baz'} + ``` + """ + + validate_assignment: bool + """ + Whether to validate the data when the model is changed. Defaults to `False`. + + The default behavior of Pydantic is to validate the data when the model is created. + + In case the user changes the data after the model is created, the model is _not_ revalidated. + + ```python + from pydantic import BaseModel + + class User(BaseModel): + name: str + + user = User(name='John Doe') # (1)! + print(user) + #> name='John Doe' + user.name = 123 # (1)! + print(user) + #> name=123 + ``` + + 1. The validation happens only when the model is created. + 2. The validation does not happen when the data is changed. + + In case you want to revalidate the model when the data is changed, you can use `validate_assignment=True`: + + ```python + from pydantic import BaseModel, ValidationError + + class User(BaseModel, validate_assignment=True): # (1)! + name: str + + user = User(name='John Doe') # (2)! + print(user) + #> name='John Doe' + try: + user.name = 123 # (3)! + except ValidationError as e: + print(e) + ''' + 1 validation error for User + name + Input should be a valid string [type=string_type, input_value=123, input_type=int] + ''' + ``` + + 1. You can either use class keyword arguments, or `model_config` to set `validate_assignment=True`. + 2. The validation happens when the model is created. + 3. The validation _also_ happens when the data is changed. + """ + + arbitrary_types_allowed: bool + """ + Whether arbitrary types are allowed for field types. Defaults to `False`. + + ```python + from pydantic import BaseModel, ConfigDict, ValidationError + + # This is not a pydantic model, it's an arbitrary class + class Pet: + def __init__(self, name: str): + self.name = name + + class Model(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + pet: Pet + owner: str + + pet = Pet(name='Hedwig') + # A simple check of instance type is used to validate the data + model = Model(owner='Harry', pet=pet) + print(model) + #> pet=<__main__.Pet object at 0x0123456789ab> owner='Harry' + print(model.pet) + #> <__main__.Pet object at 0x0123456789ab> + print(model.pet.name) + #> Hedwig + print(type(model.pet)) + #> + try: + # If the value is not an instance of the type, it's invalid + Model(owner='Harry', pet='Hedwig') + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + pet + Input should be an instance of Pet [type=is_instance_of, input_value='Hedwig', input_type=str] + ''' + + # Nothing in the instance of the arbitrary type is checked + # Here name probably should have been a str, but it's not validated + pet2 = Pet(name=42) + model2 = Model(owner='Harry', pet=pet2) + print(model2) + #> pet=<__main__.Pet object at 0x0123456789ab> owner='Harry' + print(model2.pet) + #> <__main__.Pet object at 0x0123456789ab> + print(model2.pet.name) + #> 42 + print(type(model2.pet)) + #> + ``` + """ + + from_attributes: bool + """ + Whether to build models and look up discriminators of tagged unions using python object attributes. + """ + + loc_by_alias: bool + """Whether to use the actual key provided in the data (e.g. alias) for error `loc`s rather than the field's name. Defaults to `True`.""" + + alias_generator: Callable[[str], str] | AliasGenerator | None + """ + A callable that takes a field name and returns an alias for it + or an instance of [`AliasGenerator`][pydantic.aliases.AliasGenerator]. Defaults to `None`. + + When using a callable, the alias generator is used for both validation and serialization. + If you want to use different alias generators for validation and serialization, you can use + [`AliasGenerator`][pydantic.aliases.AliasGenerator] instead. + + If data source field names do not match your code style (e. g. CamelCase fields), + you can automatically generate aliases using `alias_generator`. Here's an example with + a basic callable: + + ```python + from pydantic import BaseModel, ConfigDict + from pydantic.alias_generators import to_pascal + + class Voice(BaseModel): + model_config = ConfigDict(alias_generator=to_pascal) + + name: str + language_code: str + + voice = Voice(Name='Filiz', LanguageCode='tr-TR') + print(voice.language_code) + #> tr-TR + print(voice.model_dump(by_alias=True)) + #> {'Name': 'Filiz', 'LanguageCode': 'tr-TR'} + ``` + + If you want to use different alias generators for validation and serialization, you can use + [`AliasGenerator`][pydantic.aliases.AliasGenerator]. + + ```python + from pydantic import AliasGenerator, BaseModel, ConfigDict + from pydantic.alias_generators import to_camel, to_pascal + + class Athlete(BaseModel): + first_name: str + last_name: str + sport: str + + model_config = ConfigDict( + alias_generator=AliasGenerator( + validation_alias=to_camel, + serialization_alias=to_pascal, + ) + ) + + athlete = Athlete(firstName='John', lastName='Doe', sport='track') + print(athlete.model_dump(by_alias=True)) + #> {'FirstName': 'John', 'LastName': 'Doe', 'Sport': 'track'} + ``` + + Note: + Pydantic offers three built-in alias generators: [`to_pascal`][pydantic.alias_generators.to_pascal], + [`to_camel`][pydantic.alias_generators.to_camel], and [`to_snake`][pydantic.alias_generators.to_snake]. + """ + + ignored_types: tuple[type, ...] + """A tuple of types that may occur as values of class attributes without annotations. This is + typically used for custom descriptors (classes that behave like `property`). If an attribute is set on a + class without an annotation and has a type that is not in this tuple (or otherwise recognized by + _pydantic_), an error will be raised. Defaults to `()`. + """ + + allow_inf_nan: bool + """Whether to allow infinity (`+inf` an `-inf`) and NaN values to float and decimal fields. Defaults to `True`.""" + + json_schema_extra: JsonDict | JsonSchemaExtraCallable | None + """A dict or callable to provide extra JSON schema properties. Defaults to `None`.""" + + json_encoders: dict[type[object], JsonEncoder] | None + """ + A `dict` of custom JSON encoders for specific types. Defaults to `None`. + + !!! warning "Deprecated" + This config option is a carryover from v1. + We originally planned to remove it in v2 but didn't have a 1:1 replacement so we are keeping it for now. + It is still deprecated and will likely be removed in the future. + """ + + # new in V2 + strict: bool + """ + _(new in V2)_ If `True`, strict validation is applied to all fields on the model. + + By default, Pydantic attempts to coerce values to the correct type, when possible. + + There are situations in which you may want to disable this behavior, and instead raise an error if a value's type + does not match the field's type annotation. + + To configure strict mode for all fields on a model, you can set `strict=True` on the model. + + ```python + from pydantic import BaseModel, ConfigDict + + class Model(BaseModel): + model_config = ConfigDict(strict=True) + + name: str + age: int + ``` + + See [Strict Mode](../concepts/strict_mode.md) for more details. + + See the [Conversion Table](../concepts/conversion_table.md) for more details on how Pydantic converts data in both + strict and lax modes. + """ + # whether instances of models and dataclasses (including subclass instances) should re-validate, default 'never' + revalidate_instances: Literal['always', 'never', 'subclass-instances'] + """ + When and how to revalidate models and dataclasses during validation. Accepts the string + values of `'never'`, `'always'` and `'subclass-instances'`. Defaults to `'never'`. + + - `'never'` will not revalidate models and dataclasses during validation + - `'always'` will revalidate models and dataclasses during validation + - `'subclass-instances'` will revalidate models and dataclasses during validation if the instance is a + subclass of the model or dataclass + + By default, model and dataclass instances are not revalidated during validation. + + ```python + from pydantic import BaseModel + + class User(BaseModel, revalidate_instances='never'): # (1)! + hobbies: list[str] + + class SubUser(User): + sins: list[str] + + class Transaction(BaseModel): + user: User + + my_user = User(hobbies=['reading']) + t = Transaction(user=my_user) + print(t) + #> user=User(hobbies=['reading']) + + my_user.hobbies = [1] # (2)! + t = Transaction(user=my_user) # (3)! + print(t) + #> user=User(hobbies=[1]) + + my_sub_user = SubUser(hobbies=['scuba diving'], sins=['lying']) + t = Transaction(user=my_sub_user) + print(t) + #> user=SubUser(hobbies=['scuba diving'], sins=['lying']) + ``` + + 1. `revalidate_instances` is set to `'never'` by **default. + 2. The assignment is not validated, unless you set `validate_assignment` to `True` in the model's config. + 3. Since `revalidate_instances` is set to `never`, this is not revalidated. + + If you want to revalidate instances during validation, you can set `revalidate_instances` to `'always'` + in the model's config. + + ```python + from pydantic import BaseModel, ValidationError + + class User(BaseModel, revalidate_instances='always'): # (1)! + hobbies: list[str] + + class SubUser(User): + sins: list[str] + + class Transaction(BaseModel): + user: User + + my_user = User(hobbies=['reading']) + t = Transaction(user=my_user) + print(t) + #> user=User(hobbies=['reading']) + + my_user.hobbies = [1] + try: + t = Transaction(user=my_user) # (2)! + except ValidationError as e: + print(e) + ''' + 1 validation error for Transaction + user.hobbies.0 + Input should be a valid string [type=string_type, input_value=1, input_type=int] + ''' + + my_sub_user = SubUser(hobbies=['scuba diving'], sins=['lying']) + t = Transaction(user=my_sub_user) + print(t) # (3)! + #> user=User(hobbies=['scuba diving']) + ``` + + 1. `revalidate_instances` is set to `'always'`. + 2. The model is revalidated, since `revalidate_instances` is set to `'always'`. + 3. Using `'never'` we would have gotten `user=SubUser(hobbies=['scuba diving'], sins=['lying'])`. + + It's also possible to set `revalidate_instances` to `'subclass-instances'` to only revalidate instances + of subclasses of the model. + + ```python + from pydantic import BaseModel + + class User(BaseModel, revalidate_instances='subclass-instances'): # (1)! + hobbies: list[str] + + class SubUser(User): + sins: list[str] + + class Transaction(BaseModel): + user: User + + my_user = User(hobbies=['reading']) + t = Transaction(user=my_user) + print(t) + #> user=User(hobbies=['reading']) + + my_user.hobbies = [1] + t = Transaction(user=my_user) # (2)! + print(t) + #> user=User(hobbies=[1]) + + my_sub_user = SubUser(hobbies=['scuba diving'], sins=['lying']) + t = Transaction(user=my_sub_user) + print(t) # (3)! + #> user=User(hobbies=['scuba diving']) + ``` + + 1. `revalidate_instances` is set to `'subclass-instances'`. + 2. This is not revalidated, since `my_user` is not a subclass of `User`. + 3. Using `'never'` we would have gotten `user=SubUser(hobbies=['scuba diving'], sins=['lying'])`. + """ + + ser_json_timedelta: Literal['iso8601', 'float'] + """ + The format of JSON serialized timedeltas. Accepts the string values of `'iso8601'` and + `'float'`. Defaults to `'iso8601'`. + + - `'iso8601'` will serialize timedeltas to ISO 8601 durations. + - `'float'` will serialize timedeltas to the total number of seconds. + """ + + ser_json_bytes: Literal['utf8', 'base64', 'hex'] + """ + The encoding of JSON serialized bytes. Defaults to `'utf8'`. + Set equal to `val_json_bytes` to get back an equal value after serialization round trip. + + - `'utf8'` will serialize bytes to UTF-8 strings. + - `'base64'` will serialize bytes to URL safe base64 strings. + - `'hex'` will serialize bytes to hexadecimal strings. + """ + + val_json_bytes: Literal['utf8', 'base64', 'hex'] + """ + The encoding of JSON serialized bytes to decode. Defaults to `'utf8'`. + Set equal to `ser_json_bytes` to get back an equal value after serialization round trip. + + - `'utf8'` will deserialize UTF-8 strings to bytes. + - `'base64'` will deserialize URL safe base64 strings to bytes. + - `'hex'` will deserialize hexadecimal strings to bytes. + """ + + ser_json_inf_nan: Literal['null', 'constants', 'strings'] + """ + The encoding of JSON serialized infinity and NaN float values. Defaults to `'null'`. + + - `'null'` will serialize infinity and NaN values as `null`. + - `'constants'` will serialize infinity and NaN values as `Infinity` and `NaN`. + - `'strings'` will serialize infinity as string `"Infinity"` and NaN as string `"NaN"`. + """ + + # whether to validate default values during validation, default False + validate_default: bool + """Whether to validate default values during validation. Defaults to `False`.""" + + validate_return: bool + """Whether to validate the return value from call validators. Defaults to `False`.""" + + protected_namespaces: tuple[str | Pattern[str], ...] + """ + A `tuple` of strings and/or patterns that prevent models from having fields with names that conflict with them. + For strings, we match on a prefix basis. Ex, if 'dog' is in the protected namespace, 'dog_name' will be protected. + For patterns, we match on the entire field name. Ex, if `re.compile(r'^dog$')` is in the protected namespace, 'dog' will be protected, but 'dog_name' will not be. + Defaults to `('model_validate', 'model_dump',)`. + + The reason we've selected these is to prevent collisions with other validation / dumping formats + in the future - ex, `model_validate_{some_newly_supported_format}`. + + Before v2.10, Pydantic used `('model_',)` as the default value for this setting to + prevent collisions between model attributes and `BaseModel`'s own methods. This was changed + in v2.10 given feedback that this restriction was limiting in AI and data science contexts, + where it is common to have fields with names like `model_id`, `model_input`, `model_output`, etc. + + For more details, see https://github.com/pydantic/pydantic/issues/10315. + + ```python + import warnings + + from pydantic import BaseModel + + warnings.filterwarnings('error') # Raise warnings as errors + + try: + + class Model(BaseModel): + model_dump_something: str + + except UserWarning as e: + print(e) + ''' + Field "model_dump_something" in Model has conflict with protected namespace "model_dump". + + You may be able to resolve this warning by setting `model_config['protected_namespaces'] = ('model_validate',)`. + ''' + ``` + + You can customize this behavior using the `protected_namespaces` setting: + + ```python {test="skip"} + import re + import warnings + + from pydantic import BaseModel, ConfigDict + + with warnings.catch_warnings(record=True) as caught_warnings: + warnings.simplefilter('always') # Catch all warnings + + class Model(BaseModel): + safe_field: str + also_protect_field: str + protect_this: str + + model_config = ConfigDict( + protected_namespaces=( + 'protect_me_', + 'also_protect_', + re.compile('^protect_this$'), + ) + ) + + for warning in caught_warnings: + print(f'{warning.message}') + ''' + Field "also_protect_field" in Model has conflict with protected namespace "also_protect_". + You may be able to resolve this warning by setting `model_config['protected_namespaces'] = ('protect_me_', re.compile('^protect_this$'))`. + + Field "protect_this" in Model has conflict with protected namespace "re.compile('^protect_this$')". + You may be able to resolve this warning by setting `model_config['protected_namespaces'] = ('protect_me_', 'also_protect_')`. + ''' + ``` + + While Pydantic will only emit a warning when an item is in a protected namespace but does not actually have a collision, + an error _is_ raised if there is an actual collision with an existing attribute: + + ```python + from pydantic import BaseModel, ConfigDict + + try: + + class Model(BaseModel): + model_validate: str + + model_config = ConfigDict(protected_namespaces=('model_',)) + + except NameError as e: + print(e) + ''' + Field "model_validate" conflicts with member > of protected namespace "model_". + ''' + ``` + """ + + hide_input_in_errors: bool + """ + Whether to hide inputs when printing errors. Defaults to `False`. + + Pydantic shows the input value and type when it raises `ValidationError` during the validation. + + ```python + from pydantic import BaseModel, ValidationError + + class Model(BaseModel): + a: str + + try: + Model(a=123) + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + a + Input should be a valid string [type=string_type, input_value=123, input_type=int] + ''' + ``` + + You can hide the input value and type by setting the `hide_input_in_errors` config to `True`. + + ```python + from pydantic import BaseModel, ConfigDict, ValidationError + + class Model(BaseModel): + a: str + model_config = ConfigDict(hide_input_in_errors=True) + + try: + Model(a=123) + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + a + Input should be a valid string [type=string_type] + ''' + ``` + """ + + defer_build: bool + """ + Whether to defer model validator and serializer construction until the first model validation. Defaults to False. + + This can be useful to avoid the overhead of building models which are only + used nested within other models, or when you want to manually define type namespace via + [`Model.model_rebuild(_types_namespace=...)`][pydantic.BaseModel.model_rebuild]. + + Since v2.10, this setting also applies to pydantic dataclasses and TypeAdapter instances. + """ + + plugin_settings: dict[str, object] | None + """A `dict` of settings for plugins. Defaults to `None`.""" + + schema_generator: type[_GenerateSchema] | None + """ + !!! warning + `schema_generator` is deprecated in v2.10. + + Prior to v2.10, this setting was advertised as highly subject to change. + It's possible that this interface may once again become public once the internal core schema generation + API is more stable, but that will likely come after significant performance improvements have been made. + """ + + json_schema_serialization_defaults_required: bool + """ + Whether fields with default values should be marked as required in the serialization schema. Defaults to `False`. + + This ensures that the serialization schema will reflect the fact a field with a default will always be present + when serializing the model, even though it is not required for validation. + + However, there are scenarios where this may be undesirable — in particular, if you want to share the schema + between validation and serialization, and don't mind fields with defaults being marked as not required during + serialization. See [#7209](https://github.com/pydantic/pydantic/issues/7209) for more details. + + ```python + from pydantic import BaseModel, ConfigDict + + class Model(BaseModel): + a: str = 'a' + + model_config = ConfigDict(json_schema_serialization_defaults_required=True) + + print(Model.model_json_schema(mode='validation')) + ''' + { + 'properties': {'a': {'default': 'a', 'title': 'A', 'type': 'string'}}, + 'title': 'Model', + 'type': 'object', + } + ''' + print(Model.model_json_schema(mode='serialization')) + ''' + { + 'properties': {'a': {'default': 'a', 'title': 'A', 'type': 'string'}}, + 'required': ['a'], + 'title': 'Model', + 'type': 'object', + } + ''' + ``` + """ + + json_schema_mode_override: Literal['validation', 'serialization', None] + """ + If not `None`, the specified mode will be used to generate the JSON schema regardless of what `mode` was passed to + the function call. Defaults to `None`. + + This provides a way to force the JSON schema generation to reflect a specific mode, e.g., to always use the + validation schema. + + It can be useful when using frameworks (such as FastAPI) that may generate different schemas for validation + and serialization that must both be referenced from the same schema; when this happens, we automatically append + `-Input` to the definition reference for the validation schema and `-Output` to the definition reference for the + serialization schema. By specifying a `json_schema_mode_override` though, this prevents the conflict between + the validation and serialization schemas (since both will use the specified schema), and so prevents the suffixes + from being added to the definition references. + + ```python + from pydantic import BaseModel, ConfigDict, Json + + class Model(BaseModel): + a: Json[int] # requires a string to validate, but will dump an int + + print(Model.model_json_schema(mode='serialization')) + ''' + { + 'properties': {'a': {'title': 'A', 'type': 'integer'}}, + 'required': ['a'], + 'title': 'Model', + 'type': 'object', + } + ''' + + class ForceInputModel(Model): + # the following ensures that even with mode='serialization', we + # will get the schema that would be generated for validation. + model_config = ConfigDict(json_schema_mode_override='validation') + + print(ForceInputModel.model_json_schema(mode='serialization')) + ''' + { + 'properties': { + 'a': { + 'contentMediaType': 'application/json', + 'contentSchema': {'type': 'integer'}, + 'title': 'A', + 'type': 'string', + } + }, + 'required': ['a'], + 'title': 'ForceInputModel', + 'type': 'object', + } + ''' + ``` + """ + + coerce_numbers_to_str: bool + """ + If `True`, enables automatic coercion of any `Number` type to `str` in "lax" (non-strict) mode. Defaults to `False`. + + Pydantic doesn't allow number types (`int`, `float`, `Decimal`) to be coerced as type `str` by default. + + ```python + from decimal import Decimal + + from pydantic import BaseModel, ConfigDict, ValidationError + + class Model(BaseModel): + value: str + + try: + print(Model(value=42)) + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + value + Input should be a valid string [type=string_type, input_value=42, input_type=int] + ''' + + class Model(BaseModel): + model_config = ConfigDict(coerce_numbers_to_str=True) + + value: str + + repr(Model(value=42).value) + #> "42" + repr(Model(value=42.13).value) + #> "42.13" + repr(Model(value=Decimal('42.13')).value) + #> "42.13" + ``` + """ + + regex_engine: Literal['rust-regex', 'python-re'] + """ + The regex engine to be used for pattern validation. + Defaults to `'rust-regex'`. + + - `rust-regex` uses the [`regex`](https://docs.rs/regex) Rust crate, + which is non-backtracking and therefore more DDoS resistant, but does not support all regex features. + - `python-re` use the [`re`](https://docs.python.org/3/library/re.html) module, + which supports all regex features, but may be slower. + + !!! note + If you use a compiled regex pattern, the python-re engine will be used regardless of this setting. + This is so that flags such as `re.IGNORECASE` are respected. + + ```python + from pydantic import BaseModel, ConfigDict, Field, ValidationError + + class Model(BaseModel): + model_config = ConfigDict(regex_engine='python-re') + + value: str = Field(pattern=r'^abc(?=def)') + + print(Model(value='abcdef').value) + #> abcdef + + try: + print(Model(value='abxyzcdef')) + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + value + String should match pattern '^abc(?=def)' [type=string_pattern_mismatch, input_value='abxyzcdef', input_type=str] + ''' + ``` + """ + + validation_error_cause: bool + """ + If `True`, Python exceptions that were part of a validation failure will be shown as an exception group as a cause. Can be useful for debugging. Defaults to `False`. + + Note: + Python 3.10 and older don't support exception groups natively. <=3.10, backport must be installed: `pip install exceptiongroup`. + + Note: + The structure of validation errors are likely to change in future Pydantic versions. Pydantic offers no guarantees about their structure. Should be used for visual traceback debugging only. + """ + + use_attribute_docstrings: bool + ''' + Whether docstrings of attributes (bare string literals immediately following the attribute declaration) + should be used for field descriptions. Defaults to `False`. + + Available in Pydantic v2.7+. + + ```python + from pydantic import BaseModel, ConfigDict, Field + + + class Model(BaseModel): + model_config = ConfigDict(use_attribute_docstrings=True) + + x: str + """ + Example of an attribute docstring + """ + + y: int = Field(description="Description in Field") + """ + Description in Field overrides attribute docstring + """ + + + print(Model.model_fields["x"].description) + # > Example of an attribute docstring + print(Model.model_fields["y"].description) + # > Description in Field + ``` + This requires the source code of the class to be available at runtime. + + !!! warning "Usage with `TypedDict` and stdlib dataclasses" + Due to current limitations, attribute docstrings detection may not work as expected when using + [`TypedDict`][typing.TypedDict] and stdlib dataclasses, in particular when: + + - inheritance is being used. + - multiple classes have the same name in the same source file. + ''' + + cache_strings: bool | Literal['all', 'keys', 'none'] + """ + Whether to cache strings to avoid constructing new Python objects. Defaults to True. + + Enabling this setting should significantly improve validation performance while increasing memory usage slightly. + + - `True` or `'all'` (the default): cache all strings + - `'keys'`: cache only dictionary keys + - `False` or `'none'`: no caching + + !!! note + `True` or `'all'` is required to cache strings during general validation because + validators don't know if they're in a key or a value. + + !!! tip + If repeated strings are rare, it's recommended to use `'keys'` or `'none'` to reduce memory usage, + as the performance difference is minimal if repeated strings are rare. + """ + + validate_by_alias: bool + """ + Whether an aliased field may be populated by its alias. Defaults to `True`. + + !!! note + In v2.11, `validate_by_alias` was introduced in conjunction with [`validate_by_name`][pydantic.ConfigDict.validate_by_name] + to empower users with more fine grained validation control. In my_field='foo' + ``` + + 1. The field `'my_field'` has an alias `'my_alias'`. + 2. The model can only be populated by the attribute name `'my_field'`. + + !!! warning + You cannot set both `validate_by_alias` and `validate_by_name` to `False`. + This would make it impossible to populate an attribute. + + See [usage errors](../errors/usage_errors.md#validate-by-alias-and-name-false) for an example. + + If you set `validate_by_alias` to `False`, under the hood, Pydantic dynamically sets + `validate_by_name` to `True` to ensure that validation can still occur. + """ + + validate_by_name: bool + """ + Whether an aliased field may be populated by its name as given by the model + attribute. Defaults to `False`. + + !!! note + In v2.0-v2.10, the `populate_by_name` configuration setting was used to specify + whether or not a field could be populated by its name **and** alias. + + In v2.11, `validate_by_name` was introduced in conjunction with [`validate_by_alias`][pydantic.ConfigDict.validate_by_alias] + to empower users with more fine grained validation behavior control. + + ```python + from pydantic import BaseModel, ConfigDict, Field + + class Model(BaseModel): + model_config = ConfigDict(validate_by_name=True, validate_by_alias=True) + + my_field: str = Field(validation_alias='my_alias') # (1)! + + m = Model(my_alias='foo') # (2)! + print(m) + #> my_field='foo' + + m = Model(my_field='foo') # (3)! + print(m) + #> my_field='foo' + ``` + + 1. The field `'my_field'` has an alias `'my_alias'`. + 2. The model is populated by the alias `'my_alias'`. + 3. The model is populated by the attribute name `'my_field'`. + + !!! warning + You cannot set both `validate_by_alias` and `validate_by_name` to `False`. + This would make it impossible to populate an attribute. + + See [usage errors](../errors/usage_errors.md#validate-by-alias-and-name-false) for an example. + """ + + serialize_by_alias: bool + """ + Whether an aliased field should be serialized by its alias. Defaults to `False`. + + Note: In v2.11, `serialize_by_alias` was introduced to address the + [popular request](https://github.com/pydantic/pydantic/issues/8379) + for consistency with alias behavior for validation and serialization settings. + In v3, the default value is expected to change to `True` for consistency with the validation default. + + ```python + from pydantic import BaseModel, ConfigDict, Field + + class Model(BaseModel): + model_config = ConfigDict(serialize_by_alias=True) + + my_field: str = Field(serialization_alias='my_alias') # (1)! + + m = Model(my_field='foo') + print(m.model_dump()) # (2)! + #> {'my_alias': 'foo'} + ``` + + 1. The field `'my_field'` has an alias `'my_alias'`. + 2. The model is serialized using the alias `'my_alias'` for the `'my_field'` attribute. + """ + + +_TypeT = TypeVar('_TypeT', bound=type) + + +@overload +@deprecated('Passing `config` as a keyword argument is deprecated. Pass `config` as a positional argument instead.') +def with_config(*, config: ConfigDict) -> Callable[[_TypeT], _TypeT]: ... + + +@overload +def with_config(config: ConfigDict, /) -> Callable[[_TypeT], _TypeT]: ... + + +@overload +def with_config(**config: Unpack[ConfigDict]) -> Callable[[_TypeT], _TypeT]: ... + + +def with_config(config: ConfigDict | None = None, /, **kwargs: Any) -> Callable[[_TypeT], _TypeT]: + """!!! abstract "Usage Documentation" + [Configuration with other types](../concepts/config.md#configuration-on-other-supported-types) + + A convenience decorator to set a [Pydantic configuration](config.md) on a `TypedDict` or a `dataclass` from the standard library. + + Although the configuration can be set using the `__pydantic_config__` attribute, it does not play well with type checkers, + especially with `TypedDict`. + + !!! example "Usage" + + ```python + from typing_extensions import TypedDict + + from pydantic import ConfigDict, TypeAdapter, with_config + + @with_config(ConfigDict(str_to_lower=True)) + class TD(TypedDict): + x: str + + ta = TypeAdapter(TD) + + print(ta.validate_python({'x': 'ABC'})) + #> {'x': 'abc'} + ``` + """ + if config is not None and kwargs: + raise ValueError('Cannot specify both `config` and keyword arguments') + + if len(kwargs) == 1 and (kwargs_conf := kwargs.get('config')) is not None: + warnings.warn( + 'Passing `config` as a keyword argument is deprecated. Pass `config` as a positional argument instead', + category=PydanticDeprecatedSince211, + stacklevel=2, + ) + final_config = cast(ConfigDict, kwargs_conf) + else: + final_config = config if config is not None else cast(ConfigDict, kwargs) + + def inner(class_: _TypeT, /) -> _TypeT: + # Ideally, we would check for `class_` to either be a `TypedDict` or a stdlib dataclass. + # However, the `@with_config` decorator can be applied *after* `@dataclass`. To avoid + # common mistakes, we at least check for `class_` to not be a Pydantic model. + from ._internal._utils import is_model_class + + if is_model_class(class_): + raise PydanticUserError( + f'Cannot use `with_config` on {class_.__name__} as it is a Pydantic model', + code='with-config-on-model', + ) + class_.__pydantic_config__ = final_config + return class_ + + return inner + + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/dataclasses.py b/lib/python3.12/site-packages/pydantic/dataclasses.py new file mode 100644 index 0000000000000000000000000000000000000000..4e42d65c0c382d5ff3daf8cec9d2934d942d2eb3 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/dataclasses.py @@ -0,0 +1,383 @@ +"""Provide an enhanced dataclass that performs validation.""" + +from __future__ import annotations as _annotations + +import dataclasses +import sys +import types +from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, NoReturn, TypeVar, overload +from warnings import warn + +from typing_extensions import TypeGuard, dataclass_transform + +from ._internal import _config, _decorators, _namespace_utils, _typing_extra +from ._internal import _dataclasses as _pydantic_dataclasses +from ._migration import getattr_migration +from .config import ConfigDict +from .errors import PydanticUserError +from .fields import Field, FieldInfo, PrivateAttr + +if TYPE_CHECKING: + from ._internal._dataclasses import PydanticDataclass + from ._internal._namespace_utils import MappingNamespace + +__all__ = 'dataclass', 'rebuild_dataclass' + +_T = TypeVar('_T') + +if sys.version_info >= (3, 10): + + @dataclass_transform(field_specifiers=(dataclasses.field, Field, PrivateAttr)) + @overload + def dataclass( + *, + init: Literal[False] = False, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool = False, + config: ConfigDict | type[object] | None = None, + validate_on_init: bool | None = None, + kw_only: bool = ..., + slots: bool = ..., + ) -> Callable[[type[_T]], type[PydanticDataclass]]: # type: ignore + ... + + @dataclass_transform(field_specifiers=(dataclasses.field, Field, PrivateAttr)) + @overload + def dataclass( + _cls: type[_T], # type: ignore + *, + init: Literal[False] = False, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool | None = None, + config: ConfigDict | type[object] | None = None, + validate_on_init: bool | None = None, + kw_only: bool = ..., + slots: bool = ..., + ) -> type[PydanticDataclass]: ... + +else: + + @dataclass_transform(field_specifiers=(dataclasses.field, Field, PrivateAttr)) + @overload + def dataclass( + *, + init: Literal[False] = False, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool | None = None, + config: ConfigDict | type[object] | None = None, + validate_on_init: bool | None = None, + ) -> Callable[[type[_T]], type[PydanticDataclass]]: # type: ignore + ... + + @dataclass_transform(field_specifiers=(dataclasses.field, Field, PrivateAttr)) + @overload + def dataclass( + _cls: type[_T], # type: ignore + *, + init: Literal[False] = False, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool | None = None, + config: ConfigDict | type[object] | None = None, + validate_on_init: bool | None = None, + ) -> type[PydanticDataclass]: ... + + +@dataclass_transform(field_specifiers=(dataclasses.field, Field, PrivateAttr)) +def dataclass( + _cls: type[_T] | None = None, + *, + init: Literal[False] = False, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool | None = None, + config: ConfigDict | type[object] | None = None, + validate_on_init: bool | None = None, + kw_only: bool = False, + slots: bool = False, +) -> Callable[[type[_T]], type[PydanticDataclass]] | type[PydanticDataclass]: + """!!! abstract "Usage Documentation" + [`dataclasses`](../concepts/dataclasses.md) + + A decorator used to create a Pydantic-enhanced dataclass, similar to the standard Python `dataclass`, + but with added validation. + + This function should be used similarly to `dataclasses.dataclass`. + + Args: + _cls: The target `dataclass`. + init: Included for signature compatibility with `dataclasses.dataclass`, and is passed through to + `dataclasses.dataclass` when appropriate. If specified, must be set to `False`, as pydantic inserts its + own `__init__` function. + repr: A boolean indicating whether to include the field in the `__repr__` output. + eq: Determines if a `__eq__` method should be generated for the class. + order: Determines if comparison magic methods should be generated, such as `__lt__`, but not `__eq__`. + unsafe_hash: Determines if a `__hash__` method should be included in the class, as in `dataclasses.dataclass`. + frozen: Determines if the generated class should be a 'frozen' `dataclass`, which does not allow its + attributes to be modified after it has been initialized. If not set, the value from the provided `config` argument will be used (and will default to `False` otherwise). + config: The Pydantic config to use for the `dataclass`. + validate_on_init: A deprecated parameter included for backwards compatibility; in V2, all Pydantic dataclasses + are validated on init. + kw_only: Determines if `__init__` method parameters must be specified by keyword only. Defaults to `False`. + slots: Determines if the generated class should be a 'slots' `dataclass`, which does not allow the addition of + new attributes after instantiation. + + Returns: + A decorator that accepts a class as its argument and returns a Pydantic `dataclass`. + + Raises: + AssertionError: Raised if `init` is not `False` or `validate_on_init` is `False`. + """ + assert init is False, 'pydantic.dataclasses.dataclass only supports init=False' + assert validate_on_init is not False, 'validate_on_init=False is no longer supported' + + if sys.version_info >= (3, 10): + kwargs = {'kw_only': kw_only, 'slots': slots} + else: + kwargs = {} + + def make_pydantic_fields_compatible(cls: type[Any]) -> None: + """Make sure that stdlib `dataclasses` understands `Field` kwargs like `kw_only` + To do that, we simply change + `x: int = pydantic.Field(..., kw_only=True)` + into + `x: int = dataclasses.field(default=pydantic.Field(..., kw_only=True), kw_only=True)` + """ + for annotation_cls in cls.__mro__: + annotations: dict[str, Any] = getattr(annotation_cls, '__annotations__', {}) + for field_name in annotations: + field_value = getattr(cls, field_name, None) + # Process only if this is an instance of `FieldInfo`. + if not isinstance(field_value, FieldInfo): + continue + + # Initialize arguments for the standard `dataclasses.field`. + field_args: dict = {'default': field_value} + + # Handle `kw_only` for Python 3.10+ + if sys.version_info >= (3, 10) and field_value.kw_only: + field_args['kw_only'] = True + + # Set `repr` attribute if it's explicitly specified to be not `True`. + if field_value.repr is not True: + field_args['repr'] = field_value.repr + + setattr(cls, field_name, dataclasses.field(**field_args)) + # In Python 3.9, when subclassing, information is pulled from cls.__dict__['__annotations__'] + # for annotations, so we must make sure it's initialized before we add to it. + if cls.__dict__.get('__annotations__') is None: + cls.__annotations__ = {} + cls.__annotations__[field_name] = annotations[field_name] + + def create_dataclass(cls: type[Any]) -> type[PydanticDataclass]: + """Create a Pydantic dataclass from a regular dataclass. + + Args: + cls: The class to create the Pydantic dataclass from. + + Returns: + A Pydantic dataclass. + """ + from ._internal._utils import is_model_class + + if is_model_class(cls): + raise PydanticUserError( + f'Cannot create a Pydantic dataclass from {cls.__name__} as it is already a Pydantic model', + code='dataclass-on-model', + ) + + original_cls = cls + + # we warn on conflicting config specifications, but only if the class doesn't have a dataclass base + # because a dataclass base might provide a __pydantic_config__ attribute that we don't want to warn about + has_dataclass_base = any(dataclasses.is_dataclass(base) for base in cls.__bases__) + if not has_dataclass_base and config is not None and hasattr(cls, '__pydantic_config__'): + warn( + f'`config` is set via both the `dataclass` decorator and `__pydantic_config__` for dataclass {cls.__name__}. ' + f'The `config` specification from `dataclass` decorator will take priority.', + category=UserWarning, + stacklevel=2, + ) + + # if config is not explicitly provided, try to read it from the type + config_dict = config if config is not None else getattr(cls, '__pydantic_config__', None) + config_wrapper = _config.ConfigWrapper(config_dict) + decorators = _decorators.DecoratorInfos.build(cls) + + # Keep track of the original __doc__ so that we can restore it after applying the dataclasses decorator + # Otherwise, classes with no __doc__ will have their signature added into the JSON schema description, + # since dataclasses.dataclass will set this as the __doc__ + original_doc = cls.__doc__ + + if _pydantic_dataclasses.is_builtin_dataclass(cls): + # Don't preserve the docstring for vanilla dataclasses, as it may include the signature + # This matches v1 behavior, and there was an explicit test for it + original_doc = None + + # We don't want to add validation to the existing std lib dataclass, so we will subclass it + # If the class is generic, we need to make sure the subclass also inherits from Generic + # with all the same parameters. + bases = (cls,) + if issubclass(cls, Generic): + generic_base = Generic[cls.__parameters__] # type: ignore + bases = bases + (generic_base,) + cls = types.new_class(cls.__name__, bases) + + make_pydantic_fields_compatible(cls) + + # Respect frozen setting from dataclass constructor and fallback to config setting if not provided + if frozen is not None: + frozen_ = frozen + if config_wrapper.frozen: + # It's not recommended to define both, as the setting from the dataclass decorator will take priority. + warn( + f'`frozen` is set via both the `dataclass` decorator and `config` for dataclass {cls.__name__!r}.' + 'This is not recommended. The `frozen` specification on `dataclass` will take priority.', + category=UserWarning, + stacklevel=2, + ) + else: + frozen_ = config_wrapper.frozen or False + + cls = dataclasses.dataclass( # type: ignore[call-overload] + cls, + # the value of init here doesn't affect anything except that it makes it easier to generate a signature + init=True, + repr=repr, + eq=eq, + order=order, + unsafe_hash=unsafe_hash, + frozen=frozen_, + **kwargs, + ) + + # This is an undocumented attribute to distinguish stdlib/Pydantic dataclasses. + # It should be set as early as possible: + cls.__is_pydantic_dataclass__ = True + + cls.__pydantic_decorators__ = decorators # type: ignore + cls.__doc__ = original_doc + cls.__module__ = original_cls.__module__ + cls.__qualname__ = original_cls.__qualname__ + cls.__pydantic_fields_complete__ = classmethod(_pydantic_fields_complete) + cls.__pydantic_complete__ = False # `complete_dataclass` will set it to `True` if successful. + # TODO `parent_namespace` is currently None, but we could do the same thing as Pydantic models: + # fetch the parent ns using `parent_frame_namespace` (if the dataclass was defined in a function), + # and possibly cache it (see the `__pydantic_parent_namespace__` logic for models). + _pydantic_dataclasses.complete_dataclass(cls, config_wrapper, raise_errors=False) + return cls + + return create_dataclass if _cls is None else create_dataclass(_cls) + + +def _pydantic_fields_complete(cls: type[PydanticDataclass]) -> bool: + """Return whether the fields where successfully collected (i.e. type hints were successfully resolves). + + This is a private property, not meant to be used outside Pydantic. + """ + return all(field_info._complete for field_info in cls.__pydantic_fields__.values()) + + +__getattr__ = getattr_migration(__name__) + +if sys.version_info < (3, 11): + # Monkeypatch dataclasses.InitVar so that typing doesn't error if it occurs as a type when evaluating type hints + # Starting in 3.11, typing.get_type_hints will not raise an error if the retrieved type hints are not callable. + + def _call_initvar(*args: Any, **kwargs: Any) -> NoReturn: + """This function does nothing but raise an error that is as similar as possible to what you'd get + if you were to try calling `InitVar[int]()` without this monkeypatch. The whole purpose is just + to ensure typing._type_check does not error if the type hint evaluates to `InitVar[]`. + """ + raise TypeError("'InitVar' object is not callable") + + dataclasses.InitVar.__call__ = _call_initvar + + +def rebuild_dataclass( + cls: type[PydanticDataclass], + *, + force: bool = False, + raise_errors: bool = True, + _parent_namespace_depth: int = 2, + _types_namespace: MappingNamespace | None = None, +) -> bool | None: + """Try to rebuild the pydantic-core schema for the dataclass. + + This may be necessary when one of the annotations is a ForwardRef which could not be resolved during + the initial attempt to build the schema, and automatic rebuilding fails. + + This is analogous to `BaseModel.model_rebuild`. + + Args: + cls: The class to rebuild the pydantic-core schema for. + force: Whether to force the rebuilding of the schema, defaults to `False`. + raise_errors: Whether to raise errors, defaults to `True`. + _parent_namespace_depth: The depth level of the parent namespace, defaults to 2. + _types_namespace: The types namespace, defaults to `None`. + + Returns: + Returns `None` if the schema is already "complete" and rebuilding was not required. + If rebuilding _was_ required, returns `True` if rebuilding was successful, otherwise `False`. + """ + if not force and cls.__pydantic_complete__: + return None + + for attr in ('__pydantic_core_schema__', '__pydantic_validator__', '__pydantic_serializer__'): + if attr in cls.__dict__: + # Deleting the validator/serializer is necessary as otherwise they can get reused in + # pydantic-core. Same applies for the core schema that can be reused in schema generation. + delattr(cls, attr) + + cls.__pydantic_complete__ = False + + if _types_namespace is not None: + rebuild_ns = _types_namespace + elif _parent_namespace_depth > 0: + rebuild_ns = _typing_extra.parent_frame_namespace(parent_depth=_parent_namespace_depth, force=True) or {} + else: + rebuild_ns = {} + + ns_resolver = _namespace_utils.NsResolver( + parent_namespace=rebuild_ns, + ) + + return _pydantic_dataclasses.complete_dataclass( + cls, + _config.ConfigWrapper(cls.__pydantic_config__, check=False), + raise_errors=raise_errors, + ns_resolver=ns_resolver, + # We could provide a different config instead (with `'defer_build'` set to `True`) + # of this explicit `_force_build` argument, but because config can come from the + # decorator parameter or the `__pydantic_config__` attribute, `complete_dataclass` + # will overwrite `__pydantic_config__` with the provided config above: + _force_build=True, + ) + + +def is_pydantic_dataclass(class_: type[Any], /) -> TypeGuard[type[PydanticDataclass]]: + """Whether a class is a pydantic dataclass. + + Args: + class_: The class. + + Returns: + `True` if the class is a pydantic dataclass, `False` otherwise. + """ + try: + return '__is_pydantic_dataclass__' in class_.__dict__ and dataclasses.is_dataclass(class_) + except AttributeError: + return False diff --git a/lib/python3.12/site-packages/pydantic/datetime_parse.py b/lib/python3.12/site-packages/pydantic/datetime_parse.py new file mode 100644 index 0000000000000000000000000000000000000000..53d52649e7620965ed194240a3becaa3ce3e3448 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/datetime_parse.py @@ -0,0 +1,5 @@ +"""The `datetime_parse` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/decorator.py b/lib/python3.12/site-packages/pydantic/decorator.py new file mode 100644 index 0000000000000000000000000000000000000000..0d97560c1b791956726b04fd66740a947647aabe --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/decorator.py @@ -0,0 +1,5 @@ +"""The `decorator` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__init__.py b/lib/python3.12/site-packages/pydantic/deprecated/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b0038b8a1430eb79a03e68b4121608254787b67 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/class_validators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/class_validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3e42d91dd6037002b7d9bd9ad67a2819829f81a Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/class_validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f1afa24175d93bbd3a5e7a8f9d573ac985e2f308 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/copy_internals.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/copy_internals.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..38b5d3b6a057805c8662df85e2e0afa24370f159 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/copy_internals.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/decorator.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/decorator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf972da0b684842a54fa49cf931052e95c1ea3f8 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/decorator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/json.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/json.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e5a3e71d6fff11b331117f6e4686162fae13e5c Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/json.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/parse.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/parse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8356b10c32cbed860a281377eed585f1bbdbfa7c Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/parse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/tools.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/tools.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e8468f2079273c3408769fc8042fc97a12fd4321 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/deprecated/__pycache__/tools.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/deprecated/class_validators.py b/lib/python3.12/site-packages/pydantic/deprecated/class_validators.py new file mode 100644 index 0000000000000000000000000000000000000000..ad92864c0f5ba6c79d5722f53b02f0932116d531 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/deprecated/class_validators.py @@ -0,0 +1,256 @@ +"""Old `@validator` and `@root_validator` function validators from V1.""" + +from __future__ import annotations as _annotations + +from functools import partial, partialmethod +from types import FunctionType +from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar, Union, overload +from warnings import warn + +from typing_extensions import Protocol, TypeAlias, deprecated + +from .._internal import _decorators, _decorators_v1 +from ..errors import PydanticUserError +from ..warnings import PydanticDeprecatedSince20 + +_ALLOW_REUSE_WARNING_MESSAGE = '`allow_reuse` is deprecated and will be ignored; it should no longer be necessary' + + +if TYPE_CHECKING: + + class _OnlyValueValidatorClsMethod(Protocol): + def __call__(self, __cls: Any, __value: Any) -> Any: ... + + class _V1ValidatorWithValuesClsMethod(Protocol): + def __call__(self, __cls: Any, __value: Any, values: dict[str, Any]) -> Any: ... + + class _V1ValidatorWithValuesKwOnlyClsMethod(Protocol): + def __call__(self, __cls: Any, __value: Any, *, values: dict[str, Any]) -> Any: ... + + class _V1ValidatorWithKwargsClsMethod(Protocol): + def __call__(self, __cls: Any, **kwargs: Any) -> Any: ... + + class _V1ValidatorWithValuesAndKwargsClsMethod(Protocol): + def __call__(self, __cls: Any, values: dict[str, Any], **kwargs: Any) -> Any: ... + + class _V1RootValidatorClsMethod(Protocol): + def __call__( + self, __cls: Any, __values: _decorators_v1.RootValidatorValues + ) -> _decorators_v1.RootValidatorValues: ... + + V1Validator = Union[ + _OnlyValueValidatorClsMethod, + _V1ValidatorWithValuesClsMethod, + _V1ValidatorWithValuesKwOnlyClsMethod, + _V1ValidatorWithKwargsClsMethod, + _V1ValidatorWithValuesAndKwargsClsMethod, + _decorators_v1.V1ValidatorWithValues, + _decorators_v1.V1ValidatorWithValuesKwOnly, + _decorators_v1.V1ValidatorWithKwargs, + _decorators_v1.V1ValidatorWithValuesAndKwargs, + ] + + V1RootValidator = Union[ + _V1RootValidatorClsMethod, + _decorators_v1.V1RootValidatorFunction, + ] + + _PartialClsOrStaticMethod: TypeAlias = Union[classmethod[Any, Any, Any], staticmethod[Any, Any], partialmethod[Any]] + + # Allow both a V1 (assumed pre=False) or V2 (assumed mode='after') validator + # We lie to type checkers and say we return the same thing we get + # but in reality we return a proxy object that _mostly_ behaves like the wrapped thing + _V1ValidatorType = TypeVar('_V1ValidatorType', V1Validator, _PartialClsOrStaticMethod) + _V1RootValidatorFunctionType = TypeVar( + '_V1RootValidatorFunctionType', + _decorators_v1.V1RootValidatorFunction, + _V1RootValidatorClsMethod, + _PartialClsOrStaticMethod, + ) +else: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + + +@deprecated( + 'Pydantic V1 style `@validator` validators are deprecated.' + ' You should migrate to Pydantic V2 style `@field_validator` validators,' + ' see the migration guide for more details', + category=None, +) +def validator( + __field: str, + *fields: str, + pre: bool = False, + each_item: bool = False, + always: bool = False, + check_fields: bool | None = None, + allow_reuse: bool = False, +) -> Callable[[_V1ValidatorType], _V1ValidatorType]: + """Decorate methods on the class indicating that they should be used to validate fields. + + Args: + __field (str): The first field the validator should be called on; this is separate + from `fields` to ensure an error is raised if you don't pass at least one. + *fields (str): Additional field(s) the validator should be called on. + pre (bool, optional): Whether this validator should be called before the standard + validators (else after). Defaults to False. + each_item (bool, optional): For complex objects (sets, lists etc.) whether to validate + individual elements rather than the whole object. Defaults to False. + always (bool, optional): Whether this method and other validators should be called even if + the value is missing. Defaults to False. + check_fields (bool | None, optional): Whether to check that the fields actually exist on the model. + Defaults to None. + allow_reuse (bool, optional): Whether to track and raise an error if another validator refers to + the decorated function. Defaults to False. + + Returns: + Callable: A decorator that can be used to decorate a + function to be used as a validator. + """ + warn( + 'Pydantic V1 style `@validator` validators are deprecated.' + ' You should migrate to Pydantic V2 style `@field_validator` validators,' + ' see the migration guide for more details', + DeprecationWarning, + stacklevel=2, + ) + + if allow_reuse is True: # pragma: no cover + warn(_ALLOW_REUSE_WARNING_MESSAGE, DeprecationWarning) + fields = __field, *fields + if isinstance(fields[0], FunctionType): + raise PydanticUserError( + '`@validator` should be used with fields and keyword arguments, not bare. ' + "E.g. usage should be `@validator('', ...)`", + code='validator-no-fields', + ) + elif not all(isinstance(field, str) for field in fields): + raise PydanticUserError( + '`@validator` fields should be passed as separate string args. ' + "E.g. usage should be `@validator('', '', ...)`", + code='validator-invalid-fields', + ) + + mode: Literal['before', 'after'] = 'before' if pre is True else 'after' + + def dec(f: Any) -> _decorators.PydanticDescriptorProxy[Any]: + if _decorators.is_instance_method_from_sig(f): + raise PydanticUserError( + '`@validator` cannot be applied to instance methods', code='validator-instance-method' + ) + # auto apply the @classmethod decorator + f = _decorators.ensure_classmethod_based_on_signature(f) + wrap = _decorators_v1.make_generic_v1_field_validator + validator_wrapper_info = _decorators.ValidatorDecoratorInfo( + fields=fields, + mode=mode, + each_item=each_item, + always=always, + check_fields=check_fields, + ) + return _decorators.PydanticDescriptorProxy(f, validator_wrapper_info, shim=wrap) + + return dec # type: ignore[return-value] + + +@overload +def root_validator( + *, + # if you don't specify `pre` the default is `pre=False` + # which means you need to specify `skip_on_failure=True` + skip_on_failure: Literal[True], + allow_reuse: bool = ..., +) -> Callable[ + [_V1RootValidatorFunctionType], + _V1RootValidatorFunctionType, +]: ... + + +@overload +def root_validator( + *, + # if you specify `pre=True` then you don't need to specify + # `skip_on_failure`, in fact it is not allowed as an argument! + pre: Literal[True], + allow_reuse: bool = ..., +) -> Callable[ + [_V1RootValidatorFunctionType], + _V1RootValidatorFunctionType, +]: ... + + +@overload +def root_validator( + *, + # if you explicitly specify `pre=False` then you + # MUST specify `skip_on_failure=True` + pre: Literal[False], + skip_on_failure: Literal[True], + allow_reuse: bool = ..., +) -> Callable[ + [_V1RootValidatorFunctionType], + _V1RootValidatorFunctionType, +]: ... + + +@deprecated( + 'Pydantic V1 style `@root_validator` validators are deprecated.' + ' You should migrate to Pydantic V2 style `@model_validator` validators,' + ' see the migration guide for more details', + category=None, +) +def root_validator( + *__args, + pre: bool = False, + skip_on_failure: bool = False, + allow_reuse: bool = False, +) -> Any: + """Decorate methods on a model indicating that they should be used to validate (and perhaps + modify) data either before or after standard model parsing/validation is performed. + + Args: + pre (bool, optional): Whether this validator should be called before the standard + validators (else after). Defaults to False. + skip_on_failure (bool, optional): Whether to stop validation and return as soon as a + failure is encountered. Defaults to False. + allow_reuse (bool, optional): Whether to track and raise an error if another validator + refers to the decorated function. Defaults to False. + + Returns: + Any: A decorator that can be used to decorate a function to be used as a root_validator. + """ + warn( + 'Pydantic V1 style `@root_validator` validators are deprecated.' + ' You should migrate to Pydantic V2 style `@model_validator` validators,' + ' see the migration guide for more details', + DeprecationWarning, + stacklevel=2, + ) + + if __args: + # Ensure a nice error is raised if someone attempts to use the bare decorator + return root_validator()(*__args) # type: ignore + + if allow_reuse is True: # pragma: no cover + warn(_ALLOW_REUSE_WARNING_MESSAGE, DeprecationWarning) + mode: Literal['before', 'after'] = 'before' if pre is True else 'after' + if pre is False and skip_on_failure is not True: + raise PydanticUserError( + 'If you use `@root_validator` with pre=False (the default) you MUST specify `skip_on_failure=True`.' + ' Note that `@root_validator` is deprecated and should be replaced with `@model_validator`.', + code='root-validator-pre-skip', + ) + + wrap = partial(_decorators_v1.make_v1_generic_root_validator, pre=pre) + + def dec(f: Callable[..., Any] | classmethod[Any, Any, Any] | staticmethod[Any, Any]) -> Any: + if _decorators.is_instance_method_from_sig(f): + raise TypeError('`@root_validator` cannot be applied to instance methods') + # auto apply the @classmethod decorator + res = _decorators.ensure_classmethod_based_on_signature(f) + dec_info = _decorators.RootValidatorDecoratorInfo(mode=mode) + return _decorators.PydanticDescriptorProxy(res, dec_info, shim=wrap) + + return dec diff --git a/lib/python3.12/site-packages/pydantic/deprecated/config.py b/lib/python3.12/site-packages/pydantic/deprecated/config.py new file mode 100644 index 0000000000000000000000000000000000000000..bd4692ace9b4594799974d6ecf8d44357d1d02c5 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/deprecated/config.py @@ -0,0 +1,72 @@ +from __future__ import annotations as _annotations + +import warnings +from typing import TYPE_CHECKING, Any, Literal + +from typing_extensions import deprecated + +from .._internal import _config +from ..warnings import PydanticDeprecatedSince20 + +if not TYPE_CHECKING: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + +__all__ = 'BaseConfig', 'Extra' + + +class _ConfigMetaclass(type): + def __getattr__(self, item: str) -> Any: + try: + obj = _config.config_defaults[item] + warnings.warn(_config.DEPRECATION_MESSAGE, DeprecationWarning) + return obj + except KeyError as exc: + raise AttributeError(f"type object '{self.__name__}' has no attribute {exc}") from exc + + +@deprecated('BaseConfig is deprecated. Use the `pydantic.ConfigDict` instead.', category=PydanticDeprecatedSince20) +class BaseConfig(metaclass=_ConfigMetaclass): + """This class is only retained for backwards compatibility. + + !!! Warning "Deprecated" + BaseConfig is deprecated. Use the [`pydantic.ConfigDict`][pydantic.ConfigDict] instead. + """ + + def __getattr__(self, item: str) -> Any: + try: + obj = super().__getattribute__(item) + warnings.warn(_config.DEPRECATION_MESSAGE, DeprecationWarning) + return obj + except AttributeError as exc: + try: + return getattr(type(self), item) + except AttributeError: + # re-raising changes the displayed text to reflect that `self` is not a type + raise AttributeError(str(exc)) from exc + + def __init_subclass__(cls, **kwargs: Any) -> None: + warnings.warn(_config.DEPRECATION_MESSAGE, DeprecationWarning) + return super().__init_subclass__(**kwargs) + + +class _ExtraMeta(type): + def __getattribute__(self, __name: str) -> Any: + # The @deprecated decorator accesses other attributes, so we only emit a warning for the expected ones + if __name in {'allow', 'ignore', 'forbid'}: + warnings.warn( + "`pydantic.config.Extra` is deprecated, use literal values instead (e.g. `extra='allow'`)", + DeprecationWarning, + stacklevel=2, + ) + return super().__getattribute__(__name) + + +@deprecated( + "Extra is deprecated. Use literal values instead (e.g. `extra='allow'`)", category=PydanticDeprecatedSince20 +) +class Extra(metaclass=_ExtraMeta): + allow: Literal['allow'] = 'allow' + ignore: Literal['ignore'] = 'ignore' + forbid: Literal['forbid'] = 'forbid' diff --git a/lib/python3.12/site-packages/pydantic/deprecated/copy_internals.py b/lib/python3.12/site-packages/pydantic/deprecated/copy_internals.py new file mode 100644 index 0000000000000000000000000000000000000000..1b0dc125f196fbdb855bc0f45724a71355cedd51 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/deprecated/copy_internals.py @@ -0,0 +1,224 @@ +from __future__ import annotations as _annotations + +import typing +from copy import deepcopy +from enum import Enum +from typing import Any + +import typing_extensions + +from .._internal import ( + _model_construction, + _typing_extra, + _utils, +) + +if typing.TYPE_CHECKING: + from .. import BaseModel + from .._internal._utils import AbstractSetIntStr, MappingIntStrAny + + AnyClassMethod = classmethod[Any, Any, Any] + TupleGenerator = typing.Generator[tuple[str, Any], None, None] + Model = typing.TypeVar('Model', bound='BaseModel') + # should be `set[int] | set[str] | dict[int, IncEx] | dict[str, IncEx] | None`, but mypy can't cope + IncEx: typing_extensions.TypeAlias = 'set[int] | set[str] | dict[int, Any] | dict[str, Any] | None' + +_object_setattr = _model_construction.object_setattr + + +def _iter( + self: BaseModel, + to_dict: bool = False, + by_alias: bool = False, + include: AbstractSetIntStr | MappingIntStrAny | None = None, + exclude: AbstractSetIntStr | MappingIntStrAny | None = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, +) -> TupleGenerator: + # Merge field set excludes with explicit exclude parameter with explicit overriding field set options. + # The extra "is not None" guards are not logically necessary but optimizes performance for the simple case. + if exclude is not None: + exclude = _utils.ValueItems.merge( + {k: v.exclude for k, v in self.__pydantic_fields__.items() if v.exclude is not None}, exclude + ) + + if include is not None: + include = _utils.ValueItems.merge({k: True for k in self.__pydantic_fields__}, include, intersect=True) + + allowed_keys = _calculate_keys(self, include=include, exclude=exclude, exclude_unset=exclude_unset) # type: ignore + if allowed_keys is None and not (to_dict or by_alias or exclude_unset or exclude_defaults or exclude_none): + # huge boost for plain _iter() + yield from self.__dict__.items() + if self.__pydantic_extra__: + yield from self.__pydantic_extra__.items() + return + + value_exclude = _utils.ValueItems(self, exclude) if exclude is not None else None + value_include = _utils.ValueItems(self, include) if include is not None else None + + if self.__pydantic_extra__ is None: + items = self.__dict__.items() + else: + items = list(self.__dict__.items()) + list(self.__pydantic_extra__.items()) + + for field_key, v in items: + if (allowed_keys is not None and field_key not in allowed_keys) or (exclude_none and v is None): + continue + + if exclude_defaults: + try: + field = self.__pydantic_fields__[field_key] + except KeyError: + pass + else: + if not field.is_required() and field.default == v: + continue + + if by_alias and field_key in self.__pydantic_fields__: + dict_key = self.__pydantic_fields__[field_key].alias or field_key + else: + dict_key = field_key + + if to_dict or value_include or value_exclude: + v = _get_value( + type(self), + v, + to_dict=to_dict, + by_alias=by_alias, + include=value_include and value_include.for_element(field_key), + exclude=value_exclude and value_exclude.for_element(field_key), + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + ) + yield dict_key, v + + +def _copy_and_set_values( + self: Model, + values: dict[str, Any], + fields_set: set[str], + extra: dict[str, Any] | None = None, + private: dict[str, Any] | None = None, + *, + deep: bool, # UP006 +) -> Model: + if deep: + # chances of having empty dict here are quite low for using smart_deepcopy + values = deepcopy(values) + extra = deepcopy(extra) + private = deepcopy(private) + + cls = self.__class__ + m = cls.__new__(cls) + _object_setattr(m, '__dict__', values) + _object_setattr(m, '__pydantic_extra__', extra) + _object_setattr(m, '__pydantic_fields_set__', fields_set) + _object_setattr(m, '__pydantic_private__', private) + + return m + + +@typing.no_type_check +def _get_value( + cls: type[BaseModel], + v: Any, + to_dict: bool, + by_alias: bool, + include: AbstractSetIntStr | MappingIntStrAny | None, + exclude: AbstractSetIntStr | MappingIntStrAny | None, + exclude_unset: bool, + exclude_defaults: bool, + exclude_none: bool, +) -> Any: + from .. import BaseModel + + if isinstance(v, BaseModel): + if to_dict: + return v.model_dump( + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + include=include, # type: ignore + exclude=exclude, # type: ignore + exclude_none=exclude_none, + ) + else: + return v.copy(include=include, exclude=exclude) + + value_exclude = _utils.ValueItems(v, exclude) if exclude else None + value_include = _utils.ValueItems(v, include) if include else None + + if isinstance(v, dict): + return { + k_: _get_value( + cls, + v_, + to_dict=to_dict, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + include=value_include and value_include.for_element(k_), + exclude=value_exclude and value_exclude.for_element(k_), + exclude_none=exclude_none, + ) + for k_, v_ in v.items() + if (not value_exclude or not value_exclude.is_excluded(k_)) + and (not value_include or value_include.is_included(k_)) + } + + elif _utils.sequence_like(v): + seq_args = ( + _get_value( + cls, + v_, + to_dict=to_dict, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + include=value_include and value_include.for_element(i), + exclude=value_exclude and value_exclude.for_element(i), + exclude_none=exclude_none, + ) + for i, v_ in enumerate(v) + if (not value_exclude or not value_exclude.is_excluded(i)) + and (not value_include or value_include.is_included(i)) + ) + + return v.__class__(*seq_args) if _typing_extra.is_namedtuple(v.__class__) else v.__class__(seq_args) + + elif isinstance(v, Enum) and getattr(cls.model_config, 'use_enum_values', False): + return v.value + + else: + return v + + +def _calculate_keys( + self: BaseModel, + include: MappingIntStrAny | None, + exclude: MappingIntStrAny | None, + exclude_unset: bool, + update: dict[str, Any] | None = None, # noqa UP006 +) -> typing.AbstractSet[str] | None: + if include is None and exclude is None and exclude_unset is False: + return None + + keys: typing.AbstractSet[str] + if exclude_unset: + keys = self.__pydantic_fields_set__.copy() + else: + keys = set(self.__dict__.keys()) + keys = keys | (self.__pydantic_extra__ or {}).keys() + + if include is not None: + keys &= include.keys() + + if update: + keys -= update.keys() + + if exclude: + keys -= {k for k, v in exclude.items() if _utils.ValueItems.is_true(v)} + + return keys diff --git a/lib/python3.12/site-packages/pydantic/deprecated/decorator.py b/lib/python3.12/site-packages/pydantic/deprecated/decorator.py new file mode 100644 index 0000000000000000000000000000000000000000..e73ad209ae86b26d6fa892f1eba39cfc28980970 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/deprecated/decorator.py @@ -0,0 +1,284 @@ +import warnings +from collections.abc import Mapping +from functools import wraps +from typing import TYPE_CHECKING, Any, Callable, Optional, TypeVar, Union, overload + +from typing_extensions import deprecated + +from .._internal import _config, _typing_extra +from ..alias_generators import to_pascal +from ..errors import PydanticUserError +from ..functional_validators import field_validator +from ..main import BaseModel, create_model +from ..warnings import PydanticDeprecatedSince20 + +if not TYPE_CHECKING: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + +__all__ = ('validate_arguments',) + +if TYPE_CHECKING: + AnyCallable = Callable[..., Any] + + AnyCallableT = TypeVar('AnyCallableT', bound=AnyCallable) + ConfigType = Union[None, type[Any], dict[str, Any]] + + +@overload +def validate_arguments( + func: None = None, *, config: 'ConfigType' = None +) -> Callable[['AnyCallableT'], 'AnyCallableT']: ... + + +@overload +def validate_arguments(func: 'AnyCallableT') -> 'AnyCallableT': ... + + +@deprecated( + 'The `validate_arguments` method is deprecated; use `validate_call` instead.', + category=None, +) +def validate_arguments(func: Optional['AnyCallableT'] = None, *, config: 'ConfigType' = None) -> Any: + """Decorator to validate the arguments passed to a function.""" + warnings.warn( + 'The `validate_arguments` method is deprecated; use `validate_call` instead.', + PydanticDeprecatedSince20, + stacklevel=2, + ) + + def validate(_func: 'AnyCallable') -> 'AnyCallable': + vd = ValidatedFunction(_func, config) + + @wraps(_func) + def wrapper_function(*args: Any, **kwargs: Any) -> Any: + return vd.call(*args, **kwargs) + + wrapper_function.vd = vd # type: ignore + wrapper_function.validate = vd.init_model_instance # type: ignore + wrapper_function.raw_function = vd.raw_function # type: ignore + wrapper_function.model = vd.model # type: ignore + return wrapper_function + + if func: + return validate(func) + else: + return validate + + +ALT_V_ARGS = 'v__args' +ALT_V_KWARGS = 'v__kwargs' +V_POSITIONAL_ONLY_NAME = 'v__positional_only' +V_DUPLICATE_KWARGS = 'v__duplicate_kwargs' + + +class ValidatedFunction: + def __init__(self, function: 'AnyCallable', config: 'ConfigType'): + from inspect import Parameter, signature + + parameters: Mapping[str, Parameter] = signature(function).parameters + + if parameters.keys() & {ALT_V_ARGS, ALT_V_KWARGS, V_POSITIONAL_ONLY_NAME, V_DUPLICATE_KWARGS}: + raise PydanticUserError( + f'"{ALT_V_ARGS}", "{ALT_V_KWARGS}", "{V_POSITIONAL_ONLY_NAME}" and "{V_DUPLICATE_KWARGS}" ' + f'are not permitted as argument names when using the "{validate_arguments.__name__}" decorator', + code=None, + ) + + self.raw_function = function + self.arg_mapping: dict[int, str] = {} + self.positional_only_args: set[str] = set() + self.v_args_name = 'args' + self.v_kwargs_name = 'kwargs' + + type_hints = _typing_extra.get_type_hints(function, include_extras=True) + takes_args = False + takes_kwargs = False + fields: dict[str, tuple[Any, Any]] = {} + for i, (name, p) in enumerate(parameters.items()): + if p.annotation is p.empty: + annotation = Any + else: + annotation = type_hints[name] + + default = ... if p.default is p.empty else p.default + if p.kind == Parameter.POSITIONAL_ONLY: + self.arg_mapping[i] = name + fields[name] = annotation, default + fields[V_POSITIONAL_ONLY_NAME] = list[str], None + self.positional_only_args.add(name) + elif p.kind == Parameter.POSITIONAL_OR_KEYWORD: + self.arg_mapping[i] = name + fields[name] = annotation, default + fields[V_DUPLICATE_KWARGS] = list[str], None + elif p.kind == Parameter.KEYWORD_ONLY: + fields[name] = annotation, default + elif p.kind == Parameter.VAR_POSITIONAL: + self.v_args_name = name + fields[name] = tuple[annotation, ...], None + takes_args = True + else: + assert p.kind == Parameter.VAR_KEYWORD, p.kind + self.v_kwargs_name = name + fields[name] = dict[str, annotation], None + takes_kwargs = True + + # these checks avoid a clash between "args" and a field with that name + if not takes_args and self.v_args_name in fields: + self.v_args_name = ALT_V_ARGS + + # same with "kwargs" + if not takes_kwargs and self.v_kwargs_name in fields: + self.v_kwargs_name = ALT_V_KWARGS + + if not takes_args: + # we add the field so validation below can raise the correct exception + fields[self.v_args_name] = list[Any], None + + if not takes_kwargs: + # same with kwargs + fields[self.v_kwargs_name] = dict[Any, Any], None + + self.create_model(fields, takes_args, takes_kwargs, config) + + def init_model_instance(self, *args: Any, **kwargs: Any) -> BaseModel: + values = self.build_values(args, kwargs) + return self.model(**values) + + def call(self, *args: Any, **kwargs: Any) -> Any: + m = self.init_model_instance(*args, **kwargs) + return self.execute(m) + + def build_values(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]: + values: dict[str, Any] = {} + if args: + arg_iter = enumerate(args) + while True: + try: + i, a = next(arg_iter) + except StopIteration: + break + arg_name = self.arg_mapping.get(i) + if arg_name is not None: + values[arg_name] = a + else: + values[self.v_args_name] = [a] + [a for _, a in arg_iter] + break + + var_kwargs: dict[str, Any] = {} + wrong_positional_args = [] + duplicate_kwargs = [] + fields_alias = [ + field.alias + for name, field in self.model.__pydantic_fields__.items() + if name not in (self.v_args_name, self.v_kwargs_name) + ] + non_var_fields = set(self.model.__pydantic_fields__) - {self.v_args_name, self.v_kwargs_name} + for k, v in kwargs.items(): + if k in non_var_fields or k in fields_alias: + if k in self.positional_only_args: + wrong_positional_args.append(k) + if k in values: + duplicate_kwargs.append(k) + values[k] = v + else: + var_kwargs[k] = v + + if var_kwargs: + values[self.v_kwargs_name] = var_kwargs + if wrong_positional_args: + values[V_POSITIONAL_ONLY_NAME] = wrong_positional_args + if duplicate_kwargs: + values[V_DUPLICATE_KWARGS] = duplicate_kwargs + return values + + def execute(self, m: BaseModel) -> Any: + d = { + k: v + for k, v in m.__dict__.items() + if k in m.__pydantic_fields_set__ or m.__pydantic_fields__[k].default_factory + } + var_kwargs = d.pop(self.v_kwargs_name, {}) + + if self.v_args_name in d: + args_: list[Any] = [] + in_kwargs = False + kwargs = {} + for name, value in d.items(): + if in_kwargs: + kwargs[name] = value + elif name == self.v_args_name: + args_ += value + in_kwargs = True + else: + args_.append(value) + return self.raw_function(*args_, **kwargs, **var_kwargs) + elif self.positional_only_args: + args_ = [] + kwargs = {} + for name, value in d.items(): + if name in self.positional_only_args: + args_.append(value) + else: + kwargs[name] = value + return self.raw_function(*args_, **kwargs, **var_kwargs) + else: + return self.raw_function(**d, **var_kwargs) + + def create_model(self, fields: dict[str, Any], takes_args: bool, takes_kwargs: bool, config: 'ConfigType') -> None: + pos_args = len(self.arg_mapping) + + config_wrapper = _config.ConfigWrapper(config) + + if config_wrapper.alias_generator: + raise PydanticUserError( + 'Setting the "alias_generator" property on custom Config for ' + '@validate_arguments is not yet supported, please remove.', + code=None, + ) + if config_wrapper.extra is None: + config_wrapper.config_dict['extra'] = 'forbid' + + class DecoratorBaseModel(BaseModel): + @field_validator(self.v_args_name, check_fields=False) + @classmethod + def check_args(cls, v: Optional[list[Any]]) -> Optional[list[Any]]: + if takes_args or v is None: + return v + + raise TypeError(f'{pos_args} positional arguments expected but {pos_args + len(v)} given') + + @field_validator(self.v_kwargs_name, check_fields=False) + @classmethod + def check_kwargs(cls, v: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]: + if takes_kwargs or v is None: + return v + + plural = '' if len(v) == 1 else 's' + keys = ', '.join(map(repr, v.keys())) + raise TypeError(f'unexpected keyword argument{plural}: {keys}') + + @field_validator(V_POSITIONAL_ONLY_NAME, check_fields=False) + @classmethod + def check_positional_only(cls, v: Optional[list[str]]) -> None: + if v is None: + return + + plural = '' if len(v) == 1 else 's' + keys = ', '.join(map(repr, v)) + raise TypeError(f'positional-only argument{plural} passed as keyword argument{plural}: {keys}') + + @field_validator(V_DUPLICATE_KWARGS, check_fields=False) + @classmethod + def check_duplicate_kwargs(cls, v: Optional[list[str]]) -> None: + if v is None: + return + + plural = '' if len(v) == 1 else 's' + keys = ', '.join(map(repr, v)) + raise TypeError(f'multiple values for argument{plural}: {keys}') + + model_config = config_wrapper.config_dict + + self.model = create_model(to_pascal(self.raw_function.__name__), __base__=DecoratorBaseModel, **fields) diff --git a/lib/python3.12/site-packages/pydantic/deprecated/json.py b/lib/python3.12/site-packages/pydantic/deprecated/json.py new file mode 100644 index 0000000000000000000000000000000000000000..1e216a765d15d5dd9e379f9de4f8f91ab8063877 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/deprecated/json.py @@ -0,0 +1,141 @@ +import datetime +import warnings +from collections import deque +from decimal import Decimal +from enum import Enum +from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network +from pathlib import Path +from re import Pattern +from types import GeneratorType +from typing import TYPE_CHECKING, Any, Callable, Union +from uuid import UUID + +from typing_extensions import deprecated + +from .._internal._import_utils import import_cached_base_model +from ..color import Color +from ..networks import NameEmail +from ..types import SecretBytes, SecretStr +from ..warnings import PydanticDeprecatedSince20 + +if not TYPE_CHECKING: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + +__all__ = 'pydantic_encoder', 'custom_pydantic_encoder', 'timedelta_isoformat' + + +def isoformat(o: Union[datetime.date, datetime.time]) -> str: + return o.isoformat() + + +def decimal_encoder(dec_value: Decimal) -> Union[int, float]: + """Encodes a Decimal as int of there's no exponent, otherwise float. + + This is useful when we use ConstrainedDecimal to represent Numeric(x,0) + where a integer (but not int typed) is used. Encoding this as a float + results in failed round-tripping between encode and parse. + Our Id type is a prime example of this. + + >>> decimal_encoder(Decimal("1.0")) + 1.0 + + >>> decimal_encoder(Decimal("1")) + 1 + """ + exponent = dec_value.as_tuple().exponent + if isinstance(exponent, int) and exponent >= 0: + return int(dec_value) + else: + return float(dec_value) + + +ENCODERS_BY_TYPE: dict[type[Any], Callable[[Any], Any]] = { + bytes: lambda o: o.decode(), + Color: str, + datetime.date: isoformat, + datetime.datetime: isoformat, + datetime.time: isoformat, + datetime.timedelta: lambda td: td.total_seconds(), + Decimal: decimal_encoder, + Enum: lambda o: o.value, + frozenset: list, + deque: list, + GeneratorType: list, + IPv4Address: str, + IPv4Interface: str, + IPv4Network: str, + IPv6Address: str, + IPv6Interface: str, + IPv6Network: str, + NameEmail: str, + Path: str, + Pattern: lambda o: o.pattern, + SecretBytes: str, + SecretStr: str, + set: list, + UUID: str, +} + + +@deprecated( + '`pydantic_encoder` is deprecated, use `pydantic_core.to_jsonable_python` instead.', + category=None, +) +def pydantic_encoder(obj: Any) -> Any: + warnings.warn( + '`pydantic_encoder` is deprecated, use `pydantic_core.to_jsonable_python` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + from dataclasses import asdict, is_dataclass + + BaseModel = import_cached_base_model() + + if isinstance(obj, BaseModel): + return obj.model_dump() + elif is_dataclass(obj): + return asdict(obj) # type: ignore + + # Check the class type and its superclasses for a matching encoder + for base in obj.__class__.__mro__[:-1]: + try: + encoder = ENCODERS_BY_TYPE[base] + except KeyError: + continue + return encoder(obj) + else: # We have exited the for loop without finding a suitable encoder + raise TypeError(f"Object of type '{obj.__class__.__name__}' is not JSON serializable") + + +# TODO: Add a suggested migration path once there is a way to use custom encoders +@deprecated( + '`custom_pydantic_encoder` is deprecated, use `BaseModel.model_dump` instead.', + category=None, +) +def custom_pydantic_encoder(type_encoders: dict[Any, Callable[[type[Any]], Any]], obj: Any) -> Any: + warnings.warn( + '`custom_pydantic_encoder` is deprecated, use `BaseModel.model_dump` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + # Check the class type and its superclasses for a matching encoder + for base in obj.__class__.__mro__[:-1]: + try: + encoder = type_encoders[base] + except KeyError: + continue + + return encoder(obj) + else: # We have exited the for loop without finding a suitable encoder + return pydantic_encoder(obj) + + +@deprecated('`timedelta_isoformat` is deprecated.', category=None) +def timedelta_isoformat(td: datetime.timedelta) -> str: + """ISO 8601 encoding for Python timedelta object.""" + warnings.warn('`timedelta_isoformat` is deprecated.', category=PydanticDeprecatedSince20, stacklevel=2) + minutes, seconds = divmod(td.seconds, 60) + hours, minutes = divmod(minutes, 60) + return f'{"-" if td.days < 0 else ""}P{abs(td.days)}DT{hours:d}H{minutes:d}M{seconds:d}.{td.microseconds:06d}S' diff --git a/lib/python3.12/site-packages/pydantic/deprecated/parse.py b/lib/python3.12/site-packages/pydantic/deprecated/parse.py new file mode 100644 index 0000000000000000000000000000000000000000..2a92e62b7b2e8dab77cbe0c2dbb79c810af7f452 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/deprecated/parse.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import json +import pickle +import warnings +from enum import Enum +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable + +from typing_extensions import deprecated + +from ..warnings import PydanticDeprecatedSince20 + +if not TYPE_CHECKING: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + + +class Protocol(str, Enum): + json = 'json' + pickle = 'pickle' + + +@deprecated('`load_str_bytes` is deprecated.', category=None) +def load_str_bytes( + b: str | bytes, + *, + content_type: str | None = None, + encoding: str = 'utf8', + proto: Protocol | None = None, + allow_pickle: bool = False, + json_loads: Callable[[str], Any] = json.loads, +) -> Any: + warnings.warn('`load_str_bytes` is deprecated.', category=PydanticDeprecatedSince20, stacklevel=2) + if proto is None and content_type: + if content_type.endswith(('json', 'javascript')): + pass + elif allow_pickle and content_type.endswith('pickle'): + proto = Protocol.pickle + else: + raise TypeError(f'Unknown content-type: {content_type}') + + proto = proto or Protocol.json + + if proto == Protocol.json: + if isinstance(b, bytes): + b = b.decode(encoding) + return json_loads(b) # type: ignore + elif proto == Protocol.pickle: + if not allow_pickle: + raise RuntimeError('Trying to decode with pickle with allow_pickle=False') + bb = b if isinstance(b, bytes) else b.encode() # type: ignore + return pickle.loads(bb) + else: + raise TypeError(f'Unknown protocol: {proto}') + + +@deprecated('`load_file` is deprecated.', category=None) +def load_file( + path: str | Path, + *, + content_type: str | None = None, + encoding: str = 'utf8', + proto: Protocol | None = None, + allow_pickle: bool = False, + json_loads: Callable[[str], Any] = json.loads, +) -> Any: + warnings.warn('`load_file` is deprecated.', category=PydanticDeprecatedSince20, stacklevel=2) + path = Path(path) + b = path.read_bytes() + if content_type is None: + if path.suffix in ('.js', '.json'): + proto = Protocol.json + elif path.suffix == '.pkl': + proto = Protocol.pickle + + return load_str_bytes( + b, proto=proto, content_type=content_type, encoding=encoding, allow_pickle=allow_pickle, json_loads=json_loads + ) diff --git a/lib/python3.12/site-packages/pydantic/deprecated/tools.py b/lib/python3.12/site-packages/pydantic/deprecated/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..5ad7faef2ec69eaa737a58d3587770a612a510d8 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/deprecated/tools.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import json +import warnings +from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union + +from typing_extensions import deprecated + +from ..json_schema import DEFAULT_REF_TEMPLATE, GenerateJsonSchema +from ..type_adapter import TypeAdapter +from ..warnings import PydanticDeprecatedSince20 + +if not TYPE_CHECKING: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + +__all__ = 'parse_obj_as', 'schema_of', 'schema_json_of' + +NameFactory = Union[str, Callable[[type[Any]], str]] + + +T = TypeVar('T') + + +@deprecated( + '`parse_obj_as` is deprecated. Use `pydantic.TypeAdapter.validate_python` instead.', + category=None, +) +def parse_obj_as(type_: type[T], obj: Any, type_name: NameFactory | None = None) -> T: + warnings.warn( + '`parse_obj_as` is deprecated. Use `pydantic.TypeAdapter.validate_python` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + if type_name is not None: # pragma: no cover + warnings.warn( + 'The type_name parameter is deprecated. parse_obj_as no longer creates temporary models', + DeprecationWarning, + stacklevel=2, + ) + return TypeAdapter(type_).validate_python(obj) + + +@deprecated( + '`schema_of` is deprecated. Use `pydantic.TypeAdapter.json_schema` instead.', + category=None, +) +def schema_of( + type_: Any, + *, + title: NameFactory | None = None, + by_alias: bool = True, + ref_template: str = DEFAULT_REF_TEMPLATE, + schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, +) -> dict[str, Any]: + """Generate a JSON schema (as dict) for the passed model or dynamically generated one.""" + warnings.warn( + '`schema_of` is deprecated. Use `pydantic.TypeAdapter.json_schema` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + res = TypeAdapter(type_).json_schema( + by_alias=by_alias, + schema_generator=schema_generator, + ref_template=ref_template, + ) + if title is not None: + if isinstance(title, str): + res['title'] = title + else: + warnings.warn( + 'Passing a callable for the `title` parameter is deprecated and no longer supported', + DeprecationWarning, + stacklevel=2, + ) + res['title'] = title(type_) + return res + + +@deprecated( + '`schema_json_of` is deprecated. Use `pydantic.TypeAdapter.json_schema` instead.', + category=None, +) +def schema_json_of( + type_: Any, + *, + title: NameFactory | None = None, + by_alias: bool = True, + ref_template: str = DEFAULT_REF_TEMPLATE, + schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, + **dumps_kwargs: Any, +) -> str: + """Generate a JSON schema (as JSON) for the passed model or dynamically generated one.""" + warnings.warn( + '`schema_json_of` is deprecated. Use `pydantic.TypeAdapter.json_schema` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + return json.dumps( + schema_of(type_, title=title, by_alias=by_alias, ref_template=ref_template, schema_generator=schema_generator), + **dumps_kwargs, + ) diff --git a/lib/python3.12/site-packages/pydantic/env_settings.py b/lib/python3.12/site-packages/pydantic/env_settings.py new file mode 100644 index 0000000000000000000000000000000000000000..cd0b04e6a43c1bb3767dd136a19a9873fa547514 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/env_settings.py @@ -0,0 +1,5 @@ +"""The `env_settings` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/error_wrappers.py b/lib/python3.12/site-packages/pydantic/error_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..2985419abf23f5e529af43883f1d365452e4190a --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/error_wrappers.py @@ -0,0 +1,5 @@ +"""The `error_wrappers` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/errors.py b/lib/python3.12/site-packages/pydantic/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..f22706822f3fa542f69810f8e497857d5be36a65 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/errors.py @@ -0,0 +1,189 @@ +"""Pydantic-specific errors.""" + +from __future__ import annotations as _annotations + +import re +from typing import Any, ClassVar, Literal + +from typing_extensions import Self +from typing_inspection.introspection import Qualifier + +from pydantic._internal import _repr + +from ._migration import getattr_migration +from .version import version_short + +__all__ = ( + 'PydanticUserError', + 'PydanticUndefinedAnnotation', + 'PydanticImportError', + 'PydanticSchemaGenerationError', + 'PydanticInvalidForJsonSchema', + 'PydanticForbiddenQualifier', + 'PydanticErrorCodes', +) + +# We use this URL to allow for future flexibility about how we host the docs, while allowing for Pydantic +# code in the while with "old" URLs to still work. +# 'u' refers to "user errors" - e.g. errors caused by developers using pydantic, as opposed to validation errors. +DEV_ERROR_DOCS_URL = f'https://errors.pydantic.dev/{version_short()}/u/' +PydanticErrorCodes = Literal[ + 'class-not-fully-defined', + 'custom-json-schema', + 'decorator-missing-field', + 'discriminator-no-field', + 'discriminator-alias-type', + 'discriminator-needs-literal', + 'discriminator-alias', + 'discriminator-validator', + 'callable-discriminator-no-tag', + 'typed-dict-version', + 'model-field-overridden', + 'model-field-missing-annotation', + 'config-both', + 'removed-kwargs', + 'circular-reference-schema', + 'invalid-for-json-schema', + 'json-schema-already-used', + 'base-model-instantiated', + 'undefined-annotation', + 'schema-for-unknown-type', + 'import-error', + 'create-model-field-definitions', + 'validator-no-fields', + 'validator-invalid-fields', + 'validator-instance-method', + 'validator-input-type', + 'root-validator-pre-skip', + 'model-serializer-instance-method', + 'validator-field-config-info', + 'validator-v1-signature', + 'validator-signature', + 'field-serializer-signature', + 'model-serializer-signature', + 'multiple-field-serializers', + 'invalid-annotated-type', + 'type-adapter-config-unused', + 'root-model-extra', + 'unevaluable-type-annotation', + 'dataclass-init-false-extra-allow', + 'clashing-init-and-init-var', + 'model-config-invalid-field-name', + 'with-config-on-model', + 'dataclass-on-model', + 'validate-call-type', + 'unpack-typed-dict', + 'overlapping-unpack-typed-dict', + 'invalid-self-type', + 'validate-by-alias-and-name-false', +] + + +class PydanticErrorMixin: + """A mixin class for common functionality shared by all Pydantic-specific errors. + + Attributes: + message: A message describing the error. + code: An optional error code from PydanticErrorCodes enum. + """ + + def __init__(self, message: str, *, code: PydanticErrorCodes | None) -> None: + self.message = message + self.code = code + + def __str__(self) -> str: + if self.code is None: + return self.message + else: + return f'{self.message}\n\nFor further information visit {DEV_ERROR_DOCS_URL}{self.code}' + + +class PydanticUserError(PydanticErrorMixin, TypeError): + """An error raised due to incorrect use of Pydantic.""" + + +class PydanticUndefinedAnnotation(PydanticErrorMixin, NameError): + """A subclass of `NameError` raised when handling undefined annotations during `CoreSchema` generation. + + Attributes: + name: Name of the error. + message: Description of the error. + """ + + def __init__(self, name: str, message: str) -> None: + self.name = name + super().__init__(message=message, code='undefined-annotation') + + @classmethod + def from_name_error(cls, name_error: NameError) -> Self: + """Convert a `NameError` to a `PydanticUndefinedAnnotation` error. + + Args: + name_error: `NameError` to be converted. + + Returns: + Converted `PydanticUndefinedAnnotation` error. + """ + try: + name = name_error.name # type: ignore # python > 3.10 + except AttributeError: + name = re.search(r".*'(.+?)'", str(name_error)).group(1) # type: ignore[union-attr] + return cls(name=name, message=str(name_error)) + + +class PydanticImportError(PydanticErrorMixin, ImportError): + """An error raised when an import fails due to module changes between V1 and V2. + + Attributes: + message: Description of the error. + """ + + def __init__(self, message: str) -> None: + super().__init__(message, code='import-error') + + +class PydanticSchemaGenerationError(PydanticUserError): + """An error raised during failures to generate a `CoreSchema` for some type. + + Attributes: + message: Description of the error. + """ + + def __init__(self, message: str) -> None: + super().__init__(message, code='schema-for-unknown-type') + + +class PydanticInvalidForJsonSchema(PydanticUserError): + """An error raised during failures to generate a JSON schema for some `CoreSchema`. + + Attributes: + message: Description of the error. + """ + + def __init__(self, message: str) -> None: + super().__init__(message, code='invalid-for-json-schema') + + +class PydanticForbiddenQualifier(PydanticUserError): + """An error raised if a forbidden type qualifier is found in a type annotation.""" + + _qualifier_repr_map: ClassVar[dict[Qualifier, str]] = { + 'required': 'typing.Required', + 'not_required': 'typing.NotRequired', + 'read_only': 'typing.ReadOnly', + 'class_var': 'typing.ClassVar', + 'init_var': 'dataclasses.InitVar', + 'final': 'typing.Final', + } + + def __init__(self, qualifier: Qualifier, annotation: Any) -> None: + super().__init__( + message=( + f'The annotation {_repr.display_as_type(annotation)!r} contains the {self._qualifier_repr_map[qualifier]!r} ' + f'type qualifier, which is invalid in the context it is defined.' + ), + code=None, + ) + + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/experimental/__init__.py b/lib/python3.12/site-packages/pydantic/experimental/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4aa58c6a61ca7ac05e399e40ae80e31032aeeb4e --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/experimental/__init__.py @@ -0,0 +1,10 @@ +"""The "experimental" module of pydantic contains potential new features that are subject to change.""" + +import warnings + +from pydantic.warnings import PydanticExperimentalWarning + +warnings.warn( + 'This module is experimental, its contents are subject to change and deprecation.', + category=PydanticExperimentalWarning, +) diff --git a/lib/python3.12/site-packages/pydantic/experimental/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/experimental/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42cd7bcf2ac390336b51f19e1ed3cedd1ffaa818 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/experimental/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/experimental/__pycache__/arguments_schema.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/experimental/__pycache__/arguments_schema.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb724c562d47b2a02687b23de1dcdffabc7e6449 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/experimental/__pycache__/arguments_schema.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/experimental/__pycache__/pipeline.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/experimental/__pycache__/pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14ca1253d5f3c61fd915bf9387d5dff69e093a97 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/experimental/__pycache__/pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/experimental/arguments_schema.py b/lib/python3.12/site-packages/pydantic/experimental/arguments_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..af4a8f3be3242519b529668e321dd8393241f7ea --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/experimental/arguments_schema.py @@ -0,0 +1,44 @@ +"""Experimental module exposing a function to generate a core schema that validates callable arguments.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, Literal + +from pydantic_core import CoreSchema + +from pydantic import ConfigDict +from pydantic._internal import _config, _generate_schema, _namespace_utils + + +def generate_arguments_schema( + func: Callable[..., Any], + schema_type: Literal['arguments', 'arguments-v3'] = 'arguments-v3', + parameters_callback: Callable[[int, str, Any], Literal['skip'] | None] | None = None, + config: ConfigDict | None = None, +) -> CoreSchema: + """Generate the schema for the arguments of a function. + + Args: + func: The function to generate the schema for. + schema_type: The type of schema to generate. + parameters_callback: A callable that will be invoked for each parameter. The callback + should take three required arguments: the index, the name and the type annotation + (or [`Parameter.empty`][inspect.Parameter.empty] if not annotated) of the parameter. + The callback can optionally return `'skip'`, so that the parameter gets excluded + from the resulting schema. + config: The configuration to use. + + Returns: + The generated schema. + """ + generate_schema = _generate_schema.GenerateSchema( + _config.ConfigWrapper(config), + ns_resolver=_namespace_utils.NsResolver(namespaces_tuple=_namespace_utils.ns_for_function(func)), + ) + + if schema_type == 'arguments': + schema = generate_schema._arguments_schema(func, parameters_callback) # pyright: ignore[reportArgumentType] + else: + schema = generate_schema._arguments_v3_schema(func, parameters_callback) # pyright: ignore[reportArgumentType] + return generate_schema.clean_schema(schema) diff --git a/lib/python3.12/site-packages/pydantic/experimental/pipeline.py b/lib/python3.12/site-packages/pydantic/experimental/pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..bd63d98c1a3cf62c23662fcaf9953a0bc7ba890e --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/experimental/pipeline.py @@ -0,0 +1,667 @@ +"""Experimental pipeline API functionality. Be careful with this API, it's subject to change.""" + +from __future__ import annotations + +import datetime +import operator +import re +import sys +from collections import deque +from collections.abc import Container +from dataclasses import dataclass +from decimal import Decimal +from functools import cached_property, partial +from re import Pattern +from typing import TYPE_CHECKING, Annotated, Any, Callable, Generic, Protocol, TypeVar, Union, overload + +import annotated_types + +if TYPE_CHECKING: + from pydantic_core import core_schema as cs + + from pydantic import GetCoreSchemaHandler + +from pydantic._internal._internal_dataclass import slots_true as _slots_true + +if sys.version_info < (3, 10): + EllipsisType = type(Ellipsis) +else: + from types import EllipsisType + +__all__ = ['validate_as', 'validate_as_deferred', 'transform'] + +_slots_frozen = {**_slots_true, 'frozen': True} + + +@dataclass(**_slots_frozen) +class _ValidateAs: + tp: type[Any] + strict: bool = False + + +@dataclass +class _ValidateAsDefer: + func: Callable[[], type[Any]] + + @cached_property + def tp(self) -> type[Any]: + return self.func() + + +@dataclass(**_slots_frozen) +class _Transform: + func: Callable[[Any], Any] + + +@dataclass(**_slots_frozen) +class _PipelineOr: + left: _Pipeline[Any, Any] + right: _Pipeline[Any, Any] + + +@dataclass(**_slots_frozen) +class _PipelineAnd: + left: _Pipeline[Any, Any] + right: _Pipeline[Any, Any] + + +@dataclass(**_slots_frozen) +class _Eq: + value: Any + + +@dataclass(**_slots_frozen) +class _NotEq: + value: Any + + +@dataclass(**_slots_frozen) +class _In: + values: Container[Any] + + +@dataclass(**_slots_frozen) +class _NotIn: + values: Container[Any] + + +_ConstraintAnnotation = Union[ + annotated_types.Le, + annotated_types.Ge, + annotated_types.Lt, + annotated_types.Gt, + annotated_types.Len, + annotated_types.MultipleOf, + annotated_types.Timezone, + annotated_types.Interval, + annotated_types.Predicate, + # common predicates not included in annotated_types + _Eq, + _NotEq, + _In, + _NotIn, + # regular expressions + Pattern[str], +] + + +@dataclass(**_slots_frozen) +class _Constraint: + constraint: _ConstraintAnnotation + + +_Step = Union[_ValidateAs, _ValidateAsDefer, _Transform, _PipelineOr, _PipelineAnd, _Constraint] + +_InT = TypeVar('_InT') +_OutT = TypeVar('_OutT') +_NewOutT = TypeVar('_NewOutT') + + +class _FieldTypeMarker: + pass + + +# TODO: ultimately, make this public, see https://github.com/pydantic/pydantic/pull/9459#discussion_r1628197626 +# Also, make this frozen eventually, but that doesn't work right now because of the generic base +# Which attempts to modify __orig_base__ and such. +# We could go with a manual freeze, but that seems overkill for now. +@dataclass(**_slots_true) +class _Pipeline(Generic[_InT, _OutT]): + """Abstract representation of a chain of validation, transformation, and parsing steps.""" + + _steps: tuple[_Step, ...] + + def transform( + self, + func: Callable[[_OutT], _NewOutT], + ) -> _Pipeline[_InT, _NewOutT]: + """Transform the output of the previous step. + + If used as the first step in a pipeline, the type of the field is used. + That is, the transformation is applied to after the value is parsed to the field's type. + """ + return _Pipeline[_InT, _NewOutT](self._steps + (_Transform(func),)) + + @overload + def validate_as(self, tp: type[_NewOutT], *, strict: bool = ...) -> _Pipeline[_InT, _NewOutT]: ... + + @overload + def validate_as(self, tp: EllipsisType, *, strict: bool = ...) -> _Pipeline[_InT, Any]: # type: ignore + ... + + def validate_as(self, tp: type[_NewOutT] | EllipsisType, *, strict: bool = False) -> _Pipeline[_InT, Any]: # type: ignore + """Validate / parse the input into a new type. + + If no type is provided, the type of the field is used. + + Types are parsed in Pydantic's `lax` mode by default, + but you can enable `strict` mode by passing `strict=True`. + """ + if isinstance(tp, EllipsisType): + return _Pipeline[_InT, Any](self._steps + (_ValidateAs(_FieldTypeMarker, strict=strict),)) + return _Pipeline[_InT, _NewOutT](self._steps + (_ValidateAs(tp, strict=strict),)) + + def validate_as_deferred(self, func: Callable[[], type[_NewOutT]]) -> _Pipeline[_InT, _NewOutT]: + """Parse the input into a new type, deferring resolution of the type until the current class + is fully defined. + + This is useful when you need to reference the class in it's own type annotations. + """ + return _Pipeline[_InT, _NewOutT](self._steps + (_ValidateAsDefer(func),)) + + # constraints + @overload + def constrain(self: _Pipeline[_InT, _NewOutGe], constraint: annotated_types.Ge) -> _Pipeline[_InT, _NewOutGe]: ... + + @overload + def constrain(self: _Pipeline[_InT, _NewOutGt], constraint: annotated_types.Gt) -> _Pipeline[_InT, _NewOutGt]: ... + + @overload + def constrain(self: _Pipeline[_InT, _NewOutLe], constraint: annotated_types.Le) -> _Pipeline[_InT, _NewOutLe]: ... + + @overload + def constrain(self: _Pipeline[_InT, _NewOutLt], constraint: annotated_types.Lt) -> _Pipeline[_InT, _NewOutLt]: ... + + @overload + def constrain( + self: _Pipeline[_InT, _NewOutLen], constraint: annotated_types.Len + ) -> _Pipeline[_InT, _NewOutLen]: ... + + @overload + def constrain( + self: _Pipeline[_InT, _NewOutT], constraint: annotated_types.MultipleOf + ) -> _Pipeline[_InT, _NewOutT]: ... + + @overload + def constrain( + self: _Pipeline[_InT, _NewOutDatetime], constraint: annotated_types.Timezone + ) -> _Pipeline[_InT, _NewOutDatetime]: ... + + @overload + def constrain(self: _Pipeline[_InT, _OutT], constraint: annotated_types.Predicate) -> _Pipeline[_InT, _OutT]: ... + + @overload + def constrain( + self: _Pipeline[_InT, _NewOutInterval], constraint: annotated_types.Interval + ) -> _Pipeline[_InT, _NewOutInterval]: ... + + @overload + def constrain(self: _Pipeline[_InT, _OutT], constraint: _Eq) -> _Pipeline[_InT, _OutT]: ... + + @overload + def constrain(self: _Pipeline[_InT, _OutT], constraint: _NotEq) -> _Pipeline[_InT, _OutT]: ... + + @overload + def constrain(self: _Pipeline[_InT, _OutT], constraint: _In) -> _Pipeline[_InT, _OutT]: ... + + @overload + def constrain(self: _Pipeline[_InT, _OutT], constraint: _NotIn) -> _Pipeline[_InT, _OutT]: ... + + @overload + def constrain(self: _Pipeline[_InT, _NewOutT], constraint: Pattern[str]) -> _Pipeline[_InT, _NewOutT]: ... + + def constrain(self, constraint: _ConstraintAnnotation) -> Any: + """Constrain a value to meet a certain condition. + + We support most conditions from `annotated_types`, as well as regular expressions. + + Most of the time you'll be calling a shortcut method like `gt`, `lt`, `len`, etc + so you don't need to call this directly. + """ + return _Pipeline[_InT, _OutT](self._steps + (_Constraint(constraint),)) + + def predicate(self: _Pipeline[_InT, _NewOutT], func: Callable[[_NewOutT], bool]) -> _Pipeline[_InT, _NewOutT]: + """Constrain a value to meet a certain predicate.""" + return self.constrain(annotated_types.Predicate(func)) + + def gt(self: _Pipeline[_InT, _NewOutGt], gt: _NewOutGt) -> _Pipeline[_InT, _NewOutGt]: + """Constrain a value to be greater than a certain value.""" + return self.constrain(annotated_types.Gt(gt)) + + def lt(self: _Pipeline[_InT, _NewOutLt], lt: _NewOutLt) -> _Pipeline[_InT, _NewOutLt]: + """Constrain a value to be less than a certain value.""" + return self.constrain(annotated_types.Lt(lt)) + + def ge(self: _Pipeline[_InT, _NewOutGe], ge: _NewOutGe) -> _Pipeline[_InT, _NewOutGe]: + """Constrain a value to be greater than or equal to a certain value.""" + return self.constrain(annotated_types.Ge(ge)) + + def le(self: _Pipeline[_InT, _NewOutLe], le: _NewOutLe) -> _Pipeline[_InT, _NewOutLe]: + """Constrain a value to be less than or equal to a certain value.""" + return self.constrain(annotated_types.Le(le)) + + def len(self: _Pipeline[_InT, _NewOutLen], min_len: int, max_len: int | None = None) -> _Pipeline[_InT, _NewOutLen]: + """Constrain a value to have a certain length.""" + return self.constrain(annotated_types.Len(min_len, max_len)) + + @overload + def multiple_of(self: _Pipeline[_InT, _NewOutDiv], multiple_of: _NewOutDiv) -> _Pipeline[_InT, _NewOutDiv]: ... + + @overload + def multiple_of(self: _Pipeline[_InT, _NewOutMod], multiple_of: _NewOutMod) -> _Pipeline[_InT, _NewOutMod]: ... + + def multiple_of(self: _Pipeline[_InT, Any], multiple_of: Any) -> _Pipeline[_InT, Any]: + """Constrain a value to be a multiple of a certain number.""" + return self.constrain(annotated_types.MultipleOf(multiple_of)) + + def eq(self: _Pipeline[_InT, _OutT], value: _OutT) -> _Pipeline[_InT, _OutT]: + """Constrain a value to be equal to a certain value.""" + return self.constrain(_Eq(value)) + + def not_eq(self: _Pipeline[_InT, _OutT], value: _OutT) -> _Pipeline[_InT, _OutT]: + """Constrain a value to not be equal to a certain value.""" + return self.constrain(_NotEq(value)) + + def in_(self: _Pipeline[_InT, _OutT], values: Container[_OutT]) -> _Pipeline[_InT, _OutT]: + """Constrain a value to be in a certain set.""" + return self.constrain(_In(values)) + + def not_in(self: _Pipeline[_InT, _OutT], values: Container[_OutT]) -> _Pipeline[_InT, _OutT]: + """Constrain a value to not be in a certain set.""" + return self.constrain(_NotIn(values)) + + # timezone methods + def datetime_tz_naive(self: _Pipeline[_InT, datetime.datetime]) -> _Pipeline[_InT, datetime.datetime]: + return self.constrain(annotated_types.Timezone(None)) + + def datetime_tz_aware(self: _Pipeline[_InT, datetime.datetime]) -> _Pipeline[_InT, datetime.datetime]: + return self.constrain(annotated_types.Timezone(...)) + + def datetime_tz( + self: _Pipeline[_InT, datetime.datetime], tz: datetime.tzinfo + ) -> _Pipeline[_InT, datetime.datetime]: + return self.constrain(annotated_types.Timezone(tz)) # type: ignore + + def datetime_with_tz( + self: _Pipeline[_InT, datetime.datetime], tz: datetime.tzinfo | None + ) -> _Pipeline[_InT, datetime.datetime]: + return self.transform(partial(datetime.datetime.replace, tzinfo=tz)) + + # string methods + def str_lower(self: _Pipeline[_InT, str]) -> _Pipeline[_InT, str]: + return self.transform(str.lower) + + def str_upper(self: _Pipeline[_InT, str]) -> _Pipeline[_InT, str]: + return self.transform(str.upper) + + def str_title(self: _Pipeline[_InT, str]) -> _Pipeline[_InT, str]: + return self.transform(str.title) + + def str_strip(self: _Pipeline[_InT, str]) -> _Pipeline[_InT, str]: + return self.transform(str.strip) + + def str_pattern(self: _Pipeline[_InT, str], pattern: str) -> _Pipeline[_InT, str]: + return self.constrain(re.compile(pattern)) + + def str_contains(self: _Pipeline[_InT, str], substring: str) -> _Pipeline[_InT, str]: + return self.predicate(lambda v: substring in v) + + def str_starts_with(self: _Pipeline[_InT, str], prefix: str) -> _Pipeline[_InT, str]: + return self.predicate(lambda v: v.startswith(prefix)) + + def str_ends_with(self: _Pipeline[_InT, str], suffix: str) -> _Pipeline[_InT, str]: + return self.predicate(lambda v: v.endswith(suffix)) + + # operators + def otherwise(self, other: _Pipeline[_OtherIn, _OtherOut]) -> _Pipeline[_InT | _OtherIn, _OutT | _OtherOut]: + """Combine two validation chains, returning the result of the first chain if it succeeds, and the second chain if it fails.""" + return _Pipeline((_PipelineOr(self, other),)) + + __or__ = otherwise + + def then(self, other: _Pipeline[_OutT, _OtherOut]) -> _Pipeline[_InT, _OtherOut]: + """Pipe the result of one validation chain into another.""" + return _Pipeline((_PipelineAnd(self, other),)) + + __and__ = then + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> cs.CoreSchema: + from pydantic_core import core_schema as cs + + queue = deque(self._steps) + + s = None + + while queue: + step = queue.popleft() + s = _apply_step(step, s, handler, source_type) + + s = s or cs.any_schema() + return s + + def __supports_type__(self, _: _OutT) -> bool: + raise NotImplementedError + + +validate_as = _Pipeline[Any, Any](()).validate_as +validate_as_deferred = _Pipeline[Any, Any](()).validate_as_deferred +transform = _Pipeline[Any, Any]((_ValidateAs(_FieldTypeMarker),)).transform + + +def _check_func( + func: Callable[[Any], bool], predicate_err: str | Callable[[], str], s: cs.CoreSchema | None +) -> cs.CoreSchema: + from pydantic_core import core_schema as cs + + def handler(v: Any) -> Any: + if func(v): + return v + raise ValueError(f'Expected {predicate_err if isinstance(predicate_err, str) else predicate_err()}') + + if s is None: + return cs.no_info_plain_validator_function(handler) + else: + return cs.no_info_after_validator_function(handler, s) + + +def _apply_step(step: _Step, s: cs.CoreSchema | None, handler: GetCoreSchemaHandler, source_type: Any) -> cs.CoreSchema: + from pydantic_core import core_schema as cs + + if isinstance(step, _ValidateAs): + s = _apply_parse(s, step.tp, step.strict, handler, source_type) + elif isinstance(step, _ValidateAsDefer): + s = _apply_parse(s, step.tp, False, handler, source_type) + elif isinstance(step, _Transform): + s = _apply_transform(s, step.func, handler) + elif isinstance(step, _Constraint): + s = _apply_constraint(s, step.constraint) + elif isinstance(step, _PipelineOr): + s = cs.union_schema([handler(step.left), handler(step.right)]) + else: + assert isinstance(step, _PipelineAnd) + s = cs.chain_schema([handler(step.left), handler(step.right)]) + return s + + +def _apply_parse( + s: cs.CoreSchema | None, + tp: type[Any], + strict: bool, + handler: GetCoreSchemaHandler, + source_type: Any, +) -> cs.CoreSchema: + from pydantic_core import core_schema as cs + + from pydantic import Strict + + if tp is _FieldTypeMarker: + return cs.chain_schema([s, handler(source_type)]) if s else handler(source_type) + + if strict: + tp = Annotated[tp, Strict()] # type: ignore + + if s and s['type'] == 'any': + return handler(tp) + else: + return cs.chain_schema([s, handler(tp)]) if s else handler(tp) + + +def _apply_transform( + s: cs.CoreSchema | None, func: Callable[[Any], Any], handler: GetCoreSchemaHandler +) -> cs.CoreSchema: + from pydantic_core import core_schema as cs + + if s is None: + return cs.no_info_plain_validator_function(func) + + if s['type'] == 'str': + if func is str.strip: + s = s.copy() + s['strip_whitespace'] = True + return s + elif func is str.lower: + s = s.copy() + s['to_lower'] = True + return s + elif func is str.upper: + s = s.copy() + s['to_upper'] = True + return s + + return cs.no_info_after_validator_function(func, s) + + +def _apply_constraint( # noqa: C901 + s: cs.CoreSchema | None, constraint: _ConstraintAnnotation +) -> cs.CoreSchema: + """Apply a single constraint to a schema.""" + if isinstance(constraint, annotated_types.Gt): + gt = constraint.gt + if s and s['type'] in {'int', 'float', 'decimal'}: + s = s.copy() + if s['type'] == 'int' and isinstance(gt, int): + s['gt'] = gt + elif s['type'] == 'float' and isinstance(gt, float): + s['gt'] = gt + elif s['type'] == 'decimal' and isinstance(gt, Decimal): + s['gt'] = gt + else: + + def check_gt(v: Any) -> bool: + return v > gt + + s = _check_func(check_gt, f'> {gt}', s) + elif isinstance(constraint, annotated_types.Ge): + ge = constraint.ge + if s and s['type'] in {'int', 'float', 'decimal'}: + s = s.copy() + if s['type'] == 'int' and isinstance(ge, int): + s['ge'] = ge + elif s['type'] == 'float' and isinstance(ge, float): + s['ge'] = ge + elif s['type'] == 'decimal' and isinstance(ge, Decimal): + s['ge'] = ge + + def check_ge(v: Any) -> bool: + return v >= ge + + s = _check_func(check_ge, f'>= {ge}', s) + elif isinstance(constraint, annotated_types.Lt): + lt = constraint.lt + if s and s['type'] in {'int', 'float', 'decimal'}: + s = s.copy() + if s['type'] == 'int' and isinstance(lt, int): + s['lt'] = lt + elif s['type'] == 'float' and isinstance(lt, float): + s['lt'] = lt + elif s['type'] == 'decimal' and isinstance(lt, Decimal): + s['lt'] = lt + + def check_lt(v: Any) -> bool: + return v < lt + + s = _check_func(check_lt, f'< {lt}', s) + elif isinstance(constraint, annotated_types.Le): + le = constraint.le + if s and s['type'] in {'int', 'float', 'decimal'}: + s = s.copy() + if s['type'] == 'int' and isinstance(le, int): + s['le'] = le + elif s['type'] == 'float' and isinstance(le, float): + s['le'] = le + elif s['type'] == 'decimal' and isinstance(le, Decimal): + s['le'] = le + + def check_le(v: Any) -> bool: + return v <= le + + s = _check_func(check_le, f'<= {le}', s) + elif isinstance(constraint, annotated_types.Len): + min_len = constraint.min_length + max_len = constraint.max_length + + if s and s['type'] in {'str', 'list', 'tuple', 'set', 'frozenset', 'dict'}: + assert ( + s['type'] == 'str' + or s['type'] == 'list' + or s['type'] == 'tuple' + or s['type'] == 'set' + or s['type'] == 'dict' + or s['type'] == 'frozenset' + ) + s = s.copy() + if min_len != 0: + s['min_length'] = min_len + if max_len is not None: + s['max_length'] = max_len + + def check_len(v: Any) -> bool: + if max_len is not None: + return (min_len <= len(v)) and (len(v) <= max_len) + return min_len <= len(v) + + s = _check_func(check_len, f'length >= {min_len} and length <= {max_len}', s) + elif isinstance(constraint, annotated_types.MultipleOf): + multiple_of = constraint.multiple_of + if s and s['type'] in {'int', 'float', 'decimal'}: + s = s.copy() + if s['type'] == 'int' and isinstance(multiple_of, int): + s['multiple_of'] = multiple_of + elif s['type'] == 'float' and isinstance(multiple_of, float): + s['multiple_of'] = multiple_of + elif s['type'] == 'decimal' and isinstance(multiple_of, Decimal): + s['multiple_of'] = multiple_of + + def check_multiple_of(v: Any) -> bool: + return v % multiple_of == 0 + + s = _check_func(check_multiple_of, f'% {multiple_of} == 0', s) + elif isinstance(constraint, annotated_types.Timezone): + tz = constraint.tz + + if tz is ...: + if s and s['type'] == 'datetime': + s = s.copy() + s['tz_constraint'] = 'aware' + else: + + def check_tz_aware(v: object) -> bool: + assert isinstance(v, datetime.datetime) + return v.tzinfo is not None + + s = _check_func(check_tz_aware, 'timezone aware', s) + elif tz is None: + if s and s['type'] == 'datetime': + s = s.copy() + s['tz_constraint'] = 'naive' + else: + + def check_tz_naive(v: object) -> bool: + assert isinstance(v, datetime.datetime) + return v.tzinfo is None + + s = _check_func(check_tz_naive, 'timezone naive', s) + else: + raise NotImplementedError('Constraining to a specific timezone is not yet supported') + elif isinstance(constraint, annotated_types.Interval): + if constraint.ge: + s = _apply_constraint(s, annotated_types.Ge(constraint.ge)) + if constraint.gt: + s = _apply_constraint(s, annotated_types.Gt(constraint.gt)) + if constraint.le: + s = _apply_constraint(s, annotated_types.Le(constraint.le)) + if constraint.lt: + s = _apply_constraint(s, annotated_types.Lt(constraint.lt)) + assert s is not None + elif isinstance(constraint, annotated_types.Predicate): + func = constraint.func + + if func.__name__ == '': + # attempt to extract the source code for a lambda function + # to use as the function name in error messages + # TODO: is there a better way? should we just not do this? + import inspect + + try: + source = inspect.getsource(func).strip() + source = source.removesuffix(')') + lambda_source_code = '`' + ''.join(''.join(source.split('lambda ')[1:]).split(':')[1:]).strip() + '`' + except OSError: + # stringified annotations + lambda_source_code = 'lambda' + + s = _check_func(func, lambda_source_code, s) + else: + s = _check_func(func, func.__name__, s) + elif isinstance(constraint, _NotEq): + value = constraint.value + + def check_not_eq(v: Any) -> bool: + return operator.__ne__(v, value) + + s = _check_func(check_not_eq, f'!= {value}', s) + elif isinstance(constraint, _Eq): + value = constraint.value + + def check_eq(v: Any) -> bool: + return operator.__eq__(v, value) + + s = _check_func(check_eq, f'== {value}', s) + elif isinstance(constraint, _In): + values = constraint.values + + def check_in(v: Any) -> bool: + return operator.__contains__(values, v) + + s = _check_func(check_in, f'in {values}', s) + elif isinstance(constraint, _NotIn): + values = constraint.values + + def check_not_in(v: Any) -> bool: + return operator.__not__(operator.__contains__(values, v)) + + s = _check_func(check_not_in, f'not in {values}', s) + else: + assert isinstance(constraint, Pattern) + if s and s['type'] == 'str': + s = s.copy() + s['pattern'] = constraint.pattern + else: + + def check_pattern(v: object) -> bool: + assert isinstance(v, str) + return constraint.match(v) is not None + + s = _check_func(check_pattern, f'~ {constraint.pattern}', s) + return s + + +class _SupportsRange(annotated_types.SupportsLe, annotated_types.SupportsGe, Protocol): + pass + + +class _SupportsLen(Protocol): + def __len__(self) -> int: ... + + +_NewOutGt = TypeVar('_NewOutGt', bound=annotated_types.SupportsGt) +_NewOutGe = TypeVar('_NewOutGe', bound=annotated_types.SupportsGe) +_NewOutLt = TypeVar('_NewOutLt', bound=annotated_types.SupportsLt) +_NewOutLe = TypeVar('_NewOutLe', bound=annotated_types.SupportsLe) +_NewOutLen = TypeVar('_NewOutLen', bound=_SupportsLen) +_NewOutDiv = TypeVar('_NewOutDiv', bound=annotated_types.SupportsDiv) +_NewOutMod = TypeVar('_NewOutMod', bound=annotated_types.SupportsMod) +_NewOutDatetime = TypeVar('_NewOutDatetime', bound=datetime.datetime) +_NewOutInterval = TypeVar('_NewOutInterval', bound=_SupportsRange) +_OtherIn = TypeVar('_OtherIn') +_OtherOut = TypeVar('_OtherOut') diff --git a/lib/python3.12/site-packages/pydantic/fields.py b/lib/python3.12/site-packages/pydantic/fields.py new file mode 100644 index 0000000000000000000000000000000000000000..408fc2b3d9e4719d110d3d0049b60e283dd777e1 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/fields.py @@ -0,0 +1,1559 @@ +"""Defining fields on models.""" + +from __future__ import annotations as _annotations + +import dataclasses +import inspect +import sys +import typing +from collections.abc import Mapping +from copy import copy +from dataclasses import Field as DataclassField +from functools import cached_property +from typing import Annotated, Any, Callable, ClassVar, Literal, TypeVar, cast, overload +from warnings import warn + +import annotated_types +import typing_extensions +from pydantic_core import PydanticUndefined +from typing_extensions import Self, TypeAlias, Unpack, deprecated +from typing_inspection import typing_objects +from typing_inspection.introspection import UNKNOWN, AnnotationSource, ForbiddenQualifier, Qualifier, inspect_annotation + +from . import types +from ._internal import _decorators, _fields, _generics, _internal_dataclass, _repr, _typing_extra, _utils +from ._internal._namespace_utils import GlobalsNamespace, MappingNamespace +from .aliases import AliasChoices, AliasPath +from .config import JsonDict +from .errors import PydanticForbiddenQualifier, PydanticUserError +from .json_schema import PydanticJsonSchemaWarning +from .warnings import PydanticDeprecatedSince20 + +if typing.TYPE_CHECKING: + from ._internal._repr import ReprArgs +else: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + +__all__ = 'Field', 'PrivateAttr', 'computed_field' + + +_Unset: Any = PydanticUndefined + +if sys.version_info >= (3, 13): + import warnings + + Deprecated: TypeAlias = warnings.deprecated | deprecated +else: + Deprecated: TypeAlias = deprecated + + +class _FromFieldInfoInputs(typing_extensions.TypedDict, total=False): + """This class exists solely to add type checking for the `**kwargs` in `FieldInfo.from_field`.""" + + # TODO PEP 747: use TypeForm: + annotation: type[Any] | None + default_factory: Callable[[], Any] | Callable[[dict[str, Any]], Any] | None + alias: str | None + alias_priority: int | None + validation_alias: str | AliasPath | AliasChoices | None + serialization_alias: str | None + title: str | None + field_title_generator: Callable[[str, FieldInfo], str] | None + description: str | None + examples: list[Any] | None + exclude: bool | None + gt: annotated_types.SupportsGt | None + ge: annotated_types.SupportsGe | None + lt: annotated_types.SupportsLt | None + le: annotated_types.SupportsLe | None + multiple_of: float | None + strict: bool | None + min_length: int | None + max_length: int | None + pattern: str | typing.Pattern[str] | None + allow_inf_nan: bool | None + max_digits: int | None + decimal_places: int | None + union_mode: Literal['smart', 'left_to_right'] | None + discriminator: str | types.Discriminator | None + deprecated: Deprecated | str | bool | None + json_schema_extra: JsonDict | Callable[[JsonDict], None] | None + frozen: bool | None + validate_default: bool | None + repr: bool + init: bool | None + init_var: bool | None + kw_only: bool | None + coerce_numbers_to_str: bool | None + fail_fast: bool | None + + +class _FieldInfoInputs(_FromFieldInfoInputs, total=False): + """This class exists solely to add type checking for the `**kwargs` in `FieldInfo.__init__`.""" + + default: Any + + +class FieldInfo(_repr.Representation): + """This class holds information about a field. + + `FieldInfo` is used for any field definition regardless of whether the [`Field()`][pydantic.fields.Field] + function is explicitly used. + + !!! warning + You generally shouldn't be creating `FieldInfo` directly, you'll only need to use it when accessing + [`BaseModel`][pydantic.main.BaseModel] `.model_fields` internals. + + Attributes: + annotation: The type annotation of the field. + default: The default value of the field. + default_factory: A callable to generate the default value. The callable can either take 0 arguments + (in which case it is called as is) or a single argument containing the already validated data. + alias: The alias name of the field. + alias_priority: The priority of the field's alias. + validation_alias: The validation alias of the field. + serialization_alias: The serialization alias of the field. + title: The title of the field. + field_title_generator: A callable that takes a field name and returns title for it. + description: The description of the field. + examples: List of examples of the field. + exclude: Whether to exclude the field from the model serialization. + discriminator: Field name or Discriminator for discriminating the type in a tagged union. + deprecated: A deprecation message, an instance of `warnings.deprecated` or the `typing_extensions.deprecated` backport, + or a boolean. If `True`, a default deprecation message will be emitted when accessing the field. + json_schema_extra: A dict or callable to provide extra JSON schema properties. + frozen: Whether the field is frozen. + validate_default: Whether to validate the default value of the field. + repr: Whether to include the field in representation of the model. + init: Whether the field should be included in the constructor of the dataclass. + init_var: Whether the field should _only_ be included in the constructor of the dataclass, and not stored. + kw_only: Whether the field should be a keyword-only argument in the constructor of the dataclass. + metadata: List of metadata constraints. + """ + + annotation: type[Any] | None + default: Any + default_factory: Callable[[], Any] | Callable[[dict[str, Any]], Any] | None + alias: str | None + alias_priority: int | None + validation_alias: str | AliasPath | AliasChoices | None + serialization_alias: str | None + title: str | None + field_title_generator: Callable[[str, FieldInfo], str] | None + description: str | None + examples: list[Any] | None + exclude: bool | None + discriminator: str | types.Discriminator | None + deprecated: Deprecated | str | bool | None + json_schema_extra: JsonDict | Callable[[JsonDict], None] | None + frozen: bool | None + validate_default: bool | None + repr: bool + init: bool | None + init_var: bool | None + kw_only: bool | None + metadata: list[Any] + + __slots__ = ( + 'annotation', + 'default', + 'default_factory', + 'alias', + 'alias_priority', + 'validation_alias', + 'serialization_alias', + 'title', + 'field_title_generator', + 'description', + 'examples', + 'exclude', + 'discriminator', + 'deprecated', + 'json_schema_extra', + 'frozen', + 'validate_default', + 'repr', + 'init', + 'init_var', + 'kw_only', + 'metadata', + '_attributes_set', + '_qualifiers', + '_complete', + '_original_assignment', + '_original_annotation', + ) + + # used to convert kwargs to metadata/constraints, + # None has a special meaning - these items are collected into a `PydanticGeneralMetadata` + metadata_lookup: ClassVar[dict[str, typing.Callable[[Any], Any] | None]] = { + 'strict': types.Strict, + 'gt': annotated_types.Gt, + 'ge': annotated_types.Ge, + 'lt': annotated_types.Lt, + 'le': annotated_types.Le, + 'multiple_of': annotated_types.MultipleOf, + 'min_length': annotated_types.MinLen, + 'max_length': annotated_types.MaxLen, + 'pattern': None, + 'allow_inf_nan': None, + 'max_digits': None, + 'decimal_places': None, + 'union_mode': None, + 'coerce_numbers_to_str': None, + 'fail_fast': types.FailFast, + } + + def __init__(self, **kwargs: Unpack[_FieldInfoInputs]) -> None: + """This class should generally not be initialized directly; instead, use the `pydantic.fields.Field` function + or one of the constructor classmethods. + + See the signature of `pydantic.fields.Field` for more details about the expected arguments. + """ + self._attributes_set = {k: v for k, v in kwargs.items() if v is not _Unset} + kwargs = {k: _DefaultValues.get(k) if v is _Unset else v for k, v in kwargs.items()} # type: ignore + self.annotation = kwargs.get('annotation') + + default = kwargs.pop('default', PydanticUndefined) + if default is Ellipsis: + self.default = PydanticUndefined + self._attributes_set.pop('default', None) + else: + self.default = default + + self.default_factory = kwargs.pop('default_factory', None) + + if self.default is not PydanticUndefined and self.default_factory is not None: + raise TypeError('cannot specify both default and default_factory') + + self.alias = kwargs.pop('alias', None) + self.validation_alias = kwargs.pop('validation_alias', None) + self.serialization_alias = kwargs.pop('serialization_alias', None) + alias_is_set = any(alias is not None for alias in (self.alias, self.validation_alias, self.serialization_alias)) + self.alias_priority = kwargs.pop('alias_priority', None) or 2 if alias_is_set else None + self.title = kwargs.pop('title', None) + self.field_title_generator = kwargs.pop('field_title_generator', None) + self.description = kwargs.pop('description', None) + self.examples = kwargs.pop('examples', None) + self.exclude = kwargs.pop('exclude', None) + self.discriminator = kwargs.pop('discriminator', None) + # For compatibility with FastAPI<=0.110.0, we preserve the existing value if it is not overridden + self.deprecated = kwargs.pop('deprecated', getattr(self, 'deprecated', None)) + self.repr = kwargs.pop('repr', True) + self.json_schema_extra = kwargs.pop('json_schema_extra', None) + self.validate_default = kwargs.pop('validate_default', None) + self.frozen = kwargs.pop('frozen', None) + # currently only used on dataclasses + self.init = kwargs.pop('init', None) + self.init_var = kwargs.pop('init_var', None) + self.kw_only = kwargs.pop('kw_only', None) + + self.metadata = self._collect_metadata(kwargs) # type: ignore + + # Private attributes: + self._qualifiers: set[Qualifier] = set() + # Used to rebuild FieldInfo instances: + self._complete = True + self._original_annotation: Any = PydanticUndefined + self._original_assignment: Any = PydanticUndefined + + @staticmethod + def from_field(default: Any = PydanticUndefined, **kwargs: Unpack[_FromFieldInfoInputs]) -> FieldInfo: + """Create a new `FieldInfo` object with the `Field` function. + + Args: + default: The default value for the field. Defaults to Undefined. + **kwargs: Additional arguments dictionary. + + Raises: + TypeError: If 'annotation' is passed as a keyword argument. + + Returns: + A new FieldInfo object with the given parameters. + + Example: + This is how you can create a field with default value like this: + + ```python + import pydantic + + class MyModel(pydantic.BaseModel): + foo: int = pydantic.Field(4) + ``` + """ + if 'annotation' in kwargs: + raise TypeError('"annotation" is not permitted as a Field keyword argument') + return FieldInfo(default=default, **kwargs) + + @staticmethod + def from_annotation(annotation: type[Any], *, _source: AnnotationSource = AnnotationSource.ANY) -> FieldInfo: + """Creates a `FieldInfo` instance from a bare annotation. + + This function is used internally to create a `FieldInfo` from a bare annotation like this: + + ```python + import pydantic + + class MyModel(pydantic.BaseModel): + foo: int # <-- like this + ``` + + We also account for the case where the annotation can be an instance of `Annotated` and where + one of the (not first) arguments in `Annotated` is an instance of `FieldInfo`, e.g.: + + ```python + from typing import Annotated + + import annotated_types + + import pydantic + + class MyModel(pydantic.BaseModel): + foo: Annotated[int, annotated_types.Gt(42)] + bar: Annotated[int, pydantic.Field(gt=42)] + ``` + + Args: + annotation: An annotation object. + + Returns: + An instance of the field metadata. + """ + try: + inspected_ann = inspect_annotation( + annotation, + annotation_source=_source, + unpack_type_aliases='skip', + ) + except ForbiddenQualifier as e: + raise PydanticForbiddenQualifier(e.qualifier, annotation) + + # TODO check for classvar and error? + + # No assigned value, this happens when using a bare `Final` qualifier (also for other + # qualifiers, but they shouldn't appear here). In this case we infer the type as `Any` + # because we don't have any assigned value. + type_expr: Any = Any if inspected_ann.type is UNKNOWN else inspected_ann.type + final = 'final' in inspected_ann.qualifiers + metadata = inspected_ann.metadata + + if not metadata: + # No metadata, e.g. `field: int`, or `field: Final[str]`: + field_info = FieldInfo(annotation=type_expr, frozen=final or None) + field_info._qualifiers = inspected_ann.qualifiers + return field_info + + # With metadata, e.g. `field: Annotated[int, Field(...), Gt(1)]`: + field_info_annotations = [a for a in metadata if isinstance(a, FieldInfo)] + field_info = FieldInfo.merge_field_infos(*field_info_annotations, annotation=type_expr) + + new_field_info = field_info._copy() + new_field_info.annotation = type_expr + new_field_info.frozen = final or field_info.frozen + field_metadata: list[Any] = [] + for a in metadata: + if typing_objects.is_deprecated(a): + new_field_info.deprecated = a.message + elif not isinstance(a, FieldInfo): + field_metadata.append(a) + else: + field_metadata.extend(a.metadata) + new_field_info.metadata = field_metadata + new_field_info._qualifiers = inspected_ann.qualifiers + return new_field_info + + @staticmethod + def from_annotated_attribute( + annotation: type[Any], default: Any, *, _source: AnnotationSource = AnnotationSource.ANY + ) -> FieldInfo: + """Create `FieldInfo` from an annotation with a default value. + + This is used in cases like the following: + + ```python + from typing import Annotated + + import annotated_types + + import pydantic + + class MyModel(pydantic.BaseModel): + foo: int = 4 # <-- like this + bar: Annotated[int, annotated_types.Gt(4)] = 4 # <-- or this + spam: Annotated[int, pydantic.Field(gt=4)] = 4 # <-- or this + ``` + + Args: + annotation: The type annotation of the field. + default: The default value of the field. + + Returns: + A field object with the passed values. + """ + if annotation is default: + raise PydanticUserError( + 'Error when building FieldInfo from annotated attribute. ' + "Make sure you don't have any field name clashing with a type annotation.", + code='unevaluable-type-annotation', + ) + + try: + inspected_ann = inspect_annotation( + annotation, + annotation_source=_source, + unpack_type_aliases='skip', + ) + except ForbiddenQualifier as e: + raise PydanticForbiddenQualifier(e.qualifier, annotation) + + # TODO check for classvar and error? + + # TODO infer from the default, this can be done in v3 once we treat final fields with + # a default as proper fields and not class variables: + type_expr: Any = Any if inspected_ann.type is UNKNOWN else inspected_ann.type + final = 'final' in inspected_ann.qualifiers + metadata = inspected_ann.metadata + + if isinstance(default, FieldInfo): + # e.g. `field: int = Field(...)` + default_metadata = default.metadata.copy() + default = copy(default) + default.metadata = default_metadata + + default.annotation = type_expr + default.metadata += metadata + merged_default = FieldInfo.merge_field_infos( + *[x for x in metadata if isinstance(x, FieldInfo)], + default, + annotation=default.annotation, + ) + merged_default.frozen = final or merged_default.frozen + merged_default._qualifiers = inspected_ann.qualifiers + return merged_default + + if isinstance(default, dataclasses.Field): + # `collect_dataclass_fields()` passes the dataclass Field as a default. + pydantic_field = FieldInfo._from_dataclass_field(default) + pydantic_field.annotation = type_expr + pydantic_field.metadata += metadata + pydantic_field = FieldInfo.merge_field_infos( + *[x for x in metadata if isinstance(x, FieldInfo)], + pydantic_field, + annotation=pydantic_field.annotation, + ) + pydantic_field.frozen = final or pydantic_field.frozen + pydantic_field.init_var = 'init_var' in inspected_ann.qualifiers + pydantic_field.init = getattr(default, 'init', None) + pydantic_field.kw_only = getattr(default, 'kw_only', None) + pydantic_field._qualifiers = inspected_ann.qualifiers + return pydantic_field + + if not metadata: + # No metadata, e.g. `field: int = ...`, or `field: Final[str] = ...`: + field_info = FieldInfo(annotation=type_expr, default=default, frozen=final or None) + field_info._qualifiers = inspected_ann.qualifiers + return field_info + + # With metadata, e.g. `field: Annotated[int, Field(...), Gt(1)] = ...`: + field_infos = [a for a in metadata if isinstance(a, FieldInfo)] + field_info = FieldInfo.merge_field_infos(*field_infos, annotation=type_expr, default=default) + field_metadata: list[Any] = [] + for a in metadata: + if typing_objects.is_deprecated(a): + field_info.deprecated = a.message + elif not isinstance(a, FieldInfo): + field_metadata.append(a) + else: + field_metadata.extend(a.metadata) + field_info.metadata = field_metadata + field_info._qualifiers = inspected_ann.qualifiers + return field_info + + @staticmethod + def merge_field_infos(*field_infos: FieldInfo, **overrides: Any) -> FieldInfo: + """Merge `FieldInfo` instances keeping only explicitly set attributes. + + Later `FieldInfo` instances override earlier ones. + + Returns: + FieldInfo: A merged FieldInfo instance. + """ + if len(field_infos) == 1: + # No merging necessary, but we still need to make a copy and apply the overrides + field_info = field_infos[0]._copy() + field_info._attributes_set.update(overrides) + + default_override = overrides.pop('default', PydanticUndefined) + if default_override is Ellipsis: + default_override = PydanticUndefined + if default_override is not PydanticUndefined: + field_info.default = default_override + + for k, v in overrides.items(): + setattr(field_info, k, v) + return field_info # type: ignore + + merged_field_info_kwargs: dict[str, Any] = {} + metadata = {} + for field_info in field_infos: + attributes_set = field_info._attributes_set.copy() + + try: + json_schema_extra = attributes_set.pop('json_schema_extra') + existing_json_schema_extra = merged_field_info_kwargs.get('json_schema_extra') + + if existing_json_schema_extra is None: + merged_field_info_kwargs['json_schema_extra'] = json_schema_extra + if isinstance(existing_json_schema_extra, dict): + if isinstance(json_schema_extra, dict): + merged_field_info_kwargs['json_schema_extra'] = { + **existing_json_schema_extra, + **json_schema_extra, + } + if callable(json_schema_extra): + warn( + 'Composing `dict` and `callable` type `json_schema_extra` is not supported.' + 'The `callable` type is being ignored.' + "If you'd like support for this behavior, please open an issue on pydantic.", + PydanticJsonSchemaWarning, + ) + elif callable(json_schema_extra): + # if ever there's a case of a callable, we'll just keep the last json schema extra spec + merged_field_info_kwargs['json_schema_extra'] = json_schema_extra + except KeyError: + pass + + # later FieldInfo instances override everything except json_schema_extra from earlier FieldInfo instances + merged_field_info_kwargs.update(attributes_set) + + for x in field_info.metadata: + if not isinstance(x, FieldInfo): + metadata[type(x)] = x + + merged_field_info_kwargs.update(overrides) + field_info = FieldInfo(**merged_field_info_kwargs) + field_info.metadata = list(metadata.values()) + return field_info + + @staticmethod + def _from_dataclass_field(dc_field: DataclassField[Any]) -> FieldInfo: + """Return a new `FieldInfo` instance from a `dataclasses.Field` instance. + + Args: + dc_field: The `dataclasses.Field` instance to convert. + + Returns: + The corresponding `FieldInfo` instance. + + Raises: + TypeError: If any of the `FieldInfo` kwargs does not match the `dataclass.Field` kwargs. + """ + default = dc_field.default + if default is dataclasses.MISSING: + default = _Unset + + if dc_field.default_factory is dataclasses.MISSING: + default_factory = _Unset + else: + default_factory = dc_field.default_factory + + # use the `Field` function so in correct kwargs raise the correct `TypeError` + dc_field_metadata = {k: v for k, v in dc_field.metadata.items() if k in _FIELD_ARG_NAMES} + return Field(default=default, default_factory=default_factory, repr=dc_field.repr, **dc_field_metadata) # pyright: ignore[reportCallIssue] + + @staticmethod + def _collect_metadata(kwargs: dict[str, Any]) -> list[Any]: + """Collect annotations from kwargs. + + Args: + kwargs: Keyword arguments passed to the function. + + Returns: + A list of metadata objects - a combination of `annotated_types.BaseMetadata` and + `PydanticMetadata`. + """ + metadata: list[Any] = [] + general_metadata = {} + for key, value in list(kwargs.items()): + try: + marker = FieldInfo.metadata_lookup[key] + except KeyError: + continue + + del kwargs[key] + if value is not None: + if marker is None: + general_metadata[key] = value + else: + metadata.append(marker(value)) + if general_metadata: + metadata.append(_fields.pydantic_general_metadata(**general_metadata)) + return metadata + + def _copy(self) -> Self: + copied = copy(self) + for attr_name in ('metadata', '_attributes_set', '_qualifiers'): + # Apply "deep-copy" behavior on collections attributes: + value = getattr(copied, attr_name).copy() + setattr(copied, attr_name, value) + + return copied + + @property + def deprecation_message(self) -> str | None: + """The deprecation message to be emitted, or `None` if not set.""" + if self.deprecated is None: + return None + if isinstance(self.deprecated, bool): + return 'deprecated' if self.deprecated else None + return self.deprecated if isinstance(self.deprecated, str) else self.deprecated.message + + @property + def default_factory_takes_validated_data(self) -> bool | None: + """Whether the provided default factory callable has a validated data parameter. + + Returns `None` if no default factory is set. + """ + if self.default_factory is not None: + return _fields.takes_validated_data_argument(self.default_factory) + + @overload + def get_default( + self, *, call_default_factory: Literal[True], validated_data: dict[str, Any] | None = None + ) -> Any: ... + + @overload + def get_default(self, *, call_default_factory: Literal[False] = ...) -> Any: ... + + def get_default(self, *, call_default_factory: bool = False, validated_data: dict[str, Any] | None = None) -> Any: + """Get the default value. + + We expose an option for whether to call the default_factory (if present), as calling it may + result in side effects that we want to avoid. However, there are times when it really should + be called (namely, when instantiating a model via `model_construct`). + + Args: + call_default_factory: Whether to call the default factory or not. + validated_data: The already validated data to be passed to the default factory. + + Returns: + The default value, calling the default factory if requested or `None` if not set. + """ + if self.default_factory is None: + return _utils.smart_deepcopy(self.default) + elif call_default_factory: + if self.default_factory_takes_validated_data: + fac = cast('Callable[[dict[str, Any]], Any]', self.default_factory) + if validated_data is None: + raise ValueError( + "The default factory requires the 'validated_data' argument, which was not provided when calling 'get_default'." + ) + return fac(validated_data) + else: + fac = cast('Callable[[], Any]', self.default_factory) + return fac() + else: + return None + + def is_required(self) -> bool: + """Check if the field is required (i.e., does not have a default value or factory). + + Returns: + `True` if the field is required, `False` otherwise. + """ + return self.default is PydanticUndefined and self.default_factory is None + + def rebuild_annotation(self) -> Any: + """Attempts to rebuild the original annotation for use in function signatures. + + If metadata is present, it adds it to the original annotation using + `Annotated`. Otherwise, it returns the original annotation as-is. + + Note that because the metadata has been flattened, the original annotation + may not be reconstructed exactly as originally provided, e.g. if the original + type had unrecognized annotations, or was annotated with a call to `pydantic.Field`. + + Returns: + The rebuilt annotation. + """ + if not self.metadata: + return self.annotation + else: + # Annotated arguments must be a tuple + return Annotated[(self.annotation, *self.metadata)] # type: ignore + + def apply_typevars_map( + self, + typevars_map: Mapping[TypeVar, Any] | None, + globalns: GlobalsNamespace | None = None, + localns: MappingNamespace | None = None, + ) -> None: + """Apply a `typevars_map` to the annotation. + + This method is used when analyzing parametrized generic types to replace typevars with their concrete types. + + This method applies the `typevars_map` to the annotation in place. + + Args: + typevars_map: A dictionary mapping type variables to their concrete types. + globalns: The globals namespace to use during type annotation evaluation. + localns: The locals namespace to use during type annotation evaluation. + + See Also: + pydantic._internal._generics.replace_types is used for replacing the typevars with + their concrete types. + """ + annotation = _generics.replace_types(self.annotation, typevars_map) + annotation, evaluated = _typing_extra.try_eval_type(annotation, globalns, localns) + self.annotation = annotation + if not evaluated: + self._complete = False + self._original_annotation = self.annotation + + def __repr_args__(self) -> ReprArgs: + yield 'annotation', _repr.PlainRepr(_repr.display_as_type(self.annotation)) + yield 'required', self.is_required() + + for s in self.__slots__: + # TODO: properly make use of the protocol (https://rich.readthedocs.io/en/stable/pretty.html#rich-repr-protocol) + # By yielding a three-tuple: + if s in ( + 'annotation', + '_attributes_set', + '_qualifiers', + '_complete', + '_original_assignment', + '_original_annotation', + ): + continue + elif s == 'metadata' and not self.metadata: + continue + elif s == 'repr' and self.repr is True: + continue + if s == 'frozen' and self.frozen is False: + continue + if s == 'validation_alias' and self.validation_alias == self.alias: + continue + if s == 'serialization_alias' and self.serialization_alias == self.alias: + continue + if s == 'default' and self.default is not PydanticUndefined: + yield 'default', self.default + elif s == 'default_factory' and self.default_factory is not None: + yield 'default_factory', _repr.PlainRepr(_repr.display_as_type(self.default_factory)) + else: + value = getattr(self, s) + if value is not None and value is not PydanticUndefined: + yield s, value + + +class _EmptyKwargs(typing_extensions.TypedDict): + """This class exists solely to ensure that type checking warns about passing `**extra` in `Field`.""" + + +_DefaultValues = { + 'default': ..., + 'default_factory': None, + 'alias': None, + 'alias_priority': None, + 'validation_alias': None, + 'serialization_alias': None, + 'title': None, + 'description': None, + 'examples': None, + 'exclude': None, + 'discriminator': None, + 'json_schema_extra': None, + 'frozen': None, + 'validate_default': None, + 'repr': True, + 'init': None, + 'init_var': None, + 'kw_only': None, + 'pattern': None, + 'strict': None, + 'gt': None, + 'ge': None, + 'lt': None, + 'le': None, + 'multiple_of': None, + 'allow_inf_nan': None, + 'max_digits': None, + 'decimal_places': None, + 'min_length': None, + 'max_length': None, + 'coerce_numbers_to_str': None, +} + + +_T = TypeVar('_T') + + +# NOTE: Actual return type is 'FieldInfo', but we want to help type checkers +# to understand the magic that happens at runtime with the following overloads: +@overload # type hint the return value as `Any` to avoid type checking regressions when using `...`. +def Field( + default: ellipsis, # noqa: F821 # TODO: use `_typing_extra.EllipsisType` when we drop Py3.9 + *, + alias: str | None = _Unset, + alias_priority: int | None = _Unset, + validation_alias: str | AliasPath | AliasChoices | None = _Unset, + serialization_alias: str | None = _Unset, + title: str | None = _Unset, + field_title_generator: Callable[[str, FieldInfo], str] | None = _Unset, + description: str | None = _Unset, + examples: list[Any] | None = _Unset, + exclude: bool | None = _Unset, + discriminator: str | types.Discriminator | None = _Unset, + deprecated: Deprecated | str | bool | None = _Unset, + json_schema_extra: JsonDict | Callable[[JsonDict], None] | None = _Unset, + frozen: bool | None = _Unset, + validate_default: bool | None = _Unset, + repr: bool = _Unset, + init: bool | None = _Unset, + init_var: bool | None = _Unset, + kw_only: bool | None = _Unset, + pattern: str | typing.Pattern[str] | None = _Unset, + strict: bool | None = _Unset, + coerce_numbers_to_str: bool | None = _Unset, + gt: annotated_types.SupportsGt | None = _Unset, + ge: annotated_types.SupportsGe | None = _Unset, + lt: annotated_types.SupportsLt | None = _Unset, + le: annotated_types.SupportsLe | None = _Unset, + multiple_of: float | None = _Unset, + allow_inf_nan: bool | None = _Unset, + max_digits: int | None = _Unset, + decimal_places: int | None = _Unset, + min_length: int | None = _Unset, + max_length: int | None = _Unset, + union_mode: Literal['smart', 'left_to_right'] = _Unset, + fail_fast: bool | None = _Unset, + **extra: Unpack[_EmptyKwargs], +) -> Any: ... +@overload # `default` argument set +def Field( + default: _T, + *, + alias: str | None = _Unset, + alias_priority: int | None = _Unset, + validation_alias: str | AliasPath | AliasChoices | None = _Unset, + serialization_alias: str | None = _Unset, + title: str | None = _Unset, + field_title_generator: Callable[[str, FieldInfo], str] | None = _Unset, + description: str | None = _Unset, + examples: list[Any] | None = _Unset, + exclude: bool | None = _Unset, + discriminator: str | types.Discriminator | None = _Unset, + deprecated: Deprecated | str | bool | None = _Unset, + json_schema_extra: JsonDict | Callable[[JsonDict], None] | None = _Unset, + frozen: bool | None = _Unset, + validate_default: bool | None = _Unset, + repr: bool = _Unset, + init: bool | None = _Unset, + init_var: bool | None = _Unset, + kw_only: bool | None = _Unset, + pattern: str | typing.Pattern[str] | None = _Unset, + strict: bool | None = _Unset, + coerce_numbers_to_str: bool | None = _Unset, + gt: annotated_types.SupportsGt | None = _Unset, + ge: annotated_types.SupportsGe | None = _Unset, + lt: annotated_types.SupportsLt | None = _Unset, + le: annotated_types.SupportsLe | None = _Unset, + multiple_of: float | None = _Unset, + allow_inf_nan: bool | None = _Unset, + max_digits: int | None = _Unset, + decimal_places: int | None = _Unset, + min_length: int | None = _Unset, + max_length: int | None = _Unset, + union_mode: Literal['smart', 'left_to_right'] = _Unset, + fail_fast: bool | None = _Unset, + **extra: Unpack[_EmptyKwargs], +) -> _T: ... +@overload # `default_factory` argument set +def Field( + *, + default_factory: Callable[[], _T] | Callable[[dict[str, Any]], _T], + alias: str | None = _Unset, + alias_priority: int | None = _Unset, + validation_alias: str | AliasPath | AliasChoices | None = _Unset, + serialization_alias: str | None = _Unset, + title: str | None = _Unset, + field_title_generator: Callable[[str, FieldInfo], str] | None = _Unset, + description: str | None = _Unset, + examples: list[Any] | None = _Unset, + exclude: bool | None = _Unset, + discriminator: str | types.Discriminator | None = _Unset, + deprecated: Deprecated | str | bool | None = _Unset, + json_schema_extra: JsonDict | Callable[[JsonDict], None] | None = _Unset, + frozen: bool | None = _Unset, + validate_default: bool | None = _Unset, + repr: bool = _Unset, + init: bool | None = _Unset, + init_var: bool | None = _Unset, + kw_only: bool | None = _Unset, + pattern: str | typing.Pattern[str] | None = _Unset, + strict: bool | None = _Unset, + coerce_numbers_to_str: bool | None = _Unset, + gt: annotated_types.SupportsGt | None = _Unset, + ge: annotated_types.SupportsGe | None = _Unset, + lt: annotated_types.SupportsLt | None = _Unset, + le: annotated_types.SupportsLe | None = _Unset, + multiple_of: float | None = _Unset, + allow_inf_nan: bool | None = _Unset, + max_digits: int | None = _Unset, + decimal_places: int | None = _Unset, + min_length: int | None = _Unset, + max_length: int | None = _Unset, + union_mode: Literal['smart', 'left_to_right'] = _Unset, + fail_fast: bool | None = _Unset, + **extra: Unpack[_EmptyKwargs], +) -> _T: ... +@overload +def Field( # No default set + *, + alias: str | None = _Unset, + alias_priority: int | None = _Unset, + validation_alias: str | AliasPath | AliasChoices | None = _Unset, + serialization_alias: str | None = _Unset, + title: str | None = _Unset, + field_title_generator: Callable[[str, FieldInfo], str] | None = _Unset, + description: str | None = _Unset, + examples: list[Any] | None = _Unset, + exclude: bool | None = _Unset, + discriminator: str | types.Discriminator | None = _Unset, + deprecated: Deprecated | str | bool | None = _Unset, + json_schema_extra: JsonDict | Callable[[JsonDict], None] | None = _Unset, + frozen: bool | None = _Unset, + validate_default: bool | None = _Unset, + repr: bool = _Unset, + init: bool | None = _Unset, + init_var: bool | None = _Unset, + kw_only: bool | None = _Unset, + pattern: str | typing.Pattern[str] | None = _Unset, + strict: bool | None = _Unset, + coerce_numbers_to_str: bool | None = _Unset, + gt: annotated_types.SupportsGt | None = _Unset, + ge: annotated_types.SupportsGe | None = _Unset, + lt: annotated_types.SupportsLt | None = _Unset, + le: annotated_types.SupportsLe | None = _Unset, + multiple_of: float | None = _Unset, + allow_inf_nan: bool | None = _Unset, + max_digits: int | None = _Unset, + decimal_places: int | None = _Unset, + min_length: int | None = _Unset, + max_length: int | None = _Unset, + union_mode: Literal['smart', 'left_to_right'] = _Unset, + fail_fast: bool | None = _Unset, + **extra: Unpack[_EmptyKwargs], +) -> Any: ... +def Field( # noqa: C901 + default: Any = PydanticUndefined, + *, + default_factory: Callable[[], Any] | Callable[[dict[str, Any]], Any] | None = _Unset, + alias: str | None = _Unset, + alias_priority: int | None = _Unset, + validation_alias: str | AliasPath | AliasChoices | None = _Unset, + serialization_alias: str | None = _Unset, + title: str | None = _Unset, + field_title_generator: Callable[[str, FieldInfo], str] | None = _Unset, + description: str | None = _Unset, + examples: list[Any] | None = _Unset, + exclude: bool | None = _Unset, + discriminator: str | types.Discriminator | None = _Unset, + deprecated: Deprecated | str | bool | None = _Unset, + json_schema_extra: JsonDict | Callable[[JsonDict], None] | None = _Unset, + frozen: bool | None = _Unset, + validate_default: bool | None = _Unset, + repr: bool = _Unset, + init: bool | None = _Unset, + init_var: bool | None = _Unset, + kw_only: bool | None = _Unset, + pattern: str | typing.Pattern[str] | None = _Unset, + strict: bool | None = _Unset, + coerce_numbers_to_str: bool | None = _Unset, + gt: annotated_types.SupportsGt | None = _Unset, + ge: annotated_types.SupportsGe | None = _Unset, + lt: annotated_types.SupportsLt | None = _Unset, + le: annotated_types.SupportsLe | None = _Unset, + multiple_of: float | None = _Unset, + allow_inf_nan: bool | None = _Unset, + max_digits: int | None = _Unset, + decimal_places: int | None = _Unset, + min_length: int | None = _Unset, + max_length: int | None = _Unset, + union_mode: Literal['smart', 'left_to_right'] = _Unset, + fail_fast: bool | None = _Unset, + **extra: Unpack[_EmptyKwargs], +) -> Any: + """!!! abstract "Usage Documentation" + [Fields](../concepts/fields.md) + + Create a field for objects that can be configured. + + Used to provide extra information about a field, either for the model schema or complex validation. Some arguments + apply only to number fields (`int`, `float`, `Decimal`) and some apply only to `str`. + + Note: + - Any `_Unset` objects will be replaced by the corresponding value defined in the `_DefaultValues` dictionary. If a key for the `_Unset` object is not found in the `_DefaultValues` dictionary, it will default to `None` + + Args: + default: Default value if the field is not set. + default_factory: A callable to generate the default value. The callable can either take 0 arguments + (in which case it is called as is) or a single argument containing the already validated data. + alias: The name to use for the attribute when validating or serializing by alias. + This is often used for things like converting between snake and camel case. + alias_priority: Priority of the alias. This affects whether an alias generator is used. + validation_alias: Like `alias`, but only affects validation, not serialization. + serialization_alias: Like `alias`, but only affects serialization, not validation. + title: Human-readable title. + field_title_generator: A callable that takes a field name and returns title for it. + description: Human-readable description. + examples: Example values for this field. + exclude: Whether to exclude the field from the model serialization. + discriminator: Field name or Discriminator for discriminating the type in a tagged union. + deprecated: A deprecation message, an instance of `warnings.deprecated` or the `typing_extensions.deprecated` backport, + or a boolean. If `True`, a default deprecation message will be emitted when accessing the field. + json_schema_extra: A dict or callable to provide extra JSON schema properties. + frozen: Whether the field is frozen. If true, attempts to change the value on an instance will raise an error. + validate_default: If `True`, apply validation to the default value every time you create an instance. + Otherwise, for performance reasons, the default value of the field is trusted and not validated. + repr: A boolean indicating whether to include the field in the `__repr__` output. + init: Whether the field should be included in the constructor of the dataclass. + (Only applies to dataclasses.) + init_var: Whether the field should _only_ be included in the constructor of the dataclass. + (Only applies to dataclasses.) + kw_only: Whether the field should be a keyword-only argument in the constructor of the dataclass. + (Only applies to dataclasses.) + coerce_numbers_to_str: Whether to enable coercion of any `Number` type to `str` (not applicable in `strict` mode). + strict: If `True`, strict validation is applied to the field. + See [Strict Mode](../concepts/strict_mode.md) for details. + gt: Greater than. If set, value must be greater than this. Only applicable to numbers. + ge: Greater than or equal. If set, value must be greater than or equal to this. Only applicable to numbers. + lt: Less than. If set, value must be less than this. Only applicable to numbers. + le: Less than or equal. If set, value must be less than or equal to this. Only applicable to numbers. + multiple_of: Value must be a multiple of this. Only applicable to numbers. + min_length: Minimum length for iterables. + max_length: Maximum length for iterables. + pattern: Pattern for strings (a regular expression). + allow_inf_nan: Allow `inf`, `-inf`, `nan`. Only applicable to float and [`Decimal`][decimal.Decimal] numbers. + max_digits: Maximum number of allow digits for strings. + decimal_places: Maximum number of decimal places allowed for numbers. + union_mode: The strategy to apply when validating a union. Can be `smart` (the default), or `left_to_right`. + See [Union Mode](../concepts/unions.md#union-modes) for details. + fail_fast: If `True`, validation will stop on the first error. If `False`, all validation errors will be collected. + This option can be applied only to iterable types (list, tuple, set, and frozenset). + extra: (Deprecated) Extra fields that will be included in the JSON schema. + + !!! warning Deprecated + The `extra` kwargs is deprecated. Use `json_schema_extra` instead. + + Returns: + A new [`FieldInfo`][pydantic.fields.FieldInfo]. The return annotation is `Any` so `Field` can be used on + type-annotated fields without causing a type error. + """ + # Check deprecated and removed params from V1. This logic should eventually be removed. + const = extra.pop('const', None) # type: ignore + if const is not None: + raise PydanticUserError('`const` is removed, use `Literal` instead', code='removed-kwargs') + + min_items = extra.pop('min_items', None) # type: ignore + if min_items is not None: + warn('`min_items` is deprecated and will be removed, use `min_length` instead', DeprecationWarning) + if min_length in (None, _Unset): + min_length = min_items # type: ignore + + max_items = extra.pop('max_items', None) # type: ignore + if max_items is not None: + warn('`max_items` is deprecated and will be removed, use `max_length` instead', DeprecationWarning) + if max_length in (None, _Unset): + max_length = max_items # type: ignore + + unique_items = extra.pop('unique_items', None) # type: ignore + if unique_items is not None: + raise PydanticUserError( + ( + '`unique_items` is removed, use `Set` instead' + '(this feature is discussed in https://github.com/pydantic/pydantic-core/issues/296)' + ), + code='removed-kwargs', + ) + + allow_mutation = extra.pop('allow_mutation', None) # type: ignore + if allow_mutation is not None: + warn('`allow_mutation` is deprecated and will be removed. use `frozen` instead', DeprecationWarning) + if allow_mutation is False: + frozen = True + + regex = extra.pop('regex', None) # type: ignore + if regex is not None: + raise PydanticUserError('`regex` is removed. use `pattern` instead', code='removed-kwargs') + + if extra: + warn( + 'Using extra keyword arguments on `Field` is deprecated and will be removed.' + ' Use `json_schema_extra` instead.' + f' (Extra keys: {", ".join(k.__repr__() for k in extra.keys())})', + DeprecationWarning, + ) + if not json_schema_extra or json_schema_extra is _Unset: + json_schema_extra = extra # type: ignore + + if ( + validation_alias + and validation_alias is not _Unset + and not isinstance(validation_alias, (str, AliasChoices, AliasPath)) + ): + raise TypeError('Invalid `validation_alias` type. it should be `str`, `AliasChoices`, or `AliasPath`') + + if serialization_alias in (_Unset, None) and isinstance(alias, str): + serialization_alias = alias + + if validation_alias in (_Unset, None): + validation_alias = alias + + include = extra.pop('include', None) # type: ignore + if include is not None: + warn('`include` is deprecated and does nothing. It will be removed, use `exclude` instead', DeprecationWarning) + + return FieldInfo.from_field( + default, + default_factory=default_factory, + alias=alias, + alias_priority=alias_priority, + validation_alias=validation_alias, + serialization_alias=serialization_alias, + title=title, + field_title_generator=field_title_generator, + description=description, + examples=examples, + exclude=exclude, + discriminator=discriminator, + deprecated=deprecated, + json_schema_extra=json_schema_extra, + frozen=frozen, + pattern=pattern, + validate_default=validate_default, + repr=repr, + init=init, + init_var=init_var, + kw_only=kw_only, + coerce_numbers_to_str=coerce_numbers_to_str, + strict=strict, + gt=gt, + ge=ge, + lt=lt, + le=le, + multiple_of=multiple_of, + min_length=min_length, + max_length=max_length, + allow_inf_nan=allow_inf_nan, + max_digits=max_digits, + decimal_places=decimal_places, + union_mode=union_mode, + fail_fast=fail_fast, + ) + + +_FIELD_ARG_NAMES = set(inspect.signature(Field).parameters) +_FIELD_ARG_NAMES.remove('extra') # do not include the varkwargs parameter + + +class ModelPrivateAttr(_repr.Representation): + """A descriptor for private attributes in class models. + + !!! warning + You generally shouldn't be creating `ModelPrivateAttr` instances directly, instead use + `pydantic.fields.PrivateAttr`. (This is similar to `FieldInfo` vs. `Field`.) + + Attributes: + default: The default value of the attribute if not provided. + default_factory: A callable function that generates the default value of the + attribute if not provided. + """ + + __slots__ = ('default', 'default_factory') + + def __init__( + self, default: Any = PydanticUndefined, *, default_factory: typing.Callable[[], Any] | None = None + ) -> None: + if default is Ellipsis: + self.default = PydanticUndefined + else: + self.default = default + self.default_factory = default_factory + + if not typing.TYPE_CHECKING: + # We put `__getattr__` in a non-TYPE_CHECKING block because otherwise, mypy allows arbitrary attribute access + + def __getattr__(self, item: str) -> Any: + """This function improves compatibility with custom descriptors by ensuring delegation happens + as expected when the default value of a private attribute is a descriptor. + """ + if item in {'__get__', '__set__', '__delete__'}: + if hasattr(self.default, item): + return getattr(self.default, item) + raise AttributeError(f'{type(self).__name__!r} object has no attribute {item!r}') + + def __set_name__(self, cls: type[Any], name: str) -> None: + """Preserve `__set_name__` protocol defined in https://peps.python.org/pep-0487.""" + default = self.default + if default is PydanticUndefined: + return + set_name = getattr(default, '__set_name__', None) + if callable(set_name): + set_name(cls, name) + + def get_default(self) -> Any: + """Retrieve the default value of the object. + + If `self.default_factory` is `None`, the method will return a deep copy of the `self.default` object. + + If `self.default_factory` is not `None`, it will call `self.default_factory` and return the value returned. + + Returns: + The default value of the object. + """ + return _utils.smart_deepcopy(self.default) if self.default_factory is None else self.default_factory() + + def __eq__(self, other: Any) -> bool: + return isinstance(other, self.__class__) and (self.default, self.default_factory) == ( + other.default, + other.default_factory, + ) + + +# NOTE: Actual return type is 'ModelPrivateAttr', but we want to help type checkers +# to understand the magic that happens at runtime. +@overload # `default` argument set +def PrivateAttr( + default: _T, + *, + init: Literal[False] = False, +) -> _T: ... +@overload # `default_factory` argument set +def PrivateAttr( + *, + default_factory: Callable[[], _T], + init: Literal[False] = False, +) -> _T: ... +@overload # No default set +def PrivateAttr( + *, + init: Literal[False] = False, +) -> Any: ... +def PrivateAttr( + default: Any = PydanticUndefined, + *, + default_factory: Callable[[], Any] | None = None, + init: Literal[False] = False, +) -> Any: + """!!! abstract "Usage Documentation" + [Private Model Attributes](../concepts/models.md#private-model-attributes) + + Indicates that an attribute is intended for private use and not handled during normal validation/serialization. + + Private attributes are not validated by Pydantic, so it's up to you to ensure they are used in a type-safe manner. + + Private attributes are stored in `__private_attributes__` on the model. + + Args: + default: The attribute's default value. Defaults to Undefined. + default_factory: Callable that will be + called when a default value is needed for this attribute. + If both `default` and `default_factory` are set, an error will be raised. + init: Whether the attribute should be included in the constructor of the dataclass. Always `False`. + + Returns: + An instance of [`ModelPrivateAttr`][pydantic.fields.ModelPrivateAttr] class. + + Raises: + ValueError: If both `default` and `default_factory` are set. + """ + if default is not PydanticUndefined and default_factory is not None: + raise TypeError('cannot specify both default and default_factory') + + return ModelPrivateAttr( + default, + default_factory=default_factory, + ) + + +@dataclasses.dataclass(**_internal_dataclass.slots_true) +class ComputedFieldInfo: + """A container for data from `@computed_field` so that we can access it while building the pydantic-core schema. + + Attributes: + decorator_repr: A class variable representing the decorator string, '@computed_field'. + wrapped_property: The wrapped computed field property. + return_type: The type of the computed field property's return value. + alias: The alias of the property to be used during serialization. + alias_priority: The priority of the alias. This affects whether an alias generator is used. + title: Title of the computed field to include in the serialization JSON schema. + field_title_generator: A callable that takes a field name and returns title for it. + description: Description of the computed field to include in the serialization JSON schema. + deprecated: A deprecation message, an instance of `warnings.deprecated` or the `typing_extensions.deprecated` backport, + or a boolean. If `True`, a default deprecation message will be emitted when accessing the field. + examples: Example values of the computed field to include in the serialization JSON schema. + json_schema_extra: A dict or callable to provide extra JSON schema properties. + repr: A boolean indicating whether to include the field in the __repr__ output. + """ + + decorator_repr: ClassVar[str] = '@computed_field' + wrapped_property: property + return_type: Any + alias: str | None + alias_priority: int | None + title: str | None + field_title_generator: typing.Callable[[str, ComputedFieldInfo], str] | None + description: str | None + deprecated: Deprecated | str | bool | None + examples: list[Any] | None + json_schema_extra: JsonDict | typing.Callable[[JsonDict], None] | None + repr: bool + + @property + def deprecation_message(self) -> str | None: + """The deprecation message to be emitted, or `None` if not set.""" + if self.deprecated is None: + return None + if isinstance(self.deprecated, bool): + return 'deprecated' if self.deprecated else None + return self.deprecated if isinstance(self.deprecated, str) else self.deprecated.message + + +def _wrapped_property_is_private(property_: cached_property | property) -> bool: # type: ignore + """Returns true if provided property is private, False otherwise.""" + wrapped_name: str = '' + + if isinstance(property_, property): + wrapped_name = getattr(property_.fget, '__name__', '') + elif isinstance(property_, cached_property): # type: ignore + wrapped_name = getattr(property_.func, '__name__', '') # type: ignore + + return wrapped_name.startswith('_') and not wrapped_name.startswith('__') + + +# this should really be `property[T], cached_property[T]` but property is not generic unlike cached_property +# See https://github.com/python/typing/issues/985 and linked issues +PropertyT = typing.TypeVar('PropertyT') + + +@typing.overload +def computed_field(func: PropertyT, /) -> PropertyT: ... + + +@typing.overload +def computed_field( + *, + alias: str | None = None, + alias_priority: int | None = None, + title: str | None = None, + field_title_generator: typing.Callable[[str, ComputedFieldInfo], str] | None = None, + description: str | None = None, + deprecated: Deprecated | str | bool | None = None, + examples: list[Any] | None = None, + json_schema_extra: JsonDict | typing.Callable[[JsonDict], None] | None = None, + repr: bool = True, + return_type: Any = PydanticUndefined, +) -> typing.Callable[[PropertyT], PropertyT]: ... + + +def computed_field( + func: PropertyT | None = None, + /, + *, + alias: str | None = None, + alias_priority: int | None = None, + title: str | None = None, + field_title_generator: typing.Callable[[str, ComputedFieldInfo], str] | None = None, + description: str | None = None, + deprecated: Deprecated | str | bool | None = None, + examples: list[Any] | None = None, + json_schema_extra: JsonDict | typing.Callable[[JsonDict], None] | None = None, + repr: bool | None = None, + return_type: Any = PydanticUndefined, +) -> PropertyT | typing.Callable[[PropertyT], PropertyT]: + """!!! abstract "Usage Documentation" + [The `computed_field` decorator](../concepts/fields.md#the-computed_field-decorator) + + Decorator to include `property` and `cached_property` when serializing models or dataclasses. + + This is useful for fields that are computed from other fields, or for fields that are expensive to compute and should be cached. + + ```python + from pydantic import BaseModel, computed_field + + class Rectangle(BaseModel): + width: int + length: int + + @computed_field + @property + def area(self) -> int: + return self.width * self.length + + print(Rectangle(width=3, length=2).model_dump()) + #> {'width': 3, 'length': 2, 'area': 6} + ``` + + If applied to functions not yet decorated with `@property` or `@cached_property`, the function is + automatically wrapped with `property`. Although this is more concise, you will lose IntelliSense in your IDE, + and confuse static type checkers, thus explicit use of `@property` is recommended. + + !!! warning "Mypy Warning" + Even with the `@property` or `@cached_property` applied to your function before `@computed_field`, + mypy may throw a `Decorated property not supported` error. + See [mypy issue #1362](https://github.com/python/mypy/issues/1362), for more information. + To avoid this error message, add `# type: ignore[prop-decorator]` to the `@computed_field` line. + + [pyright](https://github.com/microsoft/pyright) supports `@computed_field` without error. + + ```python + import random + + from pydantic import BaseModel, computed_field + + class Square(BaseModel): + width: float + + @computed_field + def area(self) -> float: # converted to a `property` by `computed_field` + return round(self.width**2, 2) + + @area.setter + def area(self, new_area: float) -> None: + self.width = new_area**0.5 + + @computed_field(alias='the magic number', repr=False) + def random_number(self) -> int: + return random.randint(0, 1_000) + + square = Square(width=1.3) + + # `random_number` does not appear in representation + print(repr(square)) + #> Square(width=1.3, area=1.69) + + print(square.random_number) + #> 3 + + square.area = 4 + + print(square.model_dump_json(by_alias=True)) + #> {"width":2.0,"area":4.0,"the magic number":3} + ``` + + !!! warning "Overriding with `computed_field`" + You can't override a field from a parent class with a `computed_field` in the child class. + `mypy` complains about this behavior if allowed, and `dataclasses` doesn't allow this pattern either. + See the example below: + + ```python + from pydantic import BaseModel, computed_field + + class Parent(BaseModel): + a: str + + try: + + class Child(Parent): + @computed_field + @property + def a(self) -> str: + return 'new a' + + except TypeError as e: + print(e) + ''' + Field 'a' of class 'Child' overrides symbol of same name in a parent class. This override with a computed_field is incompatible. + ''' + ``` + + Private properties decorated with `@computed_field` have `repr=False` by default. + + ```python + from functools import cached_property + + from pydantic import BaseModel, computed_field + + class Model(BaseModel): + foo: int + + @computed_field + @cached_property + def _private_cached_property(self) -> int: + return -self.foo + + @computed_field + @property + def _private_property(self) -> int: + return -self.foo + + m = Model(foo=1) + print(repr(m)) + #> Model(foo=1) + ``` + + Args: + func: the function to wrap. + alias: alias to use when serializing this computed field, only used when `by_alias=True` + alias_priority: priority of the alias. This affects whether an alias generator is used + title: Title to use when including this computed field in JSON Schema + field_title_generator: A callable that takes a field name and returns title for it. + description: Description to use when including this computed field in JSON Schema, defaults to the function's + docstring + deprecated: A deprecation message (or an instance of `warnings.deprecated` or the `typing_extensions.deprecated` backport). + to be emitted when accessing the field. Or a boolean. This will automatically be set if the property is decorated with the + `deprecated` decorator. + examples: Example values to use when including this computed field in JSON Schema + json_schema_extra: A dict or callable to provide extra JSON schema properties. + repr: whether to include this computed field in model repr. + Default is `False` for private properties and `True` for public properties. + return_type: optional return for serialization logic to expect when serializing to JSON, if included + this must be correct, otherwise a `TypeError` is raised. + If you don't include a return type Any is used, which does runtime introspection to handle arbitrary + objects. + + Returns: + A proxy wrapper for the property. + """ + + def dec(f: Any) -> Any: + nonlocal description, deprecated, return_type, alias_priority + unwrapped = _decorators.unwrap_wrapped_function(f) + + if description is None and unwrapped.__doc__: + description = inspect.cleandoc(unwrapped.__doc__) + + if deprecated is None and hasattr(unwrapped, '__deprecated__'): + deprecated = unwrapped.__deprecated__ + + # if the function isn't already decorated with `@property` (or another descriptor), then we wrap it now + f = _decorators.ensure_property(f) + alias_priority = (alias_priority or 2) if alias is not None else None + + if repr is None: + repr_: bool = not _wrapped_property_is_private(property_=f) + else: + repr_ = repr + + dec_info = ComputedFieldInfo( + f, + return_type, + alias, + alias_priority, + title, + field_title_generator, + description, + deprecated, + examples, + json_schema_extra, + repr_, + ) + return _decorators.PydanticDescriptorProxy(f, dec_info) + + if func is None: + return dec + else: + return dec(func) diff --git a/lib/python3.12/site-packages/pydantic/functional_serializers.py b/lib/python3.12/site-packages/pydantic/functional_serializers.py new file mode 100644 index 0000000000000000000000000000000000000000..4b065e49a19d6023cb11bc55870d7db956e5f6e3 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/functional_serializers.py @@ -0,0 +1,450 @@ +"""This module contains related classes and functions for serialization.""" + +from __future__ import annotations + +import dataclasses +from functools import partial, partialmethod +from typing import TYPE_CHECKING, Annotated, Any, Callable, Literal, TypeVar, overload + +from pydantic_core import PydanticUndefined, core_schema +from pydantic_core.core_schema import SerializationInfo, SerializerFunctionWrapHandler, WhenUsed +from typing_extensions import TypeAlias + +from . import PydanticUndefinedAnnotation +from ._internal import _decorators, _internal_dataclass +from .annotated_handlers import GetCoreSchemaHandler + + +@dataclasses.dataclass(**_internal_dataclass.slots_true, frozen=True) +class PlainSerializer: + """Plain serializers use a function to modify the output of serialization. + + This is particularly helpful when you want to customize the serialization for annotated types. + Consider an input of `list`, which will be serialized into a space-delimited string. + + ```python + from typing import Annotated + + from pydantic import BaseModel, PlainSerializer + + CustomStr = Annotated[ + list, PlainSerializer(lambda x: ' '.join(x), return_type=str) + ] + + class StudentModel(BaseModel): + courses: CustomStr + + student = StudentModel(courses=['Math', 'Chemistry', 'English']) + print(student.model_dump()) + #> {'courses': 'Math Chemistry English'} + ``` + + Attributes: + func: The serializer function. + return_type: The return type for the function. If omitted it will be inferred from the type annotation. + when_used: Determines when this serializer should be used. Accepts a string with values `'always'`, + `'unless-none'`, `'json'`, and `'json-unless-none'`. Defaults to 'always'. + """ + + func: core_schema.SerializerFunction + return_type: Any = PydanticUndefined + when_used: WhenUsed = 'always' + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + """Gets the Pydantic core schema. + + Args: + source_type: The source type. + handler: The `GetCoreSchemaHandler` instance. + + Returns: + The Pydantic core schema. + """ + schema = handler(source_type) + if self.return_type is not PydanticUndefined: + return_type = self.return_type + else: + try: + # Do not pass in globals as the function could be defined in a different module. + # Instead, let `get_callable_return_type` infer the globals to use, but still pass + # in locals that may contain a parent/rebuild namespace: + return_type = _decorators.get_callable_return_type( + self.func, + localns=handler._get_types_namespace().locals, + ) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + + return_schema = None if return_type is PydanticUndefined else handler.generate_schema(return_type) + schema['serialization'] = core_schema.plain_serializer_function_ser_schema( + function=self.func, + info_arg=_decorators.inspect_annotated_serializer(self.func, 'plain'), + return_schema=return_schema, + when_used=self.when_used, + ) + return schema + + +@dataclasses.dataclass(**_internal_dataclass.slots_true, frozen=True) +class WrapSerializer: + """Wrap serializers receive the raw inputs along with a handler function that applies the standard serialization + logic, and can modify the resulting value before returning it as the final output of serialization. + + For example, here's a scenario in which a wrap serializer transforms timezones to UTC **and** utilizes the existing `datetime` serialization logic. + + ```python + from datetime import datetime, timezone + from typing import Annotated, Any + + from pydantic import BaseModel, WrapSerializer + + class EventDatetime(BaseModel): + start: datetime + end: datetime + + def convert_to_utc(value: Any, handler, info) -> dict[str, datetime]: + # Note that `handler` can actually help serialize the `value` for + # further custom serialization in case it's a subclass. + partial_result = handler(value, info) + if info.mode == 'json': + return { + k: datetime.fromisoformat(v).astimezone(timezone.utc) + for k, v in partial_result.items() + } + return {k: v.astimezone(timezone.utc) for k, v in partial_result.items()} + + UTCEventDatetime = Annotated[EventDatetime, WrapSerializer(convert_to_utc)] + + class EventModel(BaseModel): + event_datetime: UTCEventDatetime + + dt = EventDatetime( + start='2024-01-01T07:00:00-08:00', end='2024-01-03T20:00:00+06:00' + ) + event = EventModel(event_datetime=dt) + print(event.model_dump()) + ''' + { + 'event_datetime': { + 'start': datetime.datetime( + 2024, 1, 1, 15, 0, tzinfo=datetime.timezone.utc + ), + 'end': datetime.datetime( + 2024, 1, 3, 14, 0, tzinfo=datetime.timezone.utc + ), + } + } + ''' + + print(event.model_dump_json()) + ''' + {"event_datetime":{"start":"2024-01-01T15:00:00Z","end":"2024-01-03T14:00:00Z"}} + ''' + ``` + + Attributes: + func: The serializer function to be wrapped. + return_type: The return type for the function. If omitted it will be inferred from the type annotation. + when_used: Determines when this serializer should be used. Accepts a string with values `'always'`, + `'unless-none'`, `'json'`, and `'json-unless-none'`. Defaults to 'always'. + """ + + func: core_schema.WrapSerializerFunction + return_type: Any = PydanticUndefined + when_used: WhenUsed = 'always' + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + """This method is used to get the Pydantic core schema of the class. + + Args: + source_type: Source type. + handler: Core schema handler. + + Returns: + The generated core schema of the class. + """ + schema = handler(source_type) + if self.return_type is not PydanticUndefined: + return_type = self.return_type + else: + try: + # Do not pass in globals as the function could be defined in a different module. + # Instead, let `get_callable_return_type` infer the globals to use, but still pass + # in locals that may contain a parent/rebuild namespace: + return_type = _decorators.get_callable_return_type( + self.func, + localns=handler._get_types_namespace().locals, + ) + except NameError as e: + raise PydanticUndefinedAnnotation.from_name_error(e) from e + + return_schema = None if return_type is PydanticUndefined else handler.generate_schema(return_type) + schema['serialization'] = core_schema.wrap_serializer_function_ser_schema( + function=self.func, + info_arg=_decorators.inspect_annotated_serializer(self.func, 'wrap'), + return_schema=return_schema, + when_used=self.when_used, + ) + return schema + + +if TYPE_CHECKING: + _Partial: TypeAlias = 'partial[Any] | partialmethod[Any]' + + FieldPlainSerializer: TypeAlias = 'core_schema.SerializerFunction | _Partial' + """A field serializer method or function in `plain` mode.""" + + FieldWrapSerializer: TypeAlias = 'core_schema.WrapSerializerFunction | _Partial' + """A field serializer method or function in `wrap` mode.""" + + FieldSerializer: TypeAlias = 'FieldPlainSerializer | FieldWrapSerializer' + """A field serializer method or function.""" + + _FieldPlainSerializerT = TypeVar('_FieldPlainSerializerT', bound=FieldPlainSerializer) + _FieldWrapSerializerT = TypeVar('_FieldWrapSerializerT', bound=FieldWrapSerializer) + + +@overload +def field_serializer( + field: str, + /, + *fields: str, + mode: Literal['wrap'], + return_type: Any = ..., + when_used: WhenUsed = ..., + check_fields: bool | None = ..., +) -> Callable[[_FieldWrapSerializerT], _FieldWrapSerializerT]: ... + + +@overload +def field_serializer( + field: str, + /, + *fields: str, + mode: Literal['plain'] = ..., + return_type: Any = ..., + when_used: WhenUsed = ..., + check_fields: bool | None = ..., +) -> Callable[[_FieldPlainSerializerT], _FieldPlainSerializerT]: ... + + +def field_serializer( + *fields: str, + mode: Literal['plain', 'wrap'] = 'plain', + return_type: Any = PydanticUndefined, + when_used: WhenUsed = 'always', + check_fields: bool | None = None, +) -> ( + Callable[[_FieldWrapSerializerT], _FieldWrapSerializerT] + | Callable[[_FieldPlainSerializerT], _FieldPlainSerializerT] +): + """Decorator that enables custom field serialization. + + In the below example, a field of type `set` is used to mitigate duplication. A `field_serializer` is used to serialize the data as a sorted list. + + ```python + from typing import Set + + from pydantic import BaseModel, field_serializer + + class StudentModel(BaseModel): + name: str = 'Jane' + courses: Set[str] + + @field_serializer('courses', when_used='json') + def serialize_courses_in_order(self, courses: Set[str]): + return sorted(courses) + + student = StudentModel(courses={'Math', 'Chemistry', 'English'}) + print(student.model_dump_json()) + #> {"name":"Jane","courses":["Chemistry","English","Math"]} + ``` + + See [Custom serializers](../concepts/serialization.md#custom-serializers) for more information. + + Four signatures are supported: + + - `(self, value: Any, info: FieldSerializationInfo)` + - `(self, value: Any, nxt: SerializerFunctionWrapHandler, info: FieldSerializationInfo)` + - `(value: Any, info: SerializationInfo)` + - `(value: Any, nxt: SerializerFunctionWrapHandler, info: SerializationInfo)` + + Args: + fields: Which field(s) the method should be called on. + mode: The serialization mode. + + - `plain` means the function will be called instead of the default serialization logic, + - `wrap` means the function will be called with an argument to optionally call the + default serialization logic. + return_type: Optional return type for the function, if omitted it will be inferred from the type annotation. + when_used: Determines the serializer will be used for serialization. + check_fields: Whether to check that the fields actually exist on the model. + + Returns: + The decorator function. + """ + + def dec(f: FieldSerializer) -> _decorators.PydanticDescriptorProxy[Any]: + dec_info = _decorators.FieldSerializerDecoratorInfo( + fields=fields, + mode=mode, + return_type=return_type, + when_used=when_used, + check_fields=check_fields, + ) + return _decorators.PydanticDescriptorProxy(f, dec_info) # pyright: ignore[reportArgumentType] + + return dec # pyright: ignore[reportReturnType] + + +if TYPE_CHECKING: + # The first argument in the following callables represent the `self` type: + + ModelPlainSerializerWithInfo: TypeAlias = Callable[[Any, SerializationInfo], Any] + """A model serializer method with the `info` argument, in `plain` mode.""" + + ModelPlainSerializerWithoutInfo: TypeAlias = Callable[[Any], Any] + """A model serializer method without the `info` argument, in `plain` mode.""" + + ModelPlainSerializer: TypeAlias = 'ModelPlainSerializerWithInfo | ModelPlainSerializerWithoutInfo' + """A model serializer method in `plain` mode.""" + + ModelWrapSerializerWithInfo: TypeAlias = Callable[[Any, SerializerFunctionWrapHandler, SerializationInfo], Any] + """A model serializer method with the `info` argument, in `wrap` mode.""" + + ModelWrapSerializerWithoutInfo: TypeAlias = Callable[[Any, SerializerFunctionWrapHandler], Any] + """A model serializer method without the `info` argument, in `wrap` mode.""" + + ModelWrapSerializer: TypeAlias = 'ModelWrapSerializerWithInfo | ModelWrapSerializerWithoutInfo' + """A model serializer method in `wrap` mode.""" + + ModelSerializer: TypeAlias = 'ModelPlainSerializer | ModelWrapSerializer' + + _ModelPlainSerializerT = TypeVar('_ModelPlainSerializerT', bound=ModelPlainSerializer) + _ModelWrapSerializerT = TypeVar('_ModelWrapSerializerT', bound=ModelWrapSerializer) + + +@overload +def model_serializer(f: _ModelPlainSerializerT, /) -> _ModelPlainSerializerT: ... + + +@overload +def model_serializer( + *, mode: Literal['wrap'], when_used: WhenUsed = 'always', return_type: Any = ... +) -> Callable[[_ModelWrapSerializerT], _ModelWrapSerializerT]: ... + + +@overload +def model_serializer( + *, + mode: Literal['plain'] = ..., + when_used: WhenUsed = 'always', + return_type: Any = ..., +) -> Callable[[_ModelPlainSerializerT], _ModelPlainSerializerT]: ... + + +def model_serializer( + f: _ModelPlainSerializerT | _ModelWrapSerializerT | None = None, + /, + *, + mode: Literal['plain', 'wrap'] = 'plain', + when_used: WhenUsed = 'always', + return_type: Any = PydanticUndefined, +) -> ( + _ModelPlainSerializerT + | Callable[[_ModelWrapSerializerT], _ModelWrapSerializerT] + | Callable[[_ModelPlainSerializerT], _ModelPlainSerializerT] +): + """Decorator that enables custom model serialization. + + This is useful when a model need to be serialized in a customized manner, allowing for flexibility beyond just specific fields. + + An example would be to serialize temperature to the same temperature scale, such as degrees Celsius. + + ```python + from typing import Literal + + from pydantic import BaseModel, model_serializer + + class TemperatureModel(BaseModel): + unit: Literal['C', 'F'] + value: int + + @model_serializer() + def serialize_model(self): + if self.unit == 'F': + return {'unit': 'C', 'value': int((self.value - 32) / 1.8)} + return {'unit': self.unit, 'value': self.value} + + temperature = TemperatureModel(unit='F', value=212) + print(temperature.model_dump()) + #> {'unit': 'C', 'value': 100} + ``` + + Two signatures are supported for `mode='plain'`, which is the default: + + - `(self)` + - `(self, info: SerializationInfo)` + + And two other signatures for `mode='wrap'`: + + - `(self, nxt: SerializerFunctionWrapHandler)` + - `(self, nxt: SerializerFunctionWrapHandler, info: SerializationInfo)` + + See [Custom serializers](../concepts/serialization.md#custom-serializers) for more information. + + Args: + f: The function to be decorated. + mode: The serialization mode. + + - `'plain'` means the function will be called instead of the default serialization logic + - `'wrap'` means the function will be called with an argument to optionally call the default + serialization logic. + when_used: Determines when this serializer should be used. + return_type: The return type for the function. If omitted it will be inferred from the type annotation. + + Returns: + The decorator function. + """ + + def dec(f: ModelSerializer) -> _decorators.PydanticDescriptorProxy[Any]: + dec_info = _decorators.ModelSerializerDecoratorInfo(mode=mode, return_type=return_type, when_used=when_used) + return _decorators.PydanticDescriptorProxy(f, dec_info) + + if f is None: + return dec # pyright: ignore[reportReturnType] + else: + return dec(f) # pyright: ignore[reportReturnType] + + +AnyType = TypeVar('AnyType') + + +if TYPE_CHECKING: + SerializeAsAny = Annotated[AnyType, ...] # SerializeAsAny[list[str]] will be treated by type checkers as list[str] + """Force serialization to ignore whatever is defined in the schema and instead ask the object + itself how it should be serialized. + In particular, this means that when model subclasses are serialized, fields present in the subclass + but not in the original schema will be included. + """ +else: + + @dataclasses.dataclass(**_internal_dataclass.slots_true) + class SerializeAsAny: # noqa: D101 + def __class_getitem__(cls, item: Any) -> Any: + return Annotated[item, SerializeAsAny()] + + def __get_pydantic_core_schema__( + self, source_type: Any, handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + schema = handler(source_type) + schema_to_update = schema + while schema_to_update['type'] == 'definitions': + schema_to_update = schema_to_update.copy() + schema_to_update = schema_to_update['schema'] + schema_to_update['serialization'] = core_schema.wrap_serializer_function_ser_schema( + lambda x, h: h(x), schema=core_schema.any_schema() + ) + return schema + + __hash__ = object.__hash__ diff --git a/lib/python3.12/site-packages/pydantic/functional_validators.py b/lib/python3.12/site-packages/pydantic/functional_validators.py new file mode 100644 index 0000000000000000000000000000000000000000..2eed4ef8fc5fdbe738297ca030d8774dd2e1cb10 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/functional_validators.py @@ -0,0 +1,828 @@ +"""This module contains related classes and functions for validation.""" + +from __future__ import annotations as _annotations + +import dataclasses +import sys +from functools import partialmethod +from types import FunctionType +from typing import TYPE_CHECKING, Annotated, Any, Callable, Literal, TypeVar, Union, cast, overload + +from pydantic_core import PydanticUndefined, core_schema +from pydantic_core import core_schema as _core_schema +from typing_extensions import Self, TypeAlias + +from ._internal import _decorators, _generics, _internal_dataclass +from .annotated_handlers import GetCoreSchemaHandler +from .errors import PydanticUserError + +if sys.version_info < (3, 11): + from typing_extensions import Protocol +else: + from typing import Protocol + +_inspect_validator = _decorators.inspect_validator + + +@dataclasses.dataclass(frozen=True, **_internal_dataclass.slots_true) +class AfterValidator: + """!!! abstract "Usage Documentation" + [field *after* validators](../concepts/validators.md#field-after-validator) + + A metadata class that indicates that a validation should be applied **after** the inner validation logic. + + Attributes: + func: The validator function. + + Example: + ```python + from typing import Annotated + + from pydantic import AfterValidator, BaseModel, ValidationError + + MyInt = Annotated[int, AfterValidator(lambda v: v + 1)] + + class Model(BaseModel): + a: MyInt + + print(Model(a=1).a) + #> 2 + + try: + Model(a='a') + except ValidationError as e: + print(e.json(indent=2)) + ''' + [ + { + "type": "int_parsing", + "loc": [ + "a" + ], + "msg": "Input should be a valid integer, unable to parse string as an integer", + "input": "a", + "url": "https://errors.pydantic.dev/2/v/int_parsing" + } + ] + ''' + ``` + """ + + func: core_schema.NoInfoValidatorFunction | core_schema.WithInfoValidatorFunction + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + schema = handler(source_type) + info_arg = _inspect_validator(self.func, 'after') + if info_arg: + func = cast(core_schema.WithInfoValidatorFunction, self.func) + return core_schema.with_info_after_validator_function(func, schema=schema, field_name=handler.field_name) + else: + func = cast(core_schema.NoInfoValidatorFunction, self.func) + return core_schema.no_info_after_validator_function(func, schema=schema) + + @classmethod + def _from_decorator(cls, decorator: _decorators.Decorator[_decorators.FieldValidatorDecoratorInfo]) -> Self: + return cls(func=decorator.func) + + +@dataclasses.dataclass(frozen=True, **_internal_dataclass.slots_true) +class BeforeValidator: + """!!! abstract "Usage Documentation" + [field *before* validators](../concepts/validators.md#field-before-validator) + + A metadata class that indicates that a validation should be applied **before** the inner validation logic. + + Attributes: + func: The validator function. + json_schema_input_type: The input type of the function. This is only used to generate the appropriate + JSON Schema (in validation mode). + + Example: + ```python + from typing import Annotated + + from pydantic import BaseModel, BeforeValidator + + MyInt = Annotated[int, BeforeValidator(lambda v: v + 1)] + + class Model(BaseModel): + a: MyInt + + print(Model(a=1).a) + #> 2 + + try: + Model(a='a') + except TypeError as e: + print(e) + #> can only concatenate str (not "int") to str + ``` + """ + + func: core_schema.NoInfoValidatorFunction | core_schema.WithInfoValidatorFunction + json_schema_input_type: Any = PydanticUndefined + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + schema = handler(source_type) + input_schema = ( + None + if self.json_schema_input_type is PydanticUndefined + else handler.generate_schema(self.json_schema_input_type) + ) + + info_arg = _inspect_validator(self.func, 'before') + if info_arg: + func = cast(core_schema.WithInfoValidatorFunction, self.func) + return core_schema.with_info_before_validator_function( + func, + schema=schema, + field_name=handler.field_name, + json_schema_input_schema=input_schema, + ) + else: + func = cast(core_schema.NoInfoValidatorFunction, self.func) + return core_schema.no_info_before_validator_function( + func, schema=schema, json_schema_input_schema=input_schema + ) + + @classmethod + def _from_decorator(cls, decorator: _decorators.Decorator[_decorators.FieldValidatorDecoratorInfo]) -> Self: + return cls( + func=decorator.func, + json_schema_input_type=decorator.info.json_schema_input_type, + ) + + +@dataclasses.dataclass(frozen=True, **_internal_dataclass.slots_true) +class PlainValidator: + """!!! abstract "Usage Documentation" + [field *plain* validators](../concepts/validators.md#field-plain-validator) + + A metadata class that indicates that a validation should be applied **instead** of the inner validation logic. + + !!! note + Before v2.9, `PlainValidator` wasn't always compatible with JSON Schema generation for `mode='validation'`. + You can now use the `json_schema_input_type` argument to specify the input type of the function + to be used in the JSON schema when `mode='validation'` (the default). See the example below for more details. + + Attributes: + func: The validator function. + json_schema_input_type: The input type of the function. This is only used to generate the appropriate + JSON Schema (in validation mode). If not provided, will default to `Any`. + + Example: + ```python + from typing import Annotated, Union + + from pydantic import BaseModel, PlainValidator + + MyInt = Annotated[ + int, + PlainValidator( + lambda v: int(v) + 1, json_schema_input_type=Union[str, int] # (1)! + ), + ] + + class Model(BaseModel): + a: MyInt + + print(Model(a='1').a) + #> 2 + + print(Model(a=1).a) + #> 2 + ``` + + 1. In this example, we've specified the `json_schema_input_type` as `Union[str, int]` which indicates to the JSON schema + generator that in validation mode, the input type for the `a` field can be either a `str` or an `int`. + """ + + func: core_schema.NoInfoValidatorFunction | core_schema.WithInfoValidatorFunction + json_schema_input_type: Any = Any + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + # Note that for some valid uses of PlainValidator, it is not possible to generate a core schema for the + # source_type, so calling `handler(source_type)` will error, which prevents us from generating a proper + # serialization schema. To work around this for use cases that will not involve serialization, we simply + # catch any PydanticSchemaGenerationError that may be raised while attempting to build the serialization schema + # and abort any attempts to handle special serialization. + from pydantic import PydanticSchemaGenerationError + + try: + schema = handler(source_type) + # TODO if `schema['serialization']` is one of `'include-exclude-dict/sequence', + # schema validation will fail. That's why we use 'type ignore' comments below. + serialization = schema.get( + 'serialization', + core_schema.wrap_serializer_function_ser_schema( + function=lambda v, h: h(v), + schema=schema, + return_schema=handler.generate_schema(source_type), + ), + ) + except PydanticSchemaGenerationError: + serialization = None + + input_schema = handler.generate_schema(self.json_schema_input_type) + + info_arg = _inspect_validator(self.func, 'plain') + if info_arg: + func = cast(core_schema.WithInfoValidatorFunction, self.func) + return core_schema.with_info_plain_validator_function( + func, + field_name=handler.field_name, + serialization=serialization, # pyright: ignore[reportArgumentType] + json_schema_input_schema=input_schema, + ) + else: + func = cast(core_schema.NoInfoValidatorFunction, self.func) + return core_schema.no_info_plain_validator_function( + func, + serialization=serialization, # pyright: ignore[reportArgumentType] + json_schema_input_schema=input_schema, + ) + + @classmethod + def _from_decorator(cls, decorator: _decorators.Decorator[_decorators.FieldValidatorDecoratorInfo]) -> Self: + return cls( + func=decorator.func, + json_schema_input_type=decorator.info.json_schema_input_type, + ) + + +@dataclasses.dataclass(frozen=True, **_internal_dataclass.slots_true) +class WrapValidator: + """!!! abstract "Usage Documentation" + [field *wrap* validators](../concepts/validators.md#field-wrap-validator) + + A metadata class that indicates that a validation should be applied **around** the inner validation logic. + + Attributes: + func: The validator function. + json_schema_input_type: The input type of the function. This is only used to generate the appropriate + JSON Schema (in validation mode). + + ```python + from datetime import datetime + from typing import Annotated + + from pydantic import BaseModel, ValidationError, WrapValidator + + def validate_timestamp(v, handler): + if v == 'now': + # we don't want to bother with further validation, just return the new value + return datetime.now() + try: + return handler(v) + except ValidationError: + # validation failed, in this case we want to return a default value + return datetime(2000, 1, 1) + + MyTimestamp = Annotated[datetime, WrapValidator(validate_timestamp)] + + class Model(BaseModel): + a: MyTimestamp + + print(Model(a='now').a) + #> 2032-01-02 03:04:05.000006 + print(Model(a='invalid').a) + #> 2000-01-01 00:00:00 + ``` + """ + + func: core_schema.NoInfoWrapValidatorFunction | core_schema.WithInfoWrapValidatorFunction + json_schema_input_type: Any = PydanticUndefined + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + schema = handler(source_type) + input_schema = ( + None + if self.json_schema_input_type is PydanticUndefined + else handler.generate_schema(self.json_schema_input_type) + ) + + info_arg = _inspect_validator(self.func, 'wrap') + if info_arg: + func = cast(core_schema.WithInfoWrapValidatorFunction, self.func) + return core_schema.with_info_wrap_validator_function( + func, + schema=schema, + field_name=handler.field_name, + json_schema_input_schema=input_schema, + ) + else: + func = cast(core_schema.NoInfoWrapValidatorFunction, self.func) + return core_schema.no_info_wrap_validator_function( + func, + schema=schema, + json_schema_input_schema=input_schema, + ) + + @classmethod + def _from_decorator(cls, decorator: _decorators.Decorator[_decorators.FieldValidatorDecoratorInfo]) -> Self: + return cls( + func=decorator.func, + json_schema_input_type=decorator.info.json_schema_input_type, + ) + + +if TYPE_CHECKING: + + class _OnlyValueValidatorClsMethod(Protocol): + def __call__(self, cls: Any, value: Any, /) -> Any: ... + + class _V2ValidatorClsMethod(Protocol): + def __call__(self, cls: Any, value: Any, info: _core_schema.ValidationInfo, /) -> Any: ... + + class _OnlyValueWrapValidatorClsMethod(Protocol): + def __call__(self, cls: Any, value: Any, handler: _core_schema.ValidatorFunctionWrapHandler, /) -> Any: ... + + class _V2WrapValidatorClsMethod(Protocol): + def __call__( + self, + cls: Any, + value: Any, + handler: _core_schema.ValidatorFunctionWrapHandler, + info: _core_schema.ValidationInfo, + /, + ) -> Any: ... + + _V2Validator = Union[ + _V2ValidatorClsMethod, + _core_schema.WithInfoValidatorFunction, + _OnlyValueValidatorClsMethod, + _core_schema.NoInfoValidatorFunction, + ] + + _V2WrapValidator = Union[ + _V2WrapValidatorClsMethod, + _core_schema.WithInfoWrapValidatorFunction, + _OnlyValueWrapValidatorClsMethod, + _core_schema.NoInfoWrapValidatorFunction, + ] + + _PartialClsOrStaticMethod: TypeAlias = Union[classmethod[Any, Any, Any], staticmethod[Any, Any], partialmethod[Any]] + + _V2BeforeAfterOrPlainValidatorType = TypeVar( + '_V2BeforeAfterOrPlainValidatorType', + bound=Union[_V2Validator, _PartialClsOrStaticMethod], + ) + _V2WrapValidatorType = TypeVar('_V2WrapValidatorType', bound=Union[_V2WrapValidator, _PartialClsOrStaticMethod]) + +FieldValidatorModes: TypeAlias = Literal['before', 'after', 'wrap', 'plain'] + + +@overload +def field_validator( + field: str, + /, + *fields: str, + mode: Literal['wrap'], + check_fields: bool | None = ..., + json_schema_input_type: Any = ..., +) -> Callable[[_V2WrapValidatorType], _V2WrapValidatorType]: ... + + +@overload +def field_validator( + field: str, + /, + *fields: str, + mode: Literal['before', 'plain'], + check_fields: bool | None = ..., + json_schema_input_type: Any = ..., +) -> Callable[[_V2BeforeAfterOrPlainValidatorType], _V2BeforeAfterOrPlainValidatorType]: ... + + +@overload +def field_validator( + field: str, + /, + *fields: str, + mode: Literal['after'] = ..., + check_fields: bool | None = ..., +) -> Callable[[_V2BeforeAfterOrPlainValidatorType], _V2BeforeAfterOrPlainValidatorType]: ... + + +def field_validator( + field: str, + /, + *fields: str, + mode: FieldValidatorModes = 'after', + check_fields: bool | None = None, + json_schema_input_type: Any = PydanticUndefined, +) -> Callable[[Any], Any]: + """!!! abstract "Usage Documentation" + [field validators](../concepts/validators.md#field-validators) + + Decorate methods on the class indicating that they should be used to validate fields. + + Example usage: + ```python + from typing import Any + + from pydantic import ( + BaseModel, + ValidationError, + field_validator, + ) + + class Model(BaseModel): + a: str + + @field_validator('a') + @classmethod + def ensure_foobar(cls, v: Any): + if 'foobar' not in v: + raise ValueError('"foobar" not found in a') + return v + + print(repr(Model(a='this is foobar good'))) + #> Model(a='this is foobar good') + + try: + Model(a='snap') + except ValidationError as exc_info: + print(exc_info) + ''' + 1 validation error for Model + a + Value error, "foobar" not found in a [type=value_error, input_value='snap', input_type=str] + ''' + ``` + + For more in depth examples, see [Field Validators](../concepts/validators.md#field-validators). + + Args: + field: The first field the `field_validator` should be called on; this is separate + from `fields` to ensure an error is raised if you don't pass at least one. + *fields: Additional field(s) the `field_validator` should be called on. + mode: Specifies whether to validate the fields before or after validation. + check_fields: Whether to check that the fields actually exist on the model. + json_schema_input_type: The input type of the function. This is only used to generate + the appropriate JSON Schema (in validation mode) and can only specified + when `mode` is either `'before'`, `'plain'` or `'wrap'`. + + Returns: + A decorator that can be used to decorate a function to be used as a field_validator. + + Raises: + PydanticUserError: + - If `@field_validator` is used bare (with no fields). + - If the args passed to `@field_validator` as fields are not strings. + - If `@field_validator` applied to instance methods. + """ + if isinstance(field, FunctionType): + raise PydanticUserError( + '`@field_validator` should be used with fields and keyword arguments, not bare. ' + "E.g. usage should be `@validator('', ...)`", + code='validator-no-fields', + ) + + if mode not in ('before', 'plain', 'wrap') and json_schema_input_type is not PydanticUndefined: + raise PydanticUserError( + f"`json_schema_input_type` can't be used when mode is set to {mode!r}", + code='validator-input-type', + ) + + if json_schema_input_type is PydanticUndefined and mode == 'plain': + json_schema_input_type = Any + + fields = field, *fields + if not all(isinstance(field, str) for field in fields): + raise PydanticUserError( + '`@field_validator` fields should be passed as separate string args. ' + "E.g. usage should be `@validator('', '', ...)`", + code='validator-invalid-fields', + ) + + def dec( + f: Callable[..., Any] | staticmethod[Any, Any] | classmethod[Any, Any, Any], + ) -> _decorators.PydanticDescriptorProxy[Any]: + if _decorators.is_instance_method_from_sig(f): + raise PydanticUserError( + '`@field_validator` cannot be applied to instance methods', code='validator-instance-method' + ) + + # auto apply the @classmethod decorator + f = _decorators.ensure_classmethod_based_on_signature(f) + + dec_info = _decorators.FieldValidatorDecoratorInfo( + fields=fields, mode=mode, check_fields=check_fields, json_schema_input_type=json_schema_input_type + ) + return _decorators.PydanticDescriptorProxy(f, dec_info) + + return dec + + +_ModelType = TypeVar('_ModelType') +_ModelTypeCo = TypeVar('_ModelTypeCo', covariant=True) + + +class ModelWrapValidatorHandler(_core_schema.ValidatorFunctionWrapHandler, Protocol[_ModelTypeCo]): + """`@model_validator` decorated function handler argument type. This is used when `mode='wrap'`.""" + + def __call__( # noqa: D102 + self, + value: Any, + outer_location: str | int | None = None, + /, + ) -> _ModelTypeCo: # pragma: no cover + ... + + +class ModelWrapValidatorWithoutInfo(Protocol[_ModelType]): + """A `@model_validator` decorated function signature. + This is used when `mode='wrap'` and the function does not have info argument. + """ + + def __call__( # noqa: D102 + self, + cls: type[_ModelType], + # this can be a dict, a model instance + # or anything else that gets passed to validate_python + # thus validators _must_ handle all cases + value: Any, + handler: ModelWrapValidatorHandler[_ModelType], + /, + ) -> _ModelType: ... + + +class ModelWrapValidator(Protocol[_ModelType]): + """A `@model_validator` decorated function signature. This is used when `mode='wrap'`.""" + + def __call__( # noqa: D102 + self, + cls: type[_ModelType], + # this can be a dict, a model instance + # or anything else that gets passed to validate_python + # thus validators _must_ handle all cases + value: Any, + handler: ModelWrapValidatorHandler[_ModelType], + info: _core_schema.ValidationInfo, + /, + ) -> _ModelType: ... + + +class FreeModelBeforeValidatorWithoutInfo(Protocol): + """A `@model_validator` decorated function signature. + This is used when `mode='before'` and the function does not have info argument. + """ + + def __call__( # noqa: D102 + self, + # this can be a dict, a model instance + # or anything else that gets passed to validate_python + # thus validators _must_ handle all cases + value: Any, + /, + ) -> Any: ... + + +class ModelBeforeValidatorWithoutInfo(Protocol): + """A `@model_validator` decorated function signature. + This is used when `mode='before'` and the function does not have info argument. + """ + + def __call__( # noqa: D102 + self, + cls: Any, + # this can be a dict, a model instance + # or anything else that gets passed to validate_python + # thus validators _must_ handle all cases + value: Any, + /, + ) -> Any: ... + + +class FreeModelBeforeValidator(Protocol): + """A `@model_validator` decorated function signature. This is used when `mode='before'`.""" + + def __call__( # noqa: D102 + self, + # this can be a dict, a model instance + # or anything else that gets passed to validate_python + # thus validators _must_ handle all cases + value: Any, + info: _core_schema.ValidationInfo, + /, + ) -> Any: ... + + +class ModelBeforeValidator(Protocol): + """A `@model_validator` decorated function signature. This is used when `mode='before'`.""" + + def __call__( # noqa: D102 + self, + cls: Any, + # this can be a dict, a model instance + # or anything else that gets passed to validate_python + # thus validators _must_ handle all cases + value: Any, + info: _core_schema.ValidationInfo, + /, + ) -> Any: ... + + +ModelAfterValidatorWithoutInfo = Callable[[_ModelType], _ModelType] +"""A `@model_validator` decorated function signature. This is used when `mode='after'` and the function does not +have info argument. +""" + +ModelAfterValidator = Callable[[_ModelType, _core_schema.ValidationInfo], _ModelType] +"""A `@model_validator` decorated function signature. This is used when `mode='after'`.""" + +_AnyModelWrapValidator = Union[ModelWrapValidator[_ModelType], ModelWrapValidatorWithoutInfo[_ModelType]] +_AnyModelBeforeValidator = Union[ + FreeModelBeforeValidator, ModelBeforeValidator, FreeModelBeforeValidatorWithoutInfo, ModelBeforeValidatorWithoutInfo +] +_AnyModelAfterValidator = Union[ModelAfterValidator[_ModelType], ModelAfterValidatorWithoutInfo[_ModelType]] + + +@overload +def model_validator( + *, + mode: Literal['wrap'], +) -> Callable[ + [_AnyModelWrapValidator[_ModelType]], _decorators.PydanticDescriptorProxy[_decorators.ModelValidatorDecoratorInfo] +]: ... + + +@overload +def model_validator( + *, + mode: Literal['before'], +) -> Callable[ + [_AnyModelBeforeValidator], _decorators.PydanticDescriptorProxy[_decorators.ModelValidatorDecoratorInfo] +]: ... + + +@overload +def model_validator( + *, + mode: Literal['after'], +) -> Callable[ + [_AnyModelAfterValidator[_ModelType]], _decorators.PydanticDescriptorProxy[_decorators.ModelValidatorDecoratorInfo] +]: ... + + +def model_validator( + *, + mode: Literal['wrap', 'before', 'after'], +) -> Any: + """!!! abstract "Usage Documentation" + [Model Validators](../concepts/validators.md#model-validators) + + Decorate model methods for validation purposes. + + Example usage: + ```python + from typing_extensions import Self + + from pydantic import BaseModel, ValidationError, model_validator + + class Square(BaseModel): + width: float + height: float + + @model_validator(mode='after') + def verify_square(self) -> Self: + if self.width != self.height: + raise ValueError('width and height do not match') + return self + + s = Square(width=1, height=1) + print(repr(s)) + #> Square(width=1.0, height=1.0) + + try: + Square(width=1, height=2) + except ValidationError as e: + print(e) + ''' + 1 validation error for Square + Value error, width and height do not match [type=value_error, input_value={'width': 1, 'height': 2}, input_type=dict] + ''' + ``` + + For more in depth examples, see [Model Validators](../concepts/validators.md#model-validators). + + Args: + mode: A required string literal that specifies the validation mode. + It can be one of the following: 'wrap', 'before', or 'after'. + + Returns: + A decorator that can be used to decorate a function to be used as a model validator. + """ + + def dec(f: Any) -> _decorators.PydanticDescriptorProxy[Any]: + # auto apply the @classmethod decorator + f = _decorators.ensure_classmethod_based_on_signature(f) + dec_info = _decorators.ModelValidatorDecoratorInfo(mode=mode) + return _decorators.PydanticDescriptorProxy(f, dec_info) + + return dec + + +AnyType = TypeVar('AnyType') + + +if TYPE_CHECKING: + # If we add configurable attributes to IsInstance, we'd probably need to stop hiding it from type checkers like this + InstanceOf = Annotated[AnyType, ...] # `IsInstance[Sequence]` will be recognized by type checkers as `Sequence` + +else: + + @dataclasses.dataclass(**_internal_dataclass.slots_true) + class InstanceOf: + '''Generic type for annotating a type that is an instance of a given class. + + Example: + ```python + from pydantic import BaseModel, InstanceOf + + class Foo: + ... + + class Bar(BaseModel): + foo: InstanceOf[Foo] + + Bar(foo=Foo()) + try: + Bar(foo=42) + except ValidationError as e: + print(e) + """ + [ + │ { + │ │ 'type': 'is_instance_of', + │ │ 'loc': ('foo',), + │ │ 'msg': 'Input should be an instance of Foo', + │ │ 'input': 42, + │ │ 'ctx': {'class': 'Foo'}, + │ │ 'url': 'https://errors.pydantic.dev/0.38.0/v/is_instance_of' + │ } + ] + """ + ``` + ''' + + @classmethod + def __class_getitem__(cls, item: AnyType) -> AnyType: + return Annotated[item, cls()] + + @classmethod + def __get_pydantic_core_schema__(cls, source: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + from pydantic import PydanticSchemaGenerationError + + # use the generic _origin_ as the second argument to isinstance when appropriate + instance_of_schema = core_schema.is_instance_schema(_generics.get_origin(source) or source) + + try: + # Try to generate the "standard" schema, which will be used when loading from JSON + original_schema = handler(source) + except PydanticSchemaGenerationError: + # If that fails, just produce a schema that can validate from python + return instance_of_schema + else: + # Use the "original" approach to serialization + instance_of_schema['serialization'] = core_schema.wrap_serializer_function_ser_schema( + function=lambda v, h: h(v), schema=original_schema + ) + return core_schema.json_or_python_schema(python_schema=instance_of_schema, json_schema=original_schema) + + __hash__ = object.__hash__ + + +if TYPE_CHECKING: + SkipValidation = Annotated[AnyType, ...] # SkipValidation[list[str]] will be treated by type checkers as list[str] +else: + + @dataclasses.dataclass(**_internal_dataclass.slots_true) + class SkipValidation: + """If this is applied as an annotation (e.g., via `x: Annotated[int, SkipValidation]`), validation will be + skipped. You can also use `SkipValidation[int]` as a shorthand for `Annotated[int, SkipValidation]`. + + This can be useful if you want to use a type annotation for documentation/IDE/type-checking purposes, + and know that it is safe to skip validation for one or more of the fields. + + Because this converts the validation schema to `any_schema`, subsequent annotation-applied transformations + may not have the expected effects. Therefore, when used, this annotation should generally be the final + annotation applied to a type. + """ + + def __class_getitem__(cls, item: Any) -> Any: + return Annotated[item, SkipValidation()] + + @classmethod + def __get_pydantic_core_schema__(cls, source: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + original_schema = handler(source) + metadata = {'pydantic_js_annotation_functions': [lambda _c, h: h(original_schema)]} + return core_schema.any_schema( + metadata=metadata, + serialization=core_schema.wrap_serializer_function_ser_schema( + function=lambda v, h: h(v), schema=original_schema + ), + ) + + __hash__ = object.__hash__ diff --git a/lib/python3.12/site-packages/pydantic/generics.py b/lib/python3.12/site-packages/pydantic/generics.py new file mode 100644 index 0000000000000000000000000000000000000000..3f1070d08f3ac5bd554794401734eb349373ab8e --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/generics.py @@ -0,0 +1,5 @@ +"""The `generics` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/json.py b/lib/python3.12/site-packages/pydantic/json.py new file mode 100644 index 0000000000000000000000000000000000000000..bcaff9f57958b8c0938255fd8f937293c5850cbd --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/json.py @@ -0,0 +1,5 @@ +"""The `json` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/json_schema.py b/lib/python3.12/site-packages/pydantic/json_schema.py new file mode 100644 index 0000000000000000000000000000000000000000..be9595c4ffc0112f2a50f264e104609ac56d5816 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/json_schema.py @@ -0,0 +1,2695 @@ +"""!!! abstract "Usage Documentation" + [JSON Schema](../concepts/json_schema.md) + +The `json_schema` module contains classes and functions to allow the way [JSON Schema](https://json-schema.org/) +is generated to be customized. + +In general you shouldn't need to use this module directly; instead, you can use +[`BaseModel.model_json_schema`][pydantic.BaseModel.model_json_schema] and +[`TypeAdapter.json_schema`][pydantic.TypeAdapter.json_schema]. +""" + +from __future__ import annotations as _annotations + +import dataclasses +import inspect +import math +import os +import re +import warnings +from collections import Counter, defaultdict +from collections.abc import Hashable, Iterable, Sequence +from copy import deepcopy +from enum import Enum +from re import Pattern +from typing import ( + TYPE_CHECKING, + Annotated, + Any, + Callable, + Literal, + NewType, + TypeVar, + Union, + cast, + overload, +) + +import pydantic_core +from pydantic_core import CoreSchema, PydanticOmit, core_schema, to_jsonable_python +from pydantic_core.core_schema import ComputedField +from typing_extensions import TypeAlias, assert_never, deprecated, final +from typing_inspection.introspection import get_literal_values + +from pydantic.warnings import PydanticDeprecatedSince26, PydanticDeprecatedSince29 + +from ._internal import ( + _config, + _core_metadata, + _core_utils, + _decorators, + _internal_dataclass, + _mock_val_ser, + _schema_generation_shared, +) +from .annotated_handlers import GetJsonSchemaHandler +from .config import JsonDict, JsonValue +from .errors import PydanticInvalidForJsonSchema, PydanticSchemaGenerationError, PydanticUserError + +if TYPE_CHECKING: + from . import ConfigDict + from ._internal._core_utils import CoreSchemaField, CoreSchemaOrField + from ._internal._dataclasses import PydanticDataclass + from ._internal._schema_generation_shared import GetJsonSchemaFunction + from .main import BaseModel + + +CoreSchemaOrFieldType = Literal[core_schema.CoreSchemaType, core_schema.CoreSchemaFieldType] +""" +A type alias for defined schema types that represents a union of +`core_schema.CoreSchemaType` and +`core_schema.CoreSchemaFieldType`. +""" + +JsonSchemaValue = dict[str, Any] +""" +A type alias for a JSON schema value. This is a dictionary of string keys to arbitrary JSON values. +""" + +JsonSchemaMode = Literal['validation', 'serialization'] +""" +A type alias that represents the mode of a JSON schema; either 'validation' or 'serialization'. + +For some types, the inputs to validation differ from the outputs of serialization. For example, +computed fields will only be present when serializing, and should not be provided when +validating. This flag provides a way to indicate whether you want the JSON schema required +for validation inputs, or that will be matched by serialization outputs. +""" + +_MODE_TITLE_MAPPING: dict[JsonSchemaMode, str] = {'validation': 'Input', 'serialization': 'Output'} + + +JsonSchemaWarningKind = Literal['skipped-choice', 'non-serializable-default', 'skipped-discriminator'] +""" +A type alias representing the kinds of warnings that can be emitted during JSON schema generation. + +See [`GenerateJsonSchema.render_warning_message`][pydantic.json_schema.GenerateJsonSchema.render_warning_message] +for more details. +""" + + +class PydanticJsonSchemaWarning(UserWarning): + """This class is used to emit warnings produced during JSON schema generation. + See the [`GenerateJsonSchema.emit_warning`][pydantic.json_schema.GenerateJsonSchema.emit_warning] and + [`GenerateJsonSchema.render_warning_message`][pydantic.json_schema.GenerateJsonSchema.render_warning_message] + methods for more details; these can be overridden to control warning behavior. + """ + + +NoDefault = object() +"""A sentinel value used to indicate that no default value should be used when generating a JSON Schema +for a core schema with a default value. +""" + + +# ##### JSON Schema Generation ##### +DEFAULT_REF_TEMPLATE = '#/$defs/{model}' +"""The default format string used to generate reference names.""" + +# There are three types of references relevant to building JSON schemas: +# 1. core_schema "ref" values; these are not exposed as part of the JSON schema +# * these might look like the fully qualified path of a model, its id, or something similar +CoreRef = NewType('CoreRef', str) +# 2. keys of the "definitions" object that will eventually go into the JSON schema +# * by default, these look like "MyModel", though may change in the presence of collisions +# * eventually, we may want to make it easier to modify the way these names are generated +DefsRef = NewType('DefsRef', str) +# 3. the values corresponding to the "$ref" key in the schema +# * By default, these look like "#/$defs/MyModel", as in {"$ref": "#/$defs/MyModel"} +JsonRef = NewType('JsonRef', str) + +CoreModeRef = tuple[CoreRef, JsonSchemaMode] +JsonSchemaKeyT = TypeVar('JsonSchemaKeyT', bound=Hashable) + + +@dataclasses.dataclass(**_internal_dataclass.slots_true) +class _DefinitionsRemapping: + defs_remapping: dict[DefsRef, DefsRef] + json_remapping: dict[JsonRef, JsonRef] + + @staticmethod + def from_prioritized_choices( + prioritized_choices: dict[DefsRef, list[DefsRef]], + defs_to_json: dict[DefsRef, JsonRef], + definitions: dict[DefsRef, JsonSchemaValue], + ) -> _DefinitionsRemapping: + """ + This function should produce a remapping that replaces complex DefsRef with the simpler ones from the + prioritized_choices such that applying the name remapping would result in an equivalent JSON schema. + """ + # We need to iteratively simplify the definitions until we reach a fixed point. + # The reason for this is that outer definitions may reference inner definitions that get simplified + # into an equivalent reference, and the outer definitions won't be equivalent until we've simplified + # the inner definitions. + copied_definitions = deepcopy(definitions) + definitions_schema = {'$defs': copied_definitions} + for _iter in range(100): # prevent an infinite loop in the case of a bug, 100 iterations should be enough + # For every possible remapped DefsRef, collect all schemas that that DefsRef might be used for: + schemas_for_alternatives: dict[DefsRef, list[JsonSchemaValue]] = defaultdict(list) + for defs_ref in copied_definitions: + alternatives = prioritized_choices[defs_ref] + for alternative in alternatives: + schemas_for_alternatives[alternative].append(copied_definitions[defs_ref]) + + # Deduplicate the schemas for each alternative; the idea is that we only want to remap to a new DefsRef + # if it introduces no ambiguity, i.e., there is only one distinct schema for that DefsRef. + for defs_ref in schemas_for_alternatives: + schemas_for_alternatives[defs_ref] = _deduplicate_schemas(schemas_for_alternatives[defs_ref]) + + # Build the remapping + defs_remapping: dict[DefsRef, DefsRef] = {} + json_remapping: dict[JsonRef, JsonRef] = {} + for original_defs_ref in definitions: + alternatives = prioritized_choices[original_defs_ref] + # Pick the first alternative that has only one schema, since that means there is no collision + remapped_defs_ref = next(x for x in alternatives if len(schemas_for_alternatives[x]) == 1) + defs_remapping[original_defs_ref] = remapped_defs_ref + json_remapping[defs_to_json[original_defs_ref]] = defs_to_json[remapped_defs_ref] + remapping = _DefinitionsRemapping(defs_remapping, json_remapping) + new_definitions_schema = remapping.remap_json_schema({'$defs': copied_definitions}) + if definitions_schema == new_definitions_schema: + # We've reached the fixed point + return remapping + definitions_schema = new_definitions_schema + + raise PydanticInvalidForJsonSchema('Failed to simplify the JSON schema definitions') + + def remap_defs_ref(self, ref: DefsRef) -> DefsRef: + return self.defs_remapping.get(ref, ref) + + def remap_json_ref(self, ref: JsonRef) -> JsonRef: + return self.json_remapping.get(ref, ref) + + def remap_json_schema(self, schema: Any) -> Any: + """ + Recursively update the JSON schema replacing all $refs + """ + if isinstance(schema, str): + # Note: this may not really be a JsonRef; we rely on having no collisions between JsonRefs and other strings + return self.remap_json_ref(JsonRef(schema)) + elif isinstance(schema, list): + return [self.remap_json_schema(item) for item in schema] + elif isinstance(schema, dict): + for key, value in schema.items(): + if key == '$ref' and isinstance(value, str): + schema['$ref'] = self.remap_json_ref(JsonRef(value)) + elif key == '$defs': + schema['$defs'] = { + self.remap_defs_ref(DefsRef(key)): self.remap_json_schema(value) + for key, value in schema['$defs'].items() + } + else: + schema[key] = self.remap_json_schema(value) + return schema + + +class GenerateJsonSchema: + """!!! abstract "Usage Documentation" + [Customizing the JSON Schema Generation Process](../concepts/json_schema.md#customizing-the-json-schema-generation-process) + + A class for generating JSON schemas. + + This class generates JSON schemas based on configured parameters. The default schema dialect + is [https://json-schema.org/draft/2020-12/schema](https://json-schema.org/draft/2020-12/schema). + The class uses `by_alias` to configure how fields with + multiple names are handled and `ref_template` to format reference names. + + Attributes: + schema_dialect: The JSON schema dialect used to generate the schema. See + [Declaring a Dialect](https://json-schema.org/understanding-json-schema/reference/schema.html#id4) + in the JSON Schema documentation for more information about dialects. + ignored_warning_kinds: Warnings to ignore when generating the schema. `self.render_warning_message` will + do nothing if its argument `kind` is in `ignored_warning_kinds`; + this value can be modified on subclasses to easily control which warnings are emitted. + by_alias: Whether to use field aliases when generating the schema. + ref_template: The format string used when generating reference names. + core_to_json_refs: A mapping of core refs to JSON refs. + core_to_defs_refs: A mapping of core refs to definition refs. + defs_to_core_refs: A mapping of definition refs to core refs. + json_to_defs_refs: A mapping of JSON refs to definition refs. + definitions: Definitions in the schema. + + Args: + by_alias: Whether to use field aliases in the generated schemas. + ref_template: The format string to use when generating reference names. + + Raises: + JsonSchemaError: If the instance of the class is inadvertently reused after generating a schema. + """ + + schema_dialect = 'https://json-schema.org/draft/2020-12/schema' + + # `self.render_warning_message` will do nothing if its argument `kind` is in `ignored_warning_kinds`; + # this value can be modified on subclasses to easily control which warnings are emitted + ignored_warning_kinds: set[JsonSchemaWarningKind] = {'skipped-choice'} + + def __init__(self, by_alias: bool = True, ref_template: str = DEFAULT_REF_TEMPLATE): + self.by_alias = by_alias + self.ref_template = ref_template + + self.core_to_json_refs: dict[CoreModeRef, JsonRef] = {} + self.core_to_defs_refs: dict[CoreModeRef, DefsRef] = {} + self.defs_to_core_refs: dict[DefsRef, CoreModeRef] = {} + self.json_to_defs_refs: dict[JsonRef, DefsRef] = {} + + self.definitions: dict[DefsRef, JsonSchemaValue] = {} + self._config_wrapper_stack = _config.ConfigWrapperStack(_config.ConfigWrapper({})) + + self._mode: JsonSchemaMode = 'validation' + + # The following includes a mapping of a fully-unique defs ref choice to a list of preferred + # alternatives, which are generally simpler, such as only including the class name. + # At the end of schema generation, we use these to produce a JSON schema with more human-readable + # definitions, which would also work better in a generated OpenAPI client, etc. + self._prioritized_defsref_choices: dict[DefsRef, list[DefsRef]] = {} + self._collision_counter: dict[str, int] = defaultdict(int) + self._collision_index: dict[str, int] = {} + + self._schema_type_to_method = self.build_schema_type_to_method() + + # When we encounter definitions we need to try to build them immediately + # so that they are available schemas that reference them + # But it's possible that CoreSchema was never going to be used + # (e.g. because the CoreSchema that references short circuits is JSON schema generation without needing + # the reference) so instead of failing altogether if we can't build a definition we + # store the error raised and re-throw it if we end up needing that def + self._core_defs_invalid_for_json_schema: dict[DefsRef, PydanticInvalidForJsonSchema] = {} + + # This changes to True after generating a schema, to prevent issues caused by accidental reuse + # of a single instance of a schema generator + self._used = False + + @property + def _config(self) -> _config.ConfigWrapper: + return self._config_wrapper_stack.tail + + @property + def mode(self) -> JsonSchemaMode: + if self._config.json_schema_mode_override is not None: + return self._config.json_schema_mode_override + else: + return self._mode + + def build_schema_type_to_method( + self, + ) -> dict[CoreSchemaOrFieldType, Callable[[CoreSchemaOrField], JsonSchemaValue]]: + """Builds a dictionary mapping fields to methods for generating JSON schemas. + + Returns: + A dictionary containing the mapping of `CoreSchemaOrFieldType` to a handler method. + + Raises: + TypeError: If no method has been defined for generating a JSON schema for a given pydantic core schema type. + """ + mapping: dict[CoreSchemaOrFieldType, Callable[[CoreSchemaOrField], JsonSchemaValue]] = {} + core_schema_types: list[CoreSchemaOrFieldType] = list(get_literal_values(CoreSchemaOrFieldType)) + for key in core_schema_types: + method_name = f'{key.replace("-", "_")}_schema' + try: + mapping[key] = getattr(self, method_name) + except AttributeError as e: # pragma: no cover + if os.getenv('PYDANTIC_PRIVATE_ALLOW_UNHANDLED_SCHEMA_TYPES'): + continue + raise TypeError( + f'No method for generating JsonSchema for core_schema.type={key!r} ' + f'(expected: {type(self).__name__}.{method_name})' + ) from e + return mapping + + def generate_definitions( + self, inputs: Sequence[tuple[JsonSchemaKeyT, JsonSchemaMode, core_schema.CoreSchema]] + ) -> tuple[dict[tuple[JsonSchemaKeyT, JsonSchemaMode], JsonSchemaValue], dict[DefsRef, JsonSchemaValue]]: + """Generates JSON schema definitions from a list of core schemas, pairing the generated definitions with a + mapping that links the input keys to the definition references. + + Args: + inputs: A sequence of tuples, where: + + - The first element is a JSON schema key type. + - The second element is the JSON mode: either 'validation' or 'serialization'. + - The third element is a core schema. + + Returns: + A tuple where: + + - The first element is a dictionary whose keys are tuples of JSON schema key type and JSON mode, and + whose values are the JSON schema corresponding to that pair of inputs. (These schemas may have + JsonRef references to definitions that are defined in the second returned element.) + - The second element is a dictionary whose keys are definition references for the JSON schemas + from the first returned element, and whose values are the actual JSON schema definitions. + + Raises: + PydanticUserError: Raised if the JSON schema generator has already been used to generate a JSON schema. + """ + if self._used: + raise PydanticUserError( + 'This JSON schema generator has already been used to generate a JSON schema. ' + f'You must create a new instance of {type(self).__name__} to generate a new JSON schema.', + code='json-schema-already-used', + ) + + for _, mode, schema in inputs: + self._mode = mode + self.generate_inner(schema) + + definitions_remapping = self._build_definitions_remapping() + + json_schemas_map: dict[tuple[JsonSchemaKeyT, JsonSchemaMode], DefsRef] = {} + for key, mode, schema in inputs: + self._mode = mode + json_schema = self.generate_inner(schema) + json_schemas_map[(key, mode)] = definitions_remapping.remap_json_schema(json_schema) + + json_schema = {'$defs': self.definitions} + json_schema = definitions_remapping.remap_json_schema(json_schema) + self._used = True + return json_schemas_map, self.sort(json_schema['$defs']) # type: ignore + + def generate(self, schema: CoreSchema, mode: JsonSchemaMode = 'validation') -> JsonSchemaValue: + """Generates a JSON schema for a specified schema in a specified mode. + + Args: + schema: A Pydantic model. + mode: The mode in which to generate the schema. Defaults to 'validation'. + + Returns: + A JSON schema representing the specified schema. + + Raises: + PydanticUserError: If the JSON schema generator has already been used to generate a JSON schema. + """ + self._mode = mode + if self._used: + raise PydanticUserError( + 'This JSON schema generator has already been used to generate a JSON schema. ' + f'You must create a new instance of {type(self).__name__} to generate a new JSON schema.', + code='json-schema-already-used', + ) + + json_schema: JsonSchemaValue = self.generate_inner(schema) + json_ref_counts = self.get_json_ref_counts(json_schema) + + ref = cast(JsonRef, json_schema.get('$ref')) + while ref is not None: # may need to unpack multiple levels + ref_json_schema = self.get_schema_from_definitions(ref) + if json_ref_counts[ref] == 1 and ref_json_schema is not None and len(json_schema) == 1: + # "Unpack" the ref since this is the only reference and there are no sibling keys + json_schema = ref_json_schema.copy() # copy to prevent recursive dict reference + json_ref_counts[ref] -= 1 + ref = cast(JsonRef, json_schema.get('$ref')) + ref = None + + self._garbage_collect_definitions(json_schema) + definitions_remapping = self._build_definitions_remapping() + + if self.definitions: + json_schema['$defs'] = self.definitions + + json_schema = definitions_remapping.remap_json_schema(json_schema) + + # For now, we will not set the $schema key. However, if desired, this can be easily added by overriding + # this method and adding the following line after a call to super().generate(schema): + # json_schema['$schema'] = self.schema_dialect + + self._used = True + return self.sort(json_schema) + + def generate_inner(self, schema: CoreSchemaOrField) -> JsonSchemaValue: # noqa: C901 + """Generates a JSON schema for a given core schema. + + Args: + schema: The given core schema. + + Returns: + The generated JSON schema. + + TODO: the nested function definitions here seem like bad practice, I'd like to unpack these + in a future PR. It'd be great if we could shorten the call stack a bit for JSON schema generation, + and I think there's potential for that here. + """ + # If a schema with the same CoreRef has been handled, just return a reference to it + # Note that this assumes that it will _never_ be the case that the same CoreRef is used + # on types that should have different JSON schemas + if 'ref' in schema: + core_ref = CoreRef(schema['ref']) # type: ignore[typeddict-item] + core_mode_ref = (core_ref, self.mode) + if core_mode_ref in self.core_to_defs_refs and self.core_to_defs_refs[core_mode_ref] in self.definitions: + return {'$ref': self.core_to_json_refs[core_mode_ref]} + + def populate_defs(core_schema: CoreSchema, json_schema: JsonSchemaValue) -> JsonSchemaValue: + if 'ref' in core_schema: + core_ref = CoreRef(core_schema['ref']) # type: ignore[typeddict-item] + defs_ref, ref_json_schema = self.get_cache_defs_ref_schema(core_ref) + json_ref = JsonRef(ref_json_schema['$ref']) + # Replace the schema if it's not a reference to itself + # What we want to avoid is having the def be just a ref to itself + # which is what would happen if we blindly assigned any + if json_schema.get('$ref', None) != json_ref: + self.definitions[defs_ref] = json_schema + self._core_defs_invalid_for_json_schema.pop(defs_ref, None) + json_schema = ref_json_schema + return json_schema + + def handler_func(schema_or_field: CoreSchemaOrField) -> JsonSchemaValue: + """Generate a JSON schema based on the input schema. + + Args: + schema_or_field: The core schema to generate a JSON schema from. + + Returns: + The generated JSON schema. + + Raises: + TypeError: If an unexpected schema type is encountered. + """ + # Generate the core-schema-type-specific bits of the schema generation: + json_schema: JsonSchemaValue | None = None + if self.mode == 'serialization' and 'serialization' in schema_or_field: + # In this case, we skip the JSON Schema generation of the schema + # and use the `'serialization'` schema instead (canonical example: + # `Annotated[int, PlainSerializer(str)]`). + ser_schema = schema_or_field['serialization'] # type: ignore + json_schema = self.ser_schema(ser_schema) + + # It might be that the 'serialization'` is skipped depending on `when_used`. + # This is only relevant for `nullable` schemas though, so we special case here. + if ( + json_schema is not None + and ser_schema.get('when_used') in ('unless-none', 'json-unless-none') + and schema_or_field['type'] == 'nullable' + ): + json_schema = self.get_flattened_anyof([{'type': 'null'}, json_schema]) + if json_schema is None: + if _core_utils.is_core_schema(schema_or_field) or _core_utils.is_core_schema_field(schema_or_field): + generate_for_schema_type = self._schema_type_to_method[schema_or_field['type']] + json_schema = generate_for_schema_type(schema_or_field) + else: + raise TypeError(f'Unexpected schema type: schema={schema_or_field}') + + return json_schema + + current_handler = _schema_generation_shared.GenerateJsonSchemaHandler(self, handler_func) + + metadata = cast(_core_metadata.CoreMetadata, schema.get('metadata', {})) + + # TODO: I dislike that we have to wrap these basic dict updates in callables, is there any way around this? + + if js_updates := metadata.get('pydantic_js_updates'): + + def js_updates_handler_func( + schema_or_field: CoreSchemaOrField, + current_handler: GetJsonSchemaHandler = current_handler, + ) -> JsonSchemaValue: + json_schema = {**current_handler(schema_or_field), **js_updates} + return json_schema + + current_handler = _schema_generation_shared.GenerateJsonSchemaHandler(self, js_updates_handler_func) + + if js_extra := metadata.get('pydantic_js_extra'): + + def js_extra_handler_func( + schema_or_field: CoreSchemaOrField, + current_handler: GetJsonSchemaHandler = current_handler, + ) -> JsonSchemaValue: + json_schema = current_handler(schema_or_field) + if isinstance(js_extra, dict): + json_schema.update(to_jsonable_python(js_extra)) + elif callable(js_extra): + # similar to typing issue in _update_class_schema when we're working with callable js extra + js_extra(json_schema) # type: ignore + return json_schema + + current_handler = _schema_generation_shared.GenerateJsonSchemaHandler(self, js_extra_handler_func) + + for js_modify_function in metadata.get('pydantic_js_functions', ()): + + def new_handler_func( + schema_or_field: CoreSchemaOrField, + current_handler: GetJsonSchemaHandler = current_handler, + js_modify_function: GetJsonSchemaFunction = js_modify_function, + ) -> JsonSchemaValue: + json_schema = js_modify_function(schema_or_field, current_handler) + if _core_utils.is_core_schema(schema_or_field): + json_schema = populate_defs(schema_or_field, json_schema) + original_schema = current_handler.resolve_ref_schema(json_schema) + ref = json_schema.pop('$ref', None) + if ref and json_schema: + original_schema.update(json_schema) + return original_schema + + current_handler = _schema_generation_shared.GenerateJsonSchemaHandler(self, new_handler_func) + + for js_modify_function in metadata.get('pydantic_js_annotation_functions', ()): + + def new_handler_func( + schema_or_field: CoreSchemaOrField, + current_handler: GetJsonSchemaHandler = current_handler, + js_modify_function: GetJsonSchemaFunction = js_modify_function, + ) -> JsonSchemaValue: + return js_modify_function(schema_or_field, current_handler) + + current_handler = _schema_generation_shared.GenerateJsonSchemaHandler(self, new_handler_func) + + json_schema = current_handler(schema) + if _core_utils.is_core_schema(schema): + json_schema = populate_defs(schema, json_schema) + return json_schema + + def sort(self, value: JsonSchemaValue, parent_key: str | None = None) -> JsonSchemaValue: + """Override this method to customize the sorting of the JSON schema (e.g., don't sort at all, sort all keys unconditionally, etc.) + + By default, alphabetically sort the keys in the JSON schema, skipping the 'properties' and 'default' keys to preserve field definition order. + This sort is recursive, so it will sort all nested dictionaries as well. + """ + sorted_dict: dict[str, JsonSchemaValue] = {} + keys = value.keys() + if parent_key not in ('properties', 'default'): + keys = sorted(keys) + for key in keys: + sorted_dict[key] = self._sort_recursive(value[key], parent_key=key) + return sorted_dict + + def _sort_recursive(self, value: Any, parent_key: str | None = None) -> Any: + """Recursively sort a JSON schema value.""" + if isinstance(value, dict): + sorted_dict: dict[str, JsonSchemaValue] = {} + keys = value.keys() + if parent_key not in ('properties', 'default'): + keys = sorted(keys) + for key in keys: + sorted_dict[key] = self._sort_recursive(value[key], parent_key=key) + return sorted_dict + elif isinstance(value, list): + sorted_list: list[JsonSchemaValue] = [] + for item in value: + sorted_list.append(self._sort_recursive(item, parent_key)) + return sorted_list + else: + return value + + # ### Schema generation methods + + def invalid_schema(self, schema: core_schema.InvalidSchema) -> JsonSchemaValue: + """Placeholder - should never be called.""" + + raise RuntimeError('Cannot generate schema for invalid_schema. This is a bug! Please report it.') + + def any_schema(self, schema: core_schema.AnySchema) -> JsonSchemaValue: + """Generates a JSON schema that matches any value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return {} + + def none_schema(self, schema: core_schema.NoneSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches `None`. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return {'type': 'null'} + + def bool_schema(self, schema: core_schema.BoolSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a bool value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return {'type': 'boolean'} + + def int_schema(self, schema: core_schema.IntSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches an int value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema: dict[str, Any] = {'type': 'integer'} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.numeric) + json_schema = {k: v for k, v in json_schema.items() if v not in {math.inf, -math.inf}} + return json_schema + + def float_schema(self, schema: core_schema.FloatSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a float value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema: dict[str, Any] = {'type': 'number'} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.numeric) + json_schema = {k: v for k, v in json_schema.items() if v not in {math.inf, -math.inf}} + return json_schema + + def decimal_schema(self, schema: core_schema.DecimalSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a decimal value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema = self.str_schema(core_schema.str_schema()) + if self.mode == 'validation': + multiple_of = schema.get('multiple_of') + le = schema.get('le') + ge = schema.get('ge') + lt = schema.get('lt') + gt = schema.get('gt') + json_schema = { + 'anyOf': [ + self.float_schema( + core_schema.float_schema( + allow_inf_nan=schema.get('allow_inf_nan'), + multiple_of=None if multiple_of is None else float(multiple_of), + le=None if le is None else float(le), + ge=None if ge is None else float(ge), + lt=None if lt is None else float(lt), + gt=None if gt is None else float(gt), + ) + ), + json_schema, + ], + } + return json_schema + + def str_schema(self, schema: core_schema.StringSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a string value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema = {'type': 'string'} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.string) + if isinstance(json_schema.get('pattern'), Pattern): + # TODO: should we add regex flags to the pattern? + json_schema['pattern'] = json_schema.get('pattern').pattern # type: ignore + return json_schema + + def bytes_schema(self, schema: core_schema.BytesSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a bytes value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema = {'type': 'string', 'format': 'base64url' if self._config.ser_json_bytes == 'base64' else 'binary'} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.bytes) + return json_schema + + def date_schema(self, schema: core_schema.DateSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a date value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return {'type': 'string', 'format': 'date'} + + def time_schema(self, schema: core_schema.TimeSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a time value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return {'type': 'string', 'format': 'time'} + + def datetime_schema(self, schema: core_schema.DatetimeSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a datetime value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return {'type': 'string', 'format': 'date-time'} + + def timedelta_schema(self, schema: core_schema.TimedeltaSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a timedelta value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + if self._config.ser_json_timedelta == 'float': + return {'type': 'number'} + return {'type': 'string', 'format': 'duration'} + + def literal_schema(self, schema: core_schema.LiteralSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a literal value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + expected = [to_jsonable_python(v.value if isinstance(v, Enum) else v) for v in schema['expected']] + + result: dict[str, Any] = {} + if len(expected) == 1: + result['const'] = expected[0] + else: + result['enum'] = expected + + types = {type(e) for e in expected} + if types == {str}: + result['type'] = 'string' + elif types == {int}: + result['type'] = 'integer' + elif types == {float}: + result['type'] = 'number' + elif types == {bool}: + result['type'] = 'boolean' + elif types == {list}: + result['type'] = 'array' + elif types == {type(None)}: + result['type'] = 'null' + return result + + def enum_schema(self, schema: core_schema.EnumSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches an Enum value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + enum_type = schema['cls'] + description = None if not enum_type.__doc__ else inspect.cleandoc(enum_type.__doc__) + if ( + description == 'An enumeration.' + ): # This is the default value provided by enum.EnumMeta.__new__; don't use it + description = None + result: dict[str, Any] = {'title': enum_type.__name__, 'description': description} + result = {k: v for k, v in result.items() if v is not None} + + expected = [to_jsonable_python(v.value) for v in schema['members']] + + result['enum'] = expected + + types = {type(e) for e in expected} + if isinstance(enum_type, str) or types == {str}: + result['type'] = 'string' + elif isinstance(enum_type, int) or types == {int}: + result['type'] = 'integer' + elif isinstance(enum_type, float) or types == {float}: + result['type'] = 'number' + elif types == {bool}: + result['type'] = 'boolean' + elif types == {list}: + result['type'] = 'array' + + return result + + def is_instance_schema(self, schema: core_schema.IsInstanceSchema) -> JsonSchemaValue: + """Handles JSON schema generation for a core schema that checks if a value is an instance of a class. + + Unless overridden in a subclass, this raises an error. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.handle_invalid_for_json_schema(schema, f'core_schema.IsInstanceSchema ({schema["cls"]})') + + def is_subclass_schema(self, schema: core_schema.IsSubclassSchema) -> JsonSchemaValue: + """Handles JSON schema generation for a core schema that checks if a value is a subclass of a class. + + For backwards compatibility with v1, this does not raise an error, but can be overridden to change this. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + # Note: This is for compatibility with V1; you can override if you want different behavior. + return {} + + def callable_schema(self, schema: core_schema.CallableSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a callable value. + + Unless overridden in a subclass, this raises an error. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.handle_invalid_for_json_schema(schema, 'core_schema.CallableSchema') + + def list_schema(self, schema: core_schema.ListSchema) -> JsonSchemaValue: + """Returns a schema that matches a list schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + items_schema = {} if 'items_schema' not in schema else self.generate_inner(schema['items_schema']) + json_schema = {'type': 'array', 'items': items_schema} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.array) + return json_schema + + @deprecated('`tuple_positional_schema` is deprecated. Use `tuple_schema` instead.', category=None) + @final + def tuple_positional_schema(self, schema: core_schema.TupleSchema) -> JsonSchemaValue: + """Replaced by `tuple_schema`.""" + warnings.warn( + '`tuple_positional_schema` is deprecated. Use `tuple_schema` instead.', + PydanticDeprecatedSince26, + stacklevel=2, + ) + return self.tuple_schema(schema) + + @deprecated('`tuple_variable_schema` is deprecated. Use `tuple_schema` instead.', category=None) + @final + def tuple_variable_schema(self, schema: core_schema.TupleSchema) -> JsonSchemaValue: + """Replaced by `tuple_schema`.""" + warnings.warn( + '`tuple_variable_schema` is deprecated. Use `tuple_schema` instead.', + PydanticDeprecatedSince26, + stacklevel=2, + ) + return self.tuple_schema(schema) + + def tuple_schema(self, schema: core_schema.TupleSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a tuple schema e.g. `tuple[int, + str, bool]` or `tuple[int, ...]`. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema: JsonSchemaValue = {'type': 'array'} + if 'variadic_item_index' in schema: + variadic_item_index = schema['variadic_item_index'] + if variadic_item_index > 0: + json_schema['minItems'] = variadic_item_index + json_schema['prefixItems'] = [ + self.generate_inner(item) for item in schema['items_schema'][:variadic_item_index] + ] + if variadic_item_index + 1 == len(schema['items_schema']): + # if the variadic item is the last item, then represent it faithfully + json_schema['items'] = self.generate_inner(schema['items_schema'][variadic_item_index]) + else: + # otherwise, 'items' represents the schema for the variadic + # item plus the suffix, so just allow anything for simplicity + # for now + json_schema['items'] = True + else: + prefixItems = [self.generate_inner(item) for item in schema['items_schema']] + if prefixItems: + json_schema['prefixItems'] = prefixItems + json_schema['minItems'] = len(prefixItems) + json_schema['maxItems'] = len(prefixItems) + self.update_with_validations(json_schema, schema, self.ValidationsMapping.array) + return json_schema + + def set_schema(self, schema: core_schema.SetSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a set schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self._common_set_schema(schema) + + def frozenset_schema(self, schema: core_schema.FrozenSetSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a frozenset schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self._common_set_schema(schema) + + def _common_set_schema(self, schema: core_schema.SetSchema | core_schema.FrozenSetSchema) -> JsonSchemaValue: + items_schema = {} if 'items_schema' not in schema else self.generate_inner(schema['items_schema']) + json_schema = {'type': 'array', 'uniqueItems': True, 'items': items_schema} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.array) + return json_schema + + def generator_schema(self, schema: core_schema.GeneratorSchema) -> JsonSchemaValue: + """Returns a JSON schema that represents the provided GeneratorSchema. + + Args: + schema: The schema. + + Returns: + The generated JSON schema. + """ + items_schema = {} if 'items_schema' not in schema else self.generate_inner(schema['items_schema']) + json_schema = {'type': 'array', 'items': items_schema} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.array) + return json_schema + + def dict_schema(self, schema: core_schema.DictSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a dict schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema: JsonSchemaValue = {'type': 'object'} + + keys_schema = self.generate_inner(schema['keys_schema']).copy() if 'keys_schema' in schema else {} + if '$ref' not in keys_schema: + keys_pattern = keys_schema.pop('pattern', None) + # Don't give a title to patternProperties/propertyNames: + keys_schema.pop('title', None) + else: + # Here, we assume that if the keys schema is a definition reference, + # it can't be a simple string core schema (and thus no pattern can exist). + # However, this is only in practice (in theory, a definition reference core + # schema could be generated for a simple string schema). + # Note that we avoid calling `self.resolve_ref_schema`, as it might not exist yet. + keys_pattern = None + + values_schema = self.generate_inner(schema['values_schema']).copy() if 'values_schema' in schema else {} + # don't give a title to additionalProperties: + values_schema.pop('title', None) + + if values_schema or keys_pattern is not None: + if keys_pattern is None: + json_schema['additionalProperties'] = values_schema + else: + json_schema['patternProperties'] = {keys_pattern: values_schema} + else: # for `dict[str, Any]`, we allow any key and any value, since `str` is the default key type + json_schema['additionalProperties'] = True + + if ( + # The len check indicates that constraints are probably present: + (keys_schema.get('type') == 'string' and len(keys_schema) > 1) + # If this is a definition reference schema, it most likely has constraints: + or '$ref' in keys_schema + ): + keys_schema.pop('type', None) + json_schema['propertyNames'] = keys_schema + + self.update_with_validations(json_schema, schema, self.ValidationsMapping.object) + return json_schema + + def function_before_schema(self, schema: core_schema.BeforeValidatorFunctionSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a function-before schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + if self.mode == 'validation' and (input_schema := schema.get('json_schema_input_schema')): + return self.generate_inner(input_schema) + + return self.generate_inner(schema['schema']) + + def function_after_schema(self, schema: core_schema.AfterValidatorFunctionSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a function-after schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.generate_inner(schema['schema']) + + def function_plain_schema(self, schema: core_schema.PlainValidatorFunctionSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a function-plain schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + if self.mode == 'validation' and (input_schema := schema.get('json_schema_input_schema')): + return self.generate_inner(input_schema) + + return self.handle_invalid_for_json_schema( + schema, f'core_schema.PlainValidatorFunctionSchema ({schema["function"]})' + ) + + def function_wrap_schema(self, schema: core_schema.WrapValidatorFunctionSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a function-wrap schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + if self.mode == 'validation' and (input_schema := schema.get('json_schema_input_schema')): + return self.generate_inner(input_schema) + + return self.generate_inner(schema['schema']) + + def default_schema(self, schema: core_schema.WithDefaultSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema with a default value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema = self.generate_inner(schema['schema']) + + default = self.get_default_value(schema) + if default is NoDefault: + return json_schema + + # we reflect the application of custom plain, no-info serializers to defaults for + # JSON Schemas viewed in serialization mode: + # TODO: improvements along with https://github.com/pydantic/pydantic/issues/8208 + if ( + self.mode == 'serialization' + and (ser_schema := schema['schema'].get('serialization')) + and (ser_func := ser_schema.get('function')) + and ser_schema.get('type') == 'function-plain' + and not ser_schema.get('info_arg') + and not (default is None and ser_schema.get('when_used') in ('unless-none', 'json-unless-none')) + ): + try: + default = ser_func(default) # type: ignore + except Exception: + # It might be that the provided default needs to be validated (read: parsed) first + # (assuming `validate_default` is enabled). However, we can't perform + # such validation during JSON Schema generation so we don't support + # this pattern for now. + # (One example is when using `foo: ByteSize = '1MB'`, which validates and + # serializes as an int. In this case, `ser_func` is `int` and `int('1MB')` fails). + self.emit_warning( + 'non-serializable-default', + f'Unable to serialize value {default!r} with the plain serializer; excluding default from JSON schema', + ) + return json_schema + + try: + encoded_default = self.encode_default(default) + except pydantic_core.PydanticSerializationError: + self.emit_warning( + 'non-serializable-default', + f'Default value {default} is not JSON serializable; excluding default from JSON schema', + ) + # Return the inner schema, as though there was no default + return json_schema + + json_schema['default'] = encoded_default + return json_schema + + def get_default_value(self, schema: core_schema.WithDefaultSchema) -> Any: + """Get the default value to be used when generating a JSON Schema for a core schema with a default. + + The default implementation is to use the statically defined default value. This method can be overridden + if you want to make use of the default factory. + + Args: + schema: The `'with-default'` core schema. + + Returns: + The default value to use, or [`NoDefault`][pydantic.json_schema.NoDefault] if no default + value is available. + """ + return schema.get('default', NoDefault) + + def nullable_schema(self, schema: core_schema.NullableSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that allows null values. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + null_schema = {'type': 'null'} + inner_json_schema = self.generate_inner(schema['schema']) + + if inner_json_schema == null_schema: + return null_schema + else: + # Thanks to the equality check against `null_schema` above, I think 'oneOf' would also be valid here; + # I'll use 'anyOf' for now, but it could be changed it if it would work better with some external tooling + return self.get_flattened_anyof([inner_json_schema, null_schema]) + + def union_schema(self, schema: core_schema.UnionSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that allows values matching any of the given schemas. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + generated: list[JsonSchemaValue] = [] + + choices = schema['choices'] + for choice in choices: + # choice will be a tuple if an explicit label was provided + choice_schema = choice[0] if isinstance(choice, tuple) else choice + try: + generated.append(self.generate_inner(choice_schema)) + except PydanticOmit: + continue + except PydanticInvalidForJsonSchema as exc: + self.emit_warning('skipped-choice', exc.message) + if len(generated) == 1: + return generated[0] + return self.get_flattened_anyof(generated) + + def tagged_union_schema(self, schema: core_schema.TaggedUnionSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that allows values matching any of the given schemas, where + the schemas are tagged with a discriminator field that indicates which schema should be used to validate + the value. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + generated: dict[str, JsonSchemaValue] = {} + for k, v in schema['choices'].items(): + if isinstance(k, Enum): + k = k.value + try: + # Use str(k) since keys must be strings for json; while not technically correct, + # it's the closest that can be represented in valid JSON + generated[str(k)] = self.generate_inner(v).copy() + except PydanticOmit: + continue + except PydanticInvalidForJsonSchema as exc: + self.emit_warning('skipped-choice', exc.message) + + one_of_choices = _deduplicate_schemas(generated.values()) + json_schema: JsonSchemaValue = {'oneOf': one_of_choices} + + # This reflects the v1 behavior; TODO: we should make it possible to exclude OpenAPI stuff from the JSON schema + openapi_discriminator = self._extract_discriminator(schema, one_of_choices) + if openapi_discriminator is not None: + json_schema['discriminator'] = { + 'propertyName': openapi_discriminator, + 'mapping': {k: v.get('$ref', v) for k, v in generated.items()}, + } + + return json_schema + + def _extract_discriminator( + self, schema: core_schema.TaggedUnionSchema, one_of_choices: list[JsonDict] + ) -> str | None: + """Extract a compatible OpenAPI discriminator from the schema and one_of choices that end up in the final + schema.""" + openapi_discriminator: str | None = None + + if isinstance(schema['discriminator'], str): + return schema['discriminator'] + + if isinstance(schema['discriminator'], list): + # If the discriminator is a single item list containing a string, that is equivalent to the string case + if len(schema['discriminator']) == 1 and isinstance(schema['discriminator'][0], str): + return schema['discriminator'][0] + # When an alias is used that is different from the field name, the discriminator will be a list of single + # str lists, one for the attribute and one for the actual alias. The logic here will work even if there is + # more than one possible attribute, and looks for whether a single alias choice is present as a documented + # property on all choices. If so, that property will be used as the OpenAPI discriminator. + for alias_path in schema['discriminator']: + if not isinstance(alias_path, list): + break # this means that the discriminator is not a list of alias paths + if len(alias_path) != 1: + continue # this means that the "alias" does not represent a single field + alias = alias_path[0] + if not isinstance(alias, str): + continue # this means that the "alias" does not represent a field + alias_is_present_on_all_choices = True + for choice in one_of_choices: + try: + choice = self.resolve_ref_schema(choice) + except RuntimeError as exc: + # TODO: fixme - this is a workaround for the fact that we can't always resolve refs + # for tagged union choices at this point in the schema gen process, we might need to do + # another pass at the end like we do for core schemas + self.emit_warning('skipped-discriminator', str(exc)) + choice = {} + properties = choice.get('properties', {}) + if not isinstance(properties, dict) or alias not in properties: + alias_is_present_on_all_choices = False + break + if alias_is_present_on_all_choices: + openapi_discriminator = alias + break + return openapi_discriminator + + def chain_schema(self, schema: core_schema.ChainSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a core_schema.ChainSchema. + + When generating a schema for validation, we return the validation JSON schema for the first step in the chain. + For serialization, we return the serialization JSON schema for the last step in the chain. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + step_index = 0 if self.mode == 'validation' else -1 # use first step for validation, last for serialization + return self.generate_inner(schema['steps'][step_index]) + + def lax_or_strict_schema(self, schema: core_schema.LaxOrStrictSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that allows values matching either the lax schema or the + strict schema. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + # TODO: Need to read the default value off of model config or whatever + use_strict = schema.get('strict', False) # TODO: replace this default False + # If your JSON schema fails to generate it is probably + # because one of the following two branches failed. + if use_strict: + return self.generate_inner(schema['strict_schema']) + else: + return self.generate_inner(schema['lax_schema']) + + def json_or_python_schema(self, schema: core_schema.JsonOrPythonSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that allows values matching either the JSON schema or the + Python schema. + + The JSON schema is used instead of the Python schema. If you want to use the Python schema, you should override + this method. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.generate_inner(schema['json_schema']) + + def typed_dict_schema(self, schema: core_schema.TypedDictSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a typed dict. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + total = schema.get('total', True) + named_required_fields: list[tuple[str, bool, CoreSchemaField]] = [ + (name, self.field_is_required(field, total), field) + for name, field in schema['fields'].items() + if self.field_is_present(field) + ] + if self.mode == 'serialization': + named_required_fields.extend(self._name_required_computed_fields(schema.get('computed_fields', []))) + cls = schema.get('cls') + config = _get_typed_dict_config(cls) + with self._config_wrapper_stack.push(config): + json_schema = self._named_required_fields_schema(named_required_fields) + + if cls is not None: + self._update_class_schema(json_schema, cls, config) + else: + extra = config.get('extra') + if extra == 'forbid': + json_schema['additionalProperties'] = False + elif extra == 'allow': + json_schema['additionalProperties'] = True + + return json_schema + + @staticmethod + def _name_required_computed_fields( + computed_fields: list[ComputedField], + ) -> list[tuple[str, bool, core_schema.ComputedField]]: + return [(field['property_name'], True, field) for field in computed_fields] + + def _named_required_fields_schema( + self, named_required_fields: Sequence[tuple[str, bool, CoreSchemaField]] + ) -> JsonSchemaValue: + properties: dict[str, JsonSchemaValue] = {} + required_fields: list[str] = [] + for name, required, field in named_required_fields: + if self.by_alias: + name = self._get_alias_name(field, name) + try: + field_json_schema = self.generate_inner(field).copy() + except PydanticOmit: + continue + if 'title' not in field_json_schema and self.field_title_should_be_set(field): + title = self.get_title_from_name(name) + field_json_schema['title'] = title + field_json_schema = self.handle_ref_overrides(field_json_schema) + properties[name] = field_json_schema + if required: + required_fields.append(name) + + json_schema = {'type': 'object', 'properties': properties} + if required_fields: + json_schema['required'] = required_fields + return json_schema + + def _get_alias_name(self, field: CoreSchemaField, name: str) -> str: + if field['type'] == 'computed-field': + alias: Any = field.get('alias', name) + elif self.mode == 'validation': + alias = field.get('validation_alias', name) + else: + alias = field.get('serialization_alias', name) + if isinstance(alias, str): + name = alias + elif isinstance(alias, list): + alias = cast('list[str] | str', alias) + for path in alias: + if isinstance(path, list) and len(path) == 1 and isinstance(path[0], str): + # Use the first valid single-item string path; the code that constructs the alias array + # should ensure the first such item is what belongs in the JSON schema + name = path[0] + break + else: + assert_never(alias) + return name + + def typed_dict_field_schema(self, schema: core_schema.TypedDictField) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a typed dict field. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.generate_inner(schema['schema']) + + def dataclass_field_schema(self, schema: core_schema.DataclassField) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a dataclass field. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.generate_inner(schema['schema']) + + def model_field_schema(self, schema: core_schema.ModelField) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a model field. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.generate_inner(schema['schema']) + + def computed_field_schema(self, schema: core_schema.ComputedField) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a computed field. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.generate_inner(schema['return_schema']) + + def model_schema(self, schema: core_schema.ModelSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a model. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + # We do not use schema['model'].model_json_schema() here + # because it could lead to inconsistent refs handling, etc. + cls = cast('type[BaseModel]', schema['cls']) + config = cls.model_config + + with self._config_wrapper_stack.push(config): + json_schema = self.generate_inner(schema['schema']) + + self._update_class_schema(json_schema, cls, config) + + return json_schema + + def _update_class_schema(self, json_schema: JsonSchemaValue, cls: type[Any], config: ConfigDict) -> None: + """Update json_schema with the following, extracted from `config` and `cls`: + + * title + * description + * additional properties + * json_schema_extra + * deprecated + + Done in place, hence there's no return value as the original json_schema is mutated. + No ref resolving is involved here, as that's not appropriate for simple updates. + """ + from .main import BaseModel + from .root_model import RootModel + + if (config_title := config.get('title')) is not None: + json_schema.setdefault('title', config_title) + elif model_title_generator := config.get('model_title_generator'): + title = model_title_generator(cls) + if not isinstance(title, str): + raise TypeError(f'model_title_generator {model_title_generator} must return str, not {title.__class__}') + json_schema.setdefault('title', title) + if 'title' not in json_schema: + json_schema['title'] = cls.__name__ + + # BaseModel and dataclasses; don't use cls.__doc__ as it will contain the verbose class signature by default + docstring = None if cls is BaseModel or dataclasses.is_dataclass(cls) else cls.__doc__ + + if docstring: + json_schema.setdefault('description', inspect.cleandoc(docstring)) + elif issubclass(cls, RootModel) and (root_description := cls.__pydantic_fields__['root'].description): + json_schema.setdefault('description', root_description) + + extra = config.get('extra') + if 'additionalProperties' not in json_schema: + if extra == 'allow': + json_schema['additionalProperties'] = True + elif extra == 'forbid': + json_schema['additionalProperties'] = False + + json_schema_extra = config.get('json_schema_extra') + if issubclass(cls, BaseModel) and cls.__pydantic_root_model__: + root_json_schema_extra = cls.model_fields['root'].json_schema_extra + if json_schema_extra and root_json_schema_extra: + raise ValueError( + '"model_config[\'json_schema_extra\']" and "Field.json_schema_extra" on "RootModel.root"' + ' field must not be set simultaneously' + ) + if root_json_schema_extra: + json_schema_extra = root_json_schema_extra + + if isinstance(json_schema_extra, (staticmethod, classmethod)): + # In older versions of python, this is necessary to ensure staticmethod/classmethods are callable + json_schema_extra = json_schema_extra.__get__(cls) + + if isinstance(json_schema_extra, dict): + json_schema.update(json_schema_extra) + elif callable(json_schema_extra): + # FIXME: why are there type ignores here? We support two signatures for json_schema_extra callables... + if len(inspect.signature(json_schema_extra).parameters) > 1: + json_schema_extra(json_schema, cls) # type: ignore + else: + json_schema_extra(json_schema) # type: ignore + elif json_schema_extra is not None: + raise ValueError( + f"model_config['json_schema_extra']={json_schema_extra} should be a dict, callable, or None" + ) + + if hasattr(cls, '__deprecated__'): + json_schema['deprecated'] = True + + def resolve_ref_schema(self, json_schema: JsonSchemaValue) -> JsonSchemaValue: + """Resolve a JsonSchemaValue to the non-ref schema if it is a $ref schema. + + Args: + json_schema: The schema to resolve. + + Returns: + The resolved schema. + + Raises: + RuntimeError: If the schema reference can't be found in definitions. + """ + while '$ref' in json_schema: + ref = json_schema['$ref'] + schema_to_update = self.get_schema_from_definitions(JsonRef(ref)) + if schema_to_update is None: + raise RuntimeError(f'Cannot update undefined schema for $ref={ref}') + json_schema = schema_to_update + return json_schema + + def model_fields_schema(self, schema: core_schema.ModelFieldsSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a model's fields. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + named_required_fields: list[tuple[str, bool, CoreSchemaField]] = [ + (name, self.field_is_required(field, total=True), field) + for name, field in schema['fields'].items() + if self.field_is_present(field) + ] + if self.mode == 'serialization': + named_required_fields.extend(self._name_required_computed_fields(schema.get('computed_fields', []))) + json_schema = self._named_required_fields_schema(named_required_fields) + extras_schema = schema.get('extras_schema', None) + if extras_schema is not None: + schema_to_update = self.resolve_ref_schema(json_schema) + schema_to_update['additionalProperties'] = self.generate_inner(extras_schema) + return json_schema + + def field_is_present(self, field: CoreSchemaField) -> bool: + """Whether the field should be included in the generated JSON schema. + + Args: + field: The schema for the field itself. + + Returns: + `True` if the field should be included in the generated JSON schema, `False` otherwise. + """ + if self.mode == 'serialization': + # If you still want to include the field in the generated JSON schema, + # override this method and return True + return not field.get('serialization_exclude') + elif self.mode == 'validation': + return True + else: + assert_never(self.mode) + + def field_is_required( + self, + field: core_schema.ModelField | core_schema.DataclassField | core_schema.TypedDictField, + total: bool, + ) -> bool: + """Whether the field should be marked as required in the generated JSON schema. + (Note that this is irrelevant if the field is not present in the JSON schema.). + + Args: + field: The schema for the field itself. + total: Only applies to `TypedDictField`s. + Indicates if the `TypedDict` this field belongs to is total, in which case any fields that don't + explicitly specify `required=False` are required. + + Returns: + `True` if the field should be marked as required in the generated JSON schema, `False` otherwise. + """ + if self.mode == 'serialization' and self._config.json_schema_serialization_defaults_required: + return not field.get('serialization_exclude') + else: + if field['type'] == 'typed-dict-field': + return field.get('required', total) + else: + return field['schema']['type'] != 'default' + + def dataclass_args_schema(self, schema: core_schema.DataclassArgsSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a dataclass's constructor arguments. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + named_required_fields: list[tuple[str, bool, CoreSchemaField]] = [ + (field['name'], self.field_is_required(field, total=True), field) + for field in schema['fields'] + if self.field_is_present(field) + ] + if self.mode == 'serialization': + named_required_fields.extend(self._name_required_computed_fields(schema.get('computed_fields', []))) + return self._named_required_fields_schema(named_required_fields) + + def dataclass_schema(self, schema: core_schema.DataclassSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a dataclass. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + from ._internal._dataclasses import is_builtin_dataclass + + cls = schema['cls'] + config: ConfigDict = getattr(cls, '__pydantic_config__', cast('ConfigDict', {})) + + with self._config_wrapper_stack.push(config): + json_schema = self.generate_inner(schema['schema']).copy() + + self._update_class_schema(json_schema, cls, config) + + # Dataclass-specific handling of description + if is_builtin_dataclass(cls): + # vanilla dataclass; don't use cls.__doc__ as it will contain the class signature by default + description = None + else: + description = None if cls.__doc__ is None else inspect.cleandoc(cls.__doc__) + if description: + json_schema['description'] = description + + return json_schema + + def arguments_schema(self, schema: core_schema.ArgumentsSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a function's arguments. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + prefer_positional = schema.get('metadata', {}).get('pydantic_js_prefer_positional_arguments') + + arguments = schema['arguments_schema'] + kw_only_arguments = [a for a in arguments if a.get('mode') == 'keyword_only'] + kw_or_p_arguments = [a for a in arguments if a.get('mode') in {'positional_or_keyword', None}] + p_only_arguments = [a for a in arguments if a.get('mode') == 'positional_only'] + var_args_schema = schema.get('var_args_schema') + var_kwargs_schema = schema.get('var_kwargs_schema') + + if prefer_positional: + positional_possible = not kw_only_arguments and not var_kwargs_schema + if positional_possible: + return self.p_arguments_schema(p_only_arguments + kw_or_p_arguments, var_args_schema) + + keyword_possible = not p_only_arguments and not var_args_schema + if keyword_possible: + return self.kw_arguments_schema(kw_or_p_arguments + kw_only_arguments, var_kwargs_schema) + + if not prefer_positional: + positional_possible = not kw_only_arguments and not var_kwargs_schema + if positional_possible: + return self.p_arguments_schema(p_only_arguments + kw_or_p_arguments, var_args_schema) + + raise PydanticInvalidForJsonSchema( + 'Unable to generate JSON schema for arguments validator with positional-only and keyword-only arguments' + ) + + def kw_arguments_schema( + self, arguments: list[core_schema.ArgumentsParameter], var_kwargs_schema: CoreSchema | None + ) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a function's keyword arguments. + + Args: + arguments: The core schema. + + Returns: + The generated JSON schema. + """ + properties: dict[str, JsonSchemaValue] = {} + required: list[str] = [] + for argument in arguments: + name = self.get_argument_name(argument) + argument_schema = self.generate_inner(argument['schema']).copy() + argument_schema['title'] = self.get_title_from_name(name) + properties[name] = argument_schema + + if argument['schema']['type'] != 'default': + # This assumes that if the argument has a default value, + # the inner schema must be of type WithDefaultSchema. + # I believe this is true, but I am not 100% sure + required.append(name) + + json_schema: JsonSchemaValue = {'type': 'object', 'properties': properties} + if required: + json_schema['required'] = required + + if var_kwargs_schema: + additional_properties_schema = self.generate_inner(var_kwargs_schema) + if additional_properties_schema: + json_schema['additionalProperties'] = additional_properties_schema + else: + json_schema['additionalProperties'] = False + return json_schema + + def p_arguments_schema( + self, arguments: list[core_schema.ArgumentsParameter], var_args_schema: CoreSchema | None + ) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a function's positional arguments. + + Args: + arguments: The core schema. + + Returns: + The generated JSON schema. + """ + prefix_items: list[JsonSchemaValue] = [] + min_items = 0 + + for argument in arguments: + name = self.get_argument_name(argument) + + argument_schema = self.generate_inner(argument['schema']).copy() + argument_schema['title'] = self.get_title_from_name(name) + prefix_items.append(argument_schema) + + if argument['schema']['type'] != 'default': + # This assumes that if the argument has a default value, + # the inner schema must be of type WithDefaultSchema. + # I believe this is true, but I am not 100% sure + min_items += 1 + + json_schema: JsonSchemaValue = {'type': 'array'} + if prefix_items: + json_schema['prefixItems'] = prefix_items + if min_items: + json_schema['minItems'] = min_items + + if var_args_schema: + items_schema = self.generate_inner(var_args_schema) + if items_schema: + json_schema['items'] = items_schema + else: + json_schema['maxItems'] = len(prefix_items) + + return json_schema + + def get_argument_name(self, argument: core_schema.ArgumentsParameter | core_schema.ArgumentsV3Parameter) -> str: + """Retrieves the name of an argument. + + Args: + argument: The core schema. + + Returns: + The name of the argument. + """ + name = argument['name'] + if self.by_alias: + alias = argument.get('alias') + if isinstance(alias, str): + name = alias + else: + pass # might want to do something else? + return name + + def arguments_v3_schema(self, schema: core_schema.ArgumentsV3Schema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a function's arguments. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + arguments = schema['arguments_schema'] + properties: dict[str, JsonSchemaValue] = {} + required: list[str] = [] + for argument in arguments: + mode = argument.get('mode', 'positional_or_keyword') + name = self.get_argument_name(argument) + argument_schema = self.generate_inner(argument['schema']).copy() + if mode == 'var_args': + argument_schema = {'type': 'array', 'items': argument_schema} + elif mode == 'var_kwargs_uniform': + argument_schema = {'type': 'object', 'additionalProperties': argument_schema} + + argument_schema.setdefault('title', self.get_title_from_name(name)) + properties[name] = argument_schema + + if ( + (mode == 'var_kwargs_unpacked_typed_dict' and 'required' in argument_schema) + or mode not in {'var_args', 'var_kwargs_uniform', 'var_kwargs_unpacked_typed_dict'} + and argument['schema']['type'] != 'default' + ): + # This assumes that if the argument has a default value, + # the inner schema must be of type WithDefaultSchema. + # I believe this is true, but I am not 100% sure + required.append(name) + + json_schema: JsonSchemaValue = {'type': 'object', 'properties': properties} + if required: + json_schema['required'] = required + return json_schema + + def call_schema(self, schema: core_schema.CallSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a function call. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.generate_inner(schema['arguments_schema']) + + def custom_error_schema(self, schema: core_schema.CustomErrorSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a custom error. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return self.generate_inner(schema['schema']) + + def json_schema(self, schema: core_schema.JsonSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a JSON object. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + content_core_schema = schema.get('schema') or core_schema.any_schema() + content_json_schema = self.generate_inner(content_core_schema) + if self.mode == 'validation': + return {'type': 'string', 'contentMediaType': 'application/json', 'contentSchema': content_json_schema} + else: + # self.mode == 'serialization' + return content_json_schema + + def url_schema(self, schema: core_schema.UrlSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a URL. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + json_schema = {'type': 'string', 'format': 'uri', 'minLength': 1} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.string) + return json_schema + + def multi_host_url_schema(self, schema: core_schema.MultiHostUrlSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a URL that can be used with multiple hosts. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + # Note: 'multi-host-uri' is a custom/pydantic-specific format, not part of the JSON Schema spec + json_schema = {'type': 'string', 'format': 'multi-host-uri', 'minLength': 1} + self.update_with_validations(json_schema, schema, self.ValidationsMapping.string) + return json_schema + + def uuid_schema(self, schema: core_schema.UuidSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a UUID. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return {'type': 'string', 'format': 'uuid'} + + def definitions_schema(self, schema: core_schema.DefinitionsSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that defines a JSON object with definitions. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + for definition in schema['definitions']: + try: + self.generate_inner(definition) + except PydanticInvalidForJsonSchema as e: + core_ref: CoreRef = CoreRef(definition['ref']) # type: ignore + self._core_defs_invalid_for_json_schema[self.get_defs_ref((core_ref, self.mode))] = e + continue + return self.generate_inner(schema['schema']) + + def definition_ref_schema(self, schema: core_schema.DefinitionReferenceSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a schema that references a definition. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + core_ref = CoreRef(schema['schema_ref']) + _, ref_json_schema = self.get_cache_defs_ref_schema(core_ref) + return ref_json_schema + + def ser_schema( + self, schema: core_schema.SerSchema | core_schema.IncExSeqSerSchema | core_schema.IncExDictSerSchema + ) -> JsonSchemaValue | None: + """Generates a JSON schema that matches a schema that defines a serialized object. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + schema_type = schema['type'] + if schema_type == 'function-plain' or schema_type == 'function-wrap': + # PlainSerializerFunctionSerSchema or WrapSerializerFunctionSerSchema + return_schema = schema.get('return_schema') + if return_schema is not None: + return self.generate_inner(return_schema) + elif schema_type == 'format' or schema_type == 'to-string': + # FormatSerSchema or ToStringSerSchema + return self.str_schema(core_schema.str_schema()) + elif schema['type'] == 'model': + # ModelSerSchema + return self.generate_inner(schema['schema']) + return None + + def complex_schema(self, schema: core_schema.ComplexSchema) -> JsonSchemaValue: + """Generates a JSON schema that matches a complex number. + + JSON has no standard way to represent complex numbers. Complex number is not a numeric + type. Here we represent complex number as strings following the rule defined by Python. + For instance, '1+2j' is an accepted complex string. Details can be found in + [Python's `complex` documentation][complex]. + + Args: + schema: The core schema. + + Returns: + The generated JSON schema. + """ + return {'type': 'string'} + + # ### Utility methods + + def get_title_from_name(self, name: str) -> str: + """Retrieves a title from a name. + + Args: + name: The name to retrieve a title from. + + Returns: + The title. + """ + return name.title().replace('_', ' ').strip() + + def field_title_should_be_set(self, schema: CoreSchemaOrField) -> bool: + """Returns true if a field with the given schema should have a title set based on the field name. + + Intuitively, we want this to return true for schemas that wouldn't otherwise provide their own title + (e.g., int, float, str), and false for those that would (e.g., BaseModel subclasses). + + Args: + schema: The schema to check. + + Returns: + `True` if the field should have a title set, `False` otherwise. + """ + if _core_utils.is_core_schema_field(schema): + if schema['type'] == 'computed-field': + field_schema = schema['return_schema'] + else: + field_schema = schema['schema'] + return self.field_title_should_be_set(field_schema) + + elif _core_utils.is_core_schema(schema): + if schema.get('ref'): # things with refs, such as models and enums, should not have titles set + return False + if schema['type'] in {'default', 'nullable', 'definitions'}: + return self.field_title_should_be_set(schema['schema']) # type: ignore[typeddict-item] + if _core_utils.is_function_with_inner_schema(schema): + return self.field_title_should_be_set(schema['schema']) + if schema['type'] == 'definition-ref': + # Referenced schemas should not have titles set for the same reason + # schemas with refs should not + return False + return True # anything else should have title set + + else: + raise PydanticInvalidForJsonSchema(f'Unexpected schema type: schema={schema}') # pragma: no cover + + def normalize_name(self, name: str) -> str: + """Normalizes a name to be used as a key in a dictionary. + + Args: + name: The name to normalize. + + Returns: + The normalized name. + """ + return re.sub(r'[^a-zA-Z0-9.\-_]', '_', name).replace('.', '__') + + def get_defs_ref(self, core_mode_ref: CoreModeRef) -> DefsRef: + """Override this method to change the way that definitions keys are generated from a core reference. + + Args: + core_mode_ref: The core reference. + + Returns: + The definitions key. + """ + # Split the core ref into "components"; generic origins and arguments are each separate components + core_ref, mode = core_mode_ref + components = re.split(r'([\][,])', core_ref) + # Remove IDs from each component + components = [x.rsplit(':', 1)[0] for x in components] + core_ref_no_id = ''.join(components) + # Remove everything before the last period from each "component" + components = [re.sub(r'(?:[^.[\]]+\.)+((?:[^.[\]]+))', r'\1', x) for x in components] + short_ref = ''.join(components) + + mode_title = _MODE_TITLE_MAPPING[mode] + + # It is important that the generated defs_ref values be such that at least one choice will not + # be generated for any other core_ref. Currently, this should be the case because we include + # the id of the source type in the core_ref + name = DefsRef(self.normalize_name(short_ref)) + name_mode = DefsRef(self.normalize_name(short_ref) + f'-{mode_title}') + module_qualname = DefsRef(self.normalize_name(core_ref_no_id)) + module_qualname_mode = DefsRef(f'{module_qualname}-{mode_title}') + module_qualname_id = DefsRef(self.normalize_name(core_ref)) + occurrence_index = self._collision_index.get(module_qualname_id) + if occurrence_index is None: + self._collision_counter[module_qualname] += 1 + occurrence_index = self._collision_index[module_qualname_id] = self._collision_counter[module_qualname] + + module_qualname_occurrence = DefsRef(f'{module_qualname}__{occurrence_index}') + module_qualname_occurrence_mode = DefsRef(f'{module_qualname_mode}__{occurrence_index}') + + self._prioritized_defsref_choices[module_qualname_occurrence_mode] = [ + name, + name_mode, + module_qualname, + module_qualname_mode, + module_qualname_occurrence, + module_qualname_occurrence_mode, + ] + + return module_qualname_occurrence_mode + + def get_cache_defs_ref_schema(self, core_ref: CoreRef) -> tuple[DefsRef, JsonSchemaValue]: + """This method wraps the get_defs_ref method with some cache-lookup/population logic, + and returns both the produced defs_ref and the JSON schema that will refer to the right definition. + + Args: + core_ref: The core reference to get the definitions reference for. + + Returns: + A tuple of the definitions reference and the JSON schema that will refer to it. + """ + core_mode_ref = (core_ref, self.mode) + maybe_defs_ref = self.core_to_defs_refs.get(core_mode_ref) + if maybe_defs_ref is not None: + json_ref = self.core_to_json_refs[core_mode_ref] + return maybe_defs_ref, {'$ref': json_ref} + + defs_ref = self.get_defs_ref(core_mode_ref) + + # populate the ref translation mappings + self.core_to_defs_refs[core_mode_ref] = defs_ref + self.defs_to_core_refs[defs_ref] = core_mode_ref + + json_ref = JsonRef(self.ref_template.format(model=defs_ref)) + self.core_to_json_refs[core_mode_ref] = json_ref + self.json_to_defs_refs[json_ref] = defs_ref + ref_json_schema = {'$ref': json_ref} + return defs_ref, ref_json_schema + + def handle_ref_overrides(self, json_schema: JsonSchemaValue) -> JsonSchemaValue: + """Remove any sibling keys that are redundant with the referenced schema. + + Args: + json_schema: The schema to remove redundant sibling keys from. + + Returns: + The schema with redundant sibling keys removed. + """ + if '$ref' in json_schema: + # prevent modifications to the input; this copy may be safe to drop if there is significant overhead + json_schema = json_schema.copy() + + referenced_json_schema = self.get_schema_from_definitions(JsonRef(json_schema['$ref'])) + if referenced_json_schema is None: + # This can happen when building schemas for models with not-yet-defined references. + # It may be a good idea to do a recursive pass at the end of the generation to remove + # any redundant override keys. + return json_schema + for k, v in list(json_schema.items()): + if k == '$ref': + continue + if k in referenced_json_schema and referenced_json_schema[k] == v: + del json_schema[k] # redundant key + + return json_schema + + def get_schema_from_definitions(self, json_ref: JsonRef) -> JsonSchemaValue | None: + try: + def_ref = self.json_to_defs_refs[json_ref] + if def_ref in self._core_defs_invalid_for_json_schema: + raise self._core_defs_invalid_for_json_schema[def_ref] + return self.definitions.get(def_ref, None) + except KeyError: + if json_ref.startswith(('http://', 'https://')): + return None + raise + + def encode_default(self, dft: Any) -> Any: + """Encode a default value to a JSON-serializable value. + + This is used to encode default values for fields in the generated JSON schema. + + Args: + dft: The default value to encode. + + Returns: + The encoded default value. + """ + from .type_adapter import TypeAdapter, _type_has_config + + config = self._config + try: + default = ( + dft + if _type_has_config(type(dft)) + else TypeAdapter(type(dft), config=config.config_dict).dump_python( + dft, by_alias=self.by_alias, mode='json' + ) + ) + except PydanticSchemaGenerationError: + raise pydantic_core.PydanticSerializationError(f'Unable to encode default value {dft}') + + return pydantic_core.to_jsonable_python( + default, timedelta_mode=config.ser_json_timedelta, bytes_mode=config.ser_json_bytes, by_alias=self.by_alias + ) + + def update_with_validations( + self, json_schema: JsonSchemaValue, core_schema: CoreSchema, mapping: dict[str, str] + ) -> None: + """Update the json_schema with the corresponding validations specified in the core_schema, + using the provided mapping to translate keys in core_schema to the appropriate keys for a JSON schema. + + Args: + json_schema: The JSON schema to update. + core_schema: The core schema to get the validations from. + mapping: A mapping from core_schema attribute names to the corresponding JSON schema attribute names. + """ + for core_key, json_schema_key in mapping.items(): + if core_key in core_schema: + json_schema[json_schema_key] = core_schema[core_key] + + class ValidationsMapping: + """This class just contains mappings from core_schema attribute names to the corresponding + JSON schema attribute names. While I suspect it is unlikely to be necessary, you can in + principle override this class in a subclass of GenerateJsonSchema (by inheriting from + GenerateJsonSchema.ValidationsMapping) to change these mappings. + """ + + numeric = { + 'multiple_of': 'multipleOf', + 'le': 'maximum', + 'ge': 'minimum', + 'lt': 'exclusiveMaximum', + 'gt': 'exclusiveMinimum', + } + bytes = { + 'min_length': 'minLength', + 'max_length': 'maxLength', + } + string = { + 'min_length': 'minLength', + 'max_length': 'maxLength', + 'pattern': 'pattern', + } + array = { + 'min_length': 'minItems', + 'max_length': 'maxItems', + } + object = { + 'min_length': 'minProperties', + 'max_length': 'maxProperties', + } + + def get_flattened_anyof(self, schemas: list[JsonSchemaValue]) -> JsonSchemaValue: + members = [] + for schema in schemas: + if len(schema) == 1 and 'anyOf' in schema: + members.extend(schema['anyOf']) + else: + members.append(schema) + members = _deduplicate_schemas(members) + if len(members) == 1: + return members[0] + return {'anyOf': members} + + def get_json_ref_counts(self, json_schema: JsonSchemaValue) -> dict[JsonRef, int]: + """Get all values corresponding to the key '$ref' anywhere in the json_schema.""" + json_refs: dict[JsonRef, int] = Counter() + + def _add_json_refs(schema: Any) -> None: + if isinstance(schema, dict): + if '$ref' in schema: + json_ref = JsonRef(schema['$ref']) + if not isinstance(json_ref, str): + return # in this case, '$ref' might have been the name of a property + already_visited = json_ref in json_refs + json_refs[json_ref] += 1 + if already_visited: + return # prevent recursion on a definition that was already visited + try: + defs_ref = self.json_to_defs_refs[json_ref] + if defs_ref in self._core_defs_invalid_for_json_schema: + raise self._core_defs_invalid_for_json_schema[defs_ref] + _add_json_refs(self.definitions[defs_ref]) + except KeyError: + if not json_ref.startswith(('http://', 'https://')): + raise + + for k, v in schema.items(): + if k == 'examples' and isinstance(v, list): + # Skip examples that may contain arbitrary values and references + # (see the comment in `_get_all_json_refs` for more details). + continue + _add_json_refs(v) + elif isinstance(schema, list): + for v in schema: + _add_json_refs(v) + + _add_json_refs(json_schema) + return json_refs + + def handle_invalid_for_json_schema(self, schema: CoreSchemaOrField, error_info: str) -> JsonSchemaValue: + raise PydanticInvalidForJsonSchema(f'Cannot generate a JsonSchema for {error_info}') + + def emit_warning(self, kind: JsonSchemaWarningKind, detail: str) -> None: + """This method simply emits PydanticJsonSchemaWarnings based on handling in the `warning_message` method.""" + message = self.render_warning_message(kind, detail) + if message is not None: + warnings.warn(message, PydanticJsonSchemaWarning) + + def render_warning_message(self, kind: JsonSchemaWarningKind, detail: str) -> str | None: + """This method is responsible for ignoring warnings as desired, and for formatting the warning messages. + + You can override the value of `ignored_warning_kinds` in a subclass of GenerateJsonSchema + to modify what warnings are generated. If you want more control, you can override this method; + just return None in situations where you don't want warnings to be emitted. + + Args: + kind: The kind of warning to render. It can be one of the following: + + - 'skipped-choice': A choice field was skipped because it had no valid choices. + - 'non-serializable-default': A default value was skipped because it was not JSON-serializable. + detail: A string with additional details about the warning. + + Returns: + The formatted warning message, or `None` if no warning should be emitted. + """ + if kind in self.ignored_warning_kinds: + return None + return f'{detail} [{kind}]' + + def _build_definitions_remapping(self) -> _DefinitionsRemapping: + defs_to_json: dict[DefsRef, JsonRef] = {} + for defs_refs in self._prioritized_defsref_choices.values(): + for defs_ref in defs_refs: + json_ref = JsonRef(self.ref_template.format(model=defs_ref)) + defs_to_json[defs_ref] = json_ref + + return _DefinitionsRemapping.from_prioritized_choices( + self._prioritized_defsref_choices, defs_to_json, self.definitions + ) + + def _garbage_collect_definitions(self, schema: JsonSchemaValue) -> None: + visited_defs_refs: set[DefsRef] = set() + unvisited_json_refs = _get_all_json_refs(schema) + while unvisited_json_refs: + next_json_ref = unvisited_json_refs.pop() + try: + next_defs_ref = self.json_to_defs_refs[next_json_ref] + if next_defs_ref in visited_defs_refs: + continue + visited_defs_refs.add(next_defs_ref) + unvisited_json_refs.update(_get_all_json_refs(self.definitions[next_defs_ref])) + except KeyError: + if not next_json_ref.startswith(('http://', 'https://')): + raise + + self.definitions = {k: v for k, v in self.definitions.items() if k in visited_defs_refs} + + +# ##### Start JSON Schema Generation Functions ##### + + +def model_json_schema( + cls: type[BaseModel] | type[PydanticDataclass], + by_alias: bool = True, + ref_template: str = DEFAULT_REF_TEMPLATE, + schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, + mode: JsonSchemaMode = 'validation', +) -> dict[str, Any]: + """Utility function to generate a JSON Schema for a model. + + Args: + cls: The model class to generate a JSON Schema for. + by_alias: If `True` (the default), fields will be serialized according to their alias. + If `False`, fields will be serialized according to their attribute name. + ref_template: The template to use for generating JSON Schema references. + schema_generator: The class to use for generating the JSON Schema. + mode: The mode to use for generating the JSON Schema. It can be one of the following: + + - 'validation': Generate a JSON Schema for validating data. + - 'serialization': Generate a JSON Schema for serializing data. + + Returns: + The generated JSON Schema. + """ + from .main import BaseModel + + schema_generator_instance = schema_generator(by_alias=by_alias, ref_template=ref_template) + + if isinstance(cls.__pydantic_core_schema__, _mock_val_ser.MockCoreSchema): + cls.__pydantic_core_schema__.rebuild() + + if cls is BaseModel: + raise AttributeError('model_json_schema() must be called on a subclass of BaseModel, not BaseModel itself.') + + assert not isinstance(cls.__pydantic_core_schema__, _mock_val_ser.MockCoreSchema), 'this is a bug! please report it' + return schema_generator_instance.generate(cls.__pydantic_core_schema__, mode=mode) + + +def models_json_schema( + models: Sequence[tuple[type[BaseModel] | type[PydanticDataclass], JsonSchemaMode]], + *, + by_alias: bool = True, + title: str | None = None, + description: str | None = None, + ref_template: str = DEFAULT_REF_TEMPLATE, + schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, +) -> tuple[dict[tuple[type[BaseModel] | type[PydanticDataclass], JsonSchemaMode], JsonSchemaValue], JsonSchemaValue]: + """Utility function to generate a JSON Schema for multiple models. + + Args: + models: A sequence of tuples of the form (model, mode). + by_alias: Whether field aliases should be used as keys in the generated JSON Schema. + title: The title of the generated JSON Schema. + description: The description of the generated JSON Schema. + ref_template: The reference template to use for generating JSON Schema references. + schema_generator: The schema generator to use for generating the JSON Schema. + + Returns: + A tuple where: + - The first element is a dictionary whose keys are tuples of JSON schema key type and JSON mode, and + whose values are the JSON schema corresponding to that pair of inputs. (These schemas may have + JsonRef references to definitions that are defined in the second returned element.) + - The second element is a JSON schema containing all definitions referenced in the first returned + element, along with the optional title and description keys. + """ + for cls, _ in models: + if isinstance(cls.__pydantic_core_schema__, _mock_val_ser.MockCoreSchema): + cls.__pydantic_core_schema__.rebuild() + + instance = schema_generator(by_alias=by_alias, ref_template=ref_template) + inputs: list[tuple[type[BaseModel] | type[PydanticDataclass], JsonSchemaMode, CoreSchema]] = [ + (m, mode, m.__pydantic_core_schema__) for m, mode in models + ] + json_schemas_map, definitions = instance.generate_definitions(inputs) + + json_schema: dict[str, Any] = {} + if definitions: + json_schema['$defs'] = definitions + if title: + json_schema['title'] = title + if description: + json_schema['description'] = description + + return json_schemas_map, json_schema + + +# ##### End JSON Schema Generation Functions ##### + + +_HashableJsonValue: TypeAlias = Union[ + int, float, str, bool, None, tuple['_HashableJsonValue', ...], tuple[tuple[str, '_HashableJsonValue'], ...] +] + + +def _deduplicate_schemas(schemas: Iterable[JsonDict]) -> list[JsonDict]: + return list({_make_json_hashable(schema): schema for schema in schemas}.values()) + + +def _make_json_hashable(value: JsonValue) -> _HashableJsonValue: + if isinstance(value, dict): + return tuple(sorted((k, _make_json_hashable(v)) for k, v in value.items())) + elif isinstance(value, list): + return tuple(_make_json_hashable(v) for v in value) + else: + return value + + +@dataclasses.dataclass(**_internal_dataclass.slots_true) +class WithJsonSchema: + """!!! abstract "Usage Documentation" + [`WithJsonSchema` Annotation](../concepts/json_schema.md#withjsonschema-annotation) + + Add this as an annotation on a field to override the (base) JSON schema that would be generated for that field. + This provides a way to set a JSON schema for types that would otherwise raise errors when producing a JSON schema, + such as Callable, or types that have an is-instance core schema, without needing to go so far as creating a + custom subclass of pydantic.json_schema.GenerateJsonSchema. + Note that any _modifications_ to the schema that would normally be made (such as setting the title for model fields) + will still be performed. + + If `mode` is set this will only apply to that schema generation mode, allowing you + to set different json schemas for validation and serialization. + """ + + json_schema: JsonSchemaValue | None + mode: Literal['validation', 'serialization'] | None = None + + def __get_pydantic_json_schema__( + self, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler + ) -> JsonSchemaValue: + mode = self.mode or handler.mode + if mode != handler.mode: + return handler(core_schema) + if self.json_schema is None: + # This exception is handled in pydantic.json_schema.GenerateJsonSchema._named_required_fields_schema + raise PydanticOmit + else: + return self.json_schema.copy() + + def __hash__(self) -> int: + return hash(type(self.mode)) + + +class Examples: + """Add examples to a JSON schema. + + If the JSON Schema already contains examples, the provided examples + will be appended. + + If `mode` is set this will only apply to that schema generation mode, + allowing you to add different examples for validation and serialization. + """ + + @overload + @deprecated('Using a dict for `examples` is deprecated since v2.9 and will be removed in v3.0. Use a list instead.') + def __init__( + self, examples: dict[str, Any], mode: Literal['validation', 'serialization'] | None = None + ) -> None: ... + + @overload + def __init__(self, examples: list[Any], mode: Literal['validation', 'serialization'] | None = None) -> None: ... + + def __init__( + self, examples: dict[str, Any] | list[Any], mode: Literal['validation', 'serialization'] | None = None + ) -> None: + if isinstance(examples, dict): + warnings.warn( + 'Using a dict for `examples` is deprecated, use a list instead.', + PydanticDeprecatedSince29, + stacklevel=2, + ) + self.examples = examples + self.mode = mode + + def __get_pydantic_json_schema__( + self, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler + ) -> JsonSchemaValue: + mode = self.mode or handler.mode + json_schema = handler(core_schema) + if mode != handler.mode: + return json_schema + examples = json_schema.get('examples') + if examples is None: + json_schema['examples'] = to_jsonable_python(self.examples) + if isinstance(examples, dict): + if isinstance(self.examples, list): + warnings.warn( + 'Updating existing JSON Schema examples of type dict with examples of type list. ' + 'Only the existing examples values will be retained. Note that dict support for ' + 'examples is deprecated and will be removed in v3.0.', + UserWarning, + ) + json_schema['examples'] = to_jsonable_python( + [ex for value in examples.values() for ex in value] + self.examples + ) + else: + json_schema['examples'] = to_jsonable_python({**examples, **self.examples}) + if isinstance(examples, list): + if isinstance(self.examples, list): + json_schema['examples'] = to_jsonable_python(examples + self.examples) + elif isinstance(self.examples, dict): + warnings.warn( + 'Updating existing JSON Schema examples of type list with examples of type dict. ' + 'Only the examples values will be retained. Note that dict support for ' + 'examples is deprecated and will be removed in v3.0.', + UserWarning, + ) + json_schema['examples'] = to_jsonable_python( + examples + [ex for value in self.examples.values() for ex in value] + ) + + return json_schema + + def __hash__(self) -> int: + return hash(type(self.mode)) + + +def _get_all_json_refs(item: Any) -> set[JsonRef]: + """Get all the definitions references from a JSON schema.""" + refs: set[JsonRef] = set() + stack = [item] + + while stack: + current = stack.pop() + if isinstance(current, dict): + for key, value in current.items(): + if key == 'examples' and isinstance(value, list): + # Skip examples that may contain arbitrary values and references + # (e.g. `{"examples": [{"$ref": "..."}]}`). Note: checking for value + # of type list is necessary to avoid skipping valid portions of the schema, + # for instance when "examples" is used as a property key. A more robust solution + # could be found, but would require more advanced JSON Schema parsing logic. + continue + if key == '$ref' and isinstance(value, str): + refs.add(JsonRef(value)) + elif isinstance(value, dict): + stack.append(value) + elif isinstance(value, list): + stack.extend(value) + elif isinstance(current, list): + stack.extend(current) + + return refs + + +AnyType = TypeVar('AnyType') + +if TYPE_CHECKING: + SkipJsonSchema = Annotated[AnyType, ...] +else: + + @dataclasses.dataclass(**_internal_dataclass.slots_true) + class SkipJsonSchema: + """!!! abstract "Usage Documentation" + [`SkipJsonSchema` Annotation](../concepts/json_schema.md#skipjsonschema-annotation) + + Add this as an annotation on a field to skip generating a JSON schema for that field. + + Example: + ```python + from pprint import pprint + from typing import Union + + from pydantic import BaseModel + from pydantic.json_schema import SkipJsonSchema + + class Model(BaseModel): + a: Union[int, None] = None # (1)! + b: Union[int, SkipJsonSchema[None]] = None # (2)! + c: SkipJsonSchema[Union[int, None]] = None # (3)! + + pprint(Model.model_json_schema()) + ''' + { + 'properties': { + 'a': { + 'anyOf': [ + {'type': 'integer'}, + {'type': 'null'} + ], + 'default': None, + 'title': 'A' + }, + 'b': { + 'default': None, + 'title': 'B', + 'type': 'integer' + } + }, + 'title': 'Model', + 'type': 'object' + } + ''' + ``` + + 1. The integer and null types are both included in the schema for `a`. + 2. The integer type is the only type included in the schema for `b`. + 3. The entirety of the `c` field is omitted from the schema. + """ + + def __class_getitem__(cls, item: AnyType) -> AnyType: + return Annotated[item, cls()] + + def __get_pydantic_json_schema__( + self, core_schema: CoreSchema, handler: GetJsonSchemaHandler + ) -> JsonSchemaValue: + raise PydanticOmit + + def __hash__(self) -> int: + return hash(type(self)) + + +def _get_typed_dict_config(cls: type[Any] | None) -> ConfigDict: + if cls is not None: + try: + return _decorators.get_attribute_from_bases(cls, '__pydantic_config__') + except AttributeError: + pass + return {} diff --git a/lib/python3.12/site-packages/pydantic/main.py b/lib/python3.12/site-packages/pydantic/main.py new file mode 100644 index 0000000000000000000000000000000000000000..272c10aa4e66b1116b34cb96ad1121cfef500954 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/main.py @@ -0,0 +1,1773 @@ +"""Logic for creating models.""" + +# Because `dict` is in the local namespace of the `BaseModel` class, we use `Dict` for annotations. +# TODO v3 fallback to `dict` when the deprecated `dict` method gets removed. +# ruff: noqa: UP035 + +from __future__ import annotations as _annotations + +import operator +import sys +import types +import typing +import warnings +from collections.abc import Generator, Mapping +from copy import copy, deepcopy +from functools import cached_property +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ClassVar, + Dict, + Literal, + TypeVar, + Union, + cast, + overload, +) + +import pydantic_core +import typing_extensions +from pydantic_core import PydanticUndefined, ValidationError +from typing_extensions import Self, TypeAlias, Unpack + +from . import PydanticDeprecatedSince20, PydanticDeprecatedSince211 +from ._internal import ( + _config, + _decorators, + _fields, + _forward_ref, + _generics, + _mock_val_ser, + _model_construction, + _namespace_utils, + _repr, + _typing_extra, + _utils, +) +from ._migration import getattr_migration +from .aliases import AliasChoices, AliasPath +from .annotated_handlers import GetCoreSchemaHandler, GetJsonSchemaHandler +from .config import ConfigDict +from .errors import PydanticUndefinedAnnotation, PydanticUserError +from .json_schema import DEFAULT_REF_TEMPLATE, GenerateJsonSchema, JsonSchemaMode, JsonSchemaValue, model_json_schema +from .plugin._schema_validator import PluggableSchemaValidator + +if TYPE_CHECKING: + from inspect import Signature + from pathlib import Path + + from pydantic_core import CoreSchema, SchemaSerializer, SchemaValidator + + from ._internal._namespace_utils import MappingNamespace + from ._internal._utils import AbstractSetIntStr, MappingIntStrAny + from .deprecated.parse import Protocol as DeprecatedParseProtocol + from .fields import ComputedFieldInfo, FieldInfo, ModelPrivateAttr +else: + # See PyCharm issues https://youtrack.jetbrains.com/issue/PY-21915 + # and https://youtrack.jetbrains.com/issue/PY-51428 + DeprecationWarning = PydanticDeprecatedSince20 + +__all__ = 'BaseModel', 'create_model' + +# Keep these type aliases available at runtime: +TupleGenerator: TypeAlias = Generator[tuple[str, Any], None, None] +# NOTE: In reality, `bool` should be replaced by `Literal[True]` but mypy fails to correctly apply bidirectional +# type inference (e.g. when using `{'a': {'b': True}}`): +# NOTE: Keep this type alias in sync with the stub definition in `pydantic-core`: +IncEx: TypeAlias = Union[set[int], set[str], Mapping[int, Union['IncEx', bool]], Mapping[str, Union['IncEx', bool]]] + +_object_setattr = _model_construction.object_setattr + + +def _check_frozen(model_cls: type[BaseModel], name: str, value: Any) -> None: + if model_cls.model_config.get('frozen'): + error_type = 'frozen_instance' + elif getattr(model_cls.__pydantic_fields__.get(name), 'frozen', False): + error_type = 'frozen_field' + else: + return + + raise ValidationError.from_exception_data( + model_cls.__name__, [{'type': error_type, 'loc': (name,), 'input': value}] + ) + + +def _model_field_setattr_handler(model: BaseModel, name: str, val: Any) -> None: + model.__dict__[name] = val + model.__pydantic_fields_set__.add(name) + + +def _private_setattr_handler(model: BaseModel, name: str, val: Any) -> None: + if getattr(model, '__pydantic_private__', None) is None: + # While the attribute should be present at this point, this may not be the case if + # users do unusual stuff with `model_post_init()` (which is where the `__pydantic_private__` + # is initialized, by wrapping the user-defined `model_post_init()`), e.g. if they mock + # the `model_post_init()` call. Ideally we should find a better way to init private attrs. + object.__setattr__(model, '__pydantic_private__', {}) + model.__pydantic_private__[name] = val # pyright: ignore[reportOptionalSubscript] + + +_SIMPLE_SETATTR_HANDLERS: Mapping[str, Callable[[BaseModel, str, Any], None]] = { + 'model_field': _model_field_setattr_handler, + 'validate_assignment': lambda model, name, val: model.__pydantic_validator__.validate_assignment(model, name, val), # pyright: ignore[reportAssignmentType] + 'private': _private_setattr_handler, + 'cached_property': lambda model, name, val: model.__dict__.__setitem__(name, val), + 'extra_known': lambda model, name, val: _object_setattr(model, name, val), +} + + +class BaseModel(metaclass=_model_construction.ModelMetaclass): + """!!! abstract "Usage Documentation" + [Models](../concepts/models.md) + + A base class for creating Pydantic models. + + Attributes: + __class_vars__: The names of the class variables defined on the model. + __private_attributes__: Metadata about the private attributes of the model. + __signature__: The synthesized `__init__` [`Signature`][inspect.Signature] of the model. + + __pydantic_complete__: Whether model building is completed, or if there are still undefined fields. + __pydantic_core_schema__: The core schema of the model. + __pydantic_custom_init__: Whether the model has a custom `__init__` function. + __pydantic_decorators__: Metadata containing the decorators defined on the model. + This replaces `Model.__validators__` and `Model.__root_validators__` from Pydantic V1. + __pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to + __args__, __origin__, __parameters__ in typing-module generics. May eventually be replaced by these. + __pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models. + __pydantic_post_init__: The name of the post-init method for the model, if defined. + __pydantic_root_model__: Whether the model is a [`RootModel`][pydantic.root_model.RootModel]. + __pydantic_serializer__: The `pydantic-core` `SchemaSerializer` used to dump instances of the model. + __pydantic_validator__: The `pydantic-core` `SchemaValidator` used to validate instances of the model. + + __pydantic_fields__: A dictionary of field names and their corresponding [`FieldInfo`][pydantic.fields.FieldInfo] objects. + __pydantic_computed_fields__: A dictionary of computed field names and their corresponding [`ComputedFieldInfo`][pydantic.fields.ComputedFieldInfo] objects. + + __pydantic_extra__: A dictionary containing extra values, if [`extra`][pydantic.config.ConfigDict.extra] + is set to `'allow'`. + __pydantic_fields_set__: The names of fields explicitly set during instantiation. + __pydantic_private__: Values of private attributes set on the model instance. + """ + + # Note: Many of the below class vars are defined in the metaclass, but we define them here for type checking purposes. + + model_config: ClassVar[ConfigDict] = ConfigDict() + """ + Configuration for the model, should be a dictionary conforming to [`ConfigDict`][pydantic.config.ConfigDict]. + """ + + __class_vars__: ClassVar[set[str]] + """The names of the class variables defined on the model.""" + + __private_attributes__: ClassVar[Dict[str, ModelPrivateAttr]] # noqa: UP006 + """Metadata about the private attributes of the model.""" + + __signature__: ClassVar[Signature] + """The synthesized `__init__` [`Signature`][inspect.Signature] of the model.""" + + __pydantic_complete__: ClassVar[bool] = False + """Whether model building is completed, or if there are still undefined fields.""" + + __pydantic_core_schema__: ClassVar[CoreSchema] + """The core schema of the model.""" + + __pydantic_custom_init__: ClassVar[bool] + """Whether the model has a custom `__init__` method.""" + + # Must be set for `GenerateSchema.model_schema` to work for a plain `BaseModel` annotation. + __pydantic_decorators__: ClassVar[_decorators.DecoratorInfos] = _decorators.DecoratorInfos() + """Metadata containing the decorators defined on the model. + This replaces `Model.__validators__` and `Model.__root_validators__` from Pydantic V1.""" + + __pydantic_generic_metadata__: ClassVar[_generics.PydanticGenericMetadata] + """Metadata for generic models; contains data used for a similar purpose to + __args__, __origin__, __parameters__ in typing-module generics. May eventually be replaced by these.""" + + __pydantic_parent_namespace__: ClassVar[Dict[str, Any] | None] = None # noqa: UP006 + """Parent namespace of the model, used for automatic rebuilding of models.""" + + __pydantic_post_init__: ClassVar[None | Literal['model_post_init']] + """The name of the post-init method for the model, if defined.""" + + __pydantic_root_model__: ClassVar[bool] = False + """Whether the model is a [`RootModel`][pydantic.root_model.RootModel].""" + + __pydantic_serializer__: ClassVar[SchemaSerializer] + """The `pydantic-core` `SchemaSerializer` used to dump instances of the model.""" + + __pydantic_validator__: ClassVar[SchemaValidator | PluggableSchemaValidator] + """The `pydantic-core` `SchemaValidator` used to validate instances of the model.""" + + __pydantic_fields__: ClassVar[Dict[str, FieldInfo]] # noqa: UP006 + """A dictionary of field names and their corresponding [`FieldInfo`][pydantic.fields.FieldInfo] objects. + This replaces `Model.__fields__` from Pydantic V1. + """ + + __pydantic_setattr_handlers__: ClassVar[Dict[str, Callable[[BaseModel, str, Any], None]]] # noqa: UP006 + """`__setattr__` handlers. Memoizing the handlers leads to a dramatic performance improvement in `__setattr__`""" + + __pydantic_computed_fields__: ClassVar[Dict[str, ComputedFieldInfo]] # noqa: UP006 + """A dictionary of computed field names and their corresponding [`ComputedFieldInfo`][pydantic.fields.ComputedFieldInfo] objects.""" + + __pydantic_extra__: dict[str, Any] | None = _model_construction.NoInitField(init=False) + """A dictionary containing extra values, if [`extra`][pydantic.config.ConfigDict.extra] is set to `'allow'`.""" + + __pydantic_fields_set__: set[str] = _model_construction.NoInitField(init=False) + """The names of fields explicitly set during instantiation.""" + + __pydantic_private__: dict[str, Any] | None = _model_construction.NoInitField(init=False) + """Values of private attributes set on the model instance.""" + + if not TYPE_CHECKING: + # Prevent `BaseModel` from being instantiated directly + # (defined in an `if not TYPE_CHECKING` block for clarity and to avoid type checking errors): + __pydantic_core_schema__ = _mock_val_ser.MockCoreSchema( + 'Pydantic models should inherit from BaseModel, BaseModel cannot be instantiated directly', + code='base-model-instantiated', + ) + __pydantic_validator__ = _mock_val_ser.MockValSer( + 'Pydantic models should inherit from BaseModel, BaseModel cannot be instantiated directly', + val_or_ser='validator', + code='base-model-instantiated', + ) + __pydantic_serializer__ = _mock_val_ser.MockValSer( + 'Pydantic models should inherit from BaseModel, BaseModel cannot be instantiated directly', + val_or_ser='serializer', + code='base-model-instantiated', + ) + + __slots__ = '__dict__', '__pydantic_fields_set__', '__pydantic_extra__', '__pydantic_private__' + + def __init__(self, /, **data: Any) -> None: + """Create a new model by parsing and validating input data from keyword arguments. + + Raises [`ValidationError`][pydantic_core.ValidationError] if the input data cannot be + validated to form a valid model. + + `self` is explicitly positional-only to allow `self` as a field name. + """ + # `__tracebackhide__` tells pytest and some other tools to omit this function from tracebacks + __tracebackhide__ = True + validated_self = self.__pydantic_validator__.validate_python(data, self_instance=self) + if self is not validated_self: + warnings.warn( + 'A custom validator is returning a value other than `self`.\n' + "Returning anything other than `self` from a top level model validator isn't supported when validating via `__init__`.\n" + 'See the `model_validator` docs (https://docs.pydantic.dev/latest/concepts/validators/#model-validators) for more details.', + stacklevel=2, + ) + + # The following line sets a flag that we use to determine when `__init__` gets overridden by the user + __init__.__pydantic_base_init__ = True # pyright: ignore[reportFunctionMemberAccess] + + @_utils.deprecated_instance_property + @classmethod + def model_fields(cls) -> dict[str, FieldInfo]: + """A mapping of field names to their respective [`FieldInfo`][pydantic.fields.FieldInfo] instances. + + !!! warning + Accessing this attribute from a model instance is deprecated, and will not work in Pydantic V3. + Instead, you should access this attribute from the model class. + """ + return getattr(cls, '__pydantic_fields__', {}) + + @_utils.deprecated_instance_property + @classmethod + def model_computed_fields(cls) -> dict[str, ComputedFieldInfo]: + """A mapping of computed field names to their respective [`ComputedFieldInfo`][pydantic.fields.ComputedFieldInfo] instances. + + !!! warning + Accessing this attribute from a model instance is deprecated, and will not work in Pydantic V3. + Instead, you should access this attribute from the model class. + """ + return getattr(cls, '__pydantic_computed_fields__', {}) + + @property + def model_extra(self) -> dict[str, Any] | None: + """Get extra fields set during validation. + + Returns: + A dictionary of extra fields, or `None` if `config.extra` is not set to `"allow"`. + """ + return self.__pydantic_extra__ + + @property + def model_fields_set(self) -> set[str]: + """Returns the set of fields that have been explicitly set on this model instance. + + Returns: + A set of strings representing the fields that have been set, + i.e. that were not filled from defaults. + """ + return self.__pydantic_fields_set__ + + @classmethod + def model_construct(cls, _fields_set: set[str] | None = None, **values: Any) -> Self: # noqa: C901 + """Creates a new instance of the `Model` class with validated data. + + Creates a new model setting `__dict__` and `__pydantic_fields_set__` from trusted or pre-validated data. + Default values are respected, but no other validation is performed. + + !!! note + `model_construct()` generally respects the `model_config.extra` setting on the provided model. + That is, if `model_config.extra == 'allow'`, then all extra passed values are added to the model instance's `__dict__` + and `__pydantic_extra__` fields. If `model_config.extra == 'ignore'` (the default), then all extra passed values are ignored. + Because no validation is performed with a call to `model_construct()`, having `model_config.extra == 'forbid'` does not result in + an error if extra values are passed, but they will be ignored. + + Args: + _fields_set: A set of field names that were originally explicitly set during instantiation. If provided, + this is directly used for the [`model_fields_set`][pydantic.BaseModel.model_fields_set] attribute. + Otherwise, the field names from the `values` argument will be used. + values: Trusted or pre-validated data dictionary. + + Returns: + A new instance of the `Model` class with validated data. + """ + m = cls.__new__(cls) + fields_values: dict[str, Any] = {} + fields_set = set() + + for name, field in cls.__pydantic_fields__.items(): + if field.alias is not None and field.alias in values: + fields_values[name] = values.pop(field.alias) + fields_set.add(name) + + if (name not in fields_set) and (field.validation_alias is not None): + validation_aliases: list[str | AliasPath] = ( + field.validation_alias.choices + if isinstance(field.validation_alias, AliasChoices) + else [field.validation_alias] + ) + + for alias in validation_aliases: + if isinstance(alias, str) and alias in values: + fields_values[name] = values.pop(alias) + fields_set.add(name) + break + elif isinstance(alias, AliasPath): + value = alias.search_dict_for_path(values) + if value is not PydanticUndefined: + fields_values[name] = value + fields_set.add(name) + break + + if name not in fields_set: + if name in values: + fields_values[name] = values.pop(name) + fields_set.add(name) + elif not field.is_required(): + fields_values[name] = field.get_default(call_default_factory=True, validated_data=fields_values) + if _fields_set is None: + _fields_set = fields_set + + _extra: dict[str, Any] | None = values if cls.model_config.get('extra') == 'allow' else None + _object_setattr(m, '__dict__', fields_values) + _object_setattr(m, '__pydantic_fields_set__', _fields_set) + if not cls.__pydantic_root_model__: + _object_setattr(m, '__pydantic_extra__', _extra) + + if cls.__pydantic_post_init__: + m.model_post_init(None) + # update private attributes with values set + if hasattr(m, '__pydantic_private__') and m.__pydantic_private__ is not None: + for k, v in values.items(): + if k in m.__private_attributes__: + m.__pydantic_private__[k] = v + + elif not cls.__pydantic_root_model__: + # Note: if there are any private attributes, cls.__pydantic_post_init__ would exist + # Since it doesn't, that means that `__pydantic_private__` should be set to None + _object_setattr(m, '__pydantic_private__', None) + + return m + + def model_copy(self, *, update: Mapping[str, Any] | None = None, deep: bool = False) -> Self: + """!!! abstract "Usage Documentation" + [`model_copy`](../concepts/serialization.md#model_copy) + + Returns a copy of the model. + + !!! note + The underlying instance's [`__dict__`][object.__dict__] attribute is copied. This + might have unexpected side effects if you store anything in it, on top of the model + fields (e.g. the value of [cached properties][functools.cached_property]). + + Args: + update: Values to change/add in the new model. Note: the data is not validated + before creating the new model. You should trust this data. + deep: Set to `True` to make a deep copy of the model. + + Returns: + New model instance. + """ + copied = self.__deepcopy__() if deep else self.__copy__() + if update: + if self.model_config.get('extra') == 'allow': + for k, v in update.items(): + if k in self.__pydantic_fields__: + copied.__dict__[k] = v + else: + if copied.__pydantic_extra__ is None: + copied.__pydantic_extra__ = {} + copied.__pydantic_extra__[k] = v + else: + copied.__dict__.update(update) + copied.__pydantic_fields_set__.update(update.keys()) + return copied + + def model_dump( + self, + *, + mode: Literal['json', 'python'] | str = 'python', + include: IncEx | None = None, + exclude: IncEx | None = None, + context: Any | None = None, + by_alias: bool | None = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + round_trip: bool = False, + warnings: bool | Literal['none', 'warn', 'error'] = True, + fallback: Callable[[Any], Any] | None = None, + serialize_as_any: bool = False, + ) -> dict[str, Any]: + """!!! abstract "Usage Documentation" + [`model_dump`](../concepts/serialization.md#modelmodel_dump) + + Generate a dictionary representation of the model, optionally specifying which fields to include or exclude. + + Args: + mode: The mode in which `to_python` should run. + If mode is 'json', the output will only contain JSON serializable types. + If mode is 'python', the output may contain non-JSON-serializable Python objects. + include: A set of fields to include in the output. + exclude: A set of fields to exclude from the output. + context: Additional context to pass to the serializer. + by_alias: Whether to use the field's alias in the dictionary key if defined. + exclude_unset: Whether to exclude fields that have not been explicitly set. + exclude_defaults: Whether to exclude fields that are set to their default value. + exclude_none: Whether to exclude fields that have a value of `None`. + round_trip: If True, dumped values should be valid as input for non-idempotent types such as Json[T]. + warnings: How to handle serialization errors. False/"none" ignores them, True/"warn" logs errors, + "error" raises a [`PydanticSerializationError`][pydantic_core.PydanticSerializationError]. + fallback: A function to call when an unknown value is encountered. If not provided, + a [`PydanticSerializationError`][pydantic_core.PydanticSerializationError] error is raised. + serialize_as_any: Whether to serialize fields with duck-typing serialization behavior. + + Returns: + A dictionary representation of the model. + """ + return self.__pydantic_serializer__.to_python( + self, + mode=mode, + by_alias=by_alias, + include=include, + exclude=exclude, + context=context, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + round_trip=round_trip, + warnings=warnings, + fallback=fallback, + serialize_as_any=serialize_as_any, + ) + + def model_dump_json( + self, + *, + indent: int | None = None, + include: IncEx | None = None, + exclude: IncEx | None = None, + context: Any | None = None, + by_alias: bool | None = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + round_trip: bool = False, + warnings: bool | Literal['none', 'warn', 'error'] = True, + fallback: Callable[[Any], Any] | None = None, + serialize_as_any: bool = False, + ) -> str: + """!!! abstract "Usage Documentation" + [`model_dump_json`](../concepts/serialization.md#modelmodel_dump_json) + + Generates a JSON representation of the model using Pydantic's `to_json` method. + + Args: + indent: Indentation to use in the JSON output. If None is passed, the output will be compact. + include: Field(s) to include in the JSON output. + exclude: Field(s) to exclude from the JSON output. + context: Additional context to pass to the serializer. + by_alias: Whether to serialize using field aliases. + exclude_unset: Whether to exclude fields that have not been explicitly set. + exclude_defaults: Whether to exclude fields that are set to their default value. + exclude_none: Whether to exclude fields that have a value of `None`. + round_trip: If True, dumped values should be valid as input for non-idempotent types such as Json[T]. + warnings: How to handle serialization errors. False/"none" ignores them, True/"warn" logs errors, + "error" raises a [`PydanticSerializationError`][pydantic_core.PydanticSerializationError]. + fallback: A function to call when an unknown value is encountered. If not provided, + a [`PydanticSerializationError`][pydantic_core.PydanticSerializationError] error is raised. + serialize_as_any: Whether to serialize fields with duck-typing serialization behavior. + + Returns: + A JSON string representation of the model. + """ + return self.__pydantic_serializer__.to_json( + self, + indent=indent, + include=include, + exclude=exclude, + context=context, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + round_trip=round_trip, + warnings=warnings, + fallback=fallback, + serialize_as_any=serialize_as_any, + ).decode() + + @classmethod + def model_json_schema( + cls, + by_alias: bool = True, + ref_template: str = DEFAULT_REF_TEMPLATE, + schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, + mode: JsonSchemaMode = 'validation', + ) -> dict[str, Any]: + """Generates a JSON schema for a model class. + + Args: + by_alias: Whether to use attribute aliases or not. + ref_template: The reference template. + schema_generator: To override the logic used to generate the JSON schema, as a subclass of + `GenerateJsonSchema` with your desired modifications + mode: The mode in which to generate the schema. + + Returns: + The JSON schema for the given model class. + """ + return model_json_schema( + cls, by_alias=by_alias, ref_template=ref_template, schema_generator=schema_generator, mode=mode + ) + + @classmethod + def model_parametrized_name(cls, params: tuple[type[Any], ...]) -> str: + """Compute the class name for parametrizations of generic classes. + + This method can be overridden to achieve a custom naming scheme for generic BaseModels. + + Args: + params: Tuple of types of the class. Given a generic class + `Model` with 2 type variables and a concrete model `Model[str, int]`, + the value `(str, int)` would be passed to `params`. + + Returns: + String representing the new class where `params` are passed to `cls` as type variables. + + Raises: + TypeError: Raised when trying to generate concrete names for non-generic models. + """ + if not issubclass(cls, typing.Generic): + raise TypeError('Concrete names should only be generated for generic models.') + + # Any strings received should represent forward references, so we handle them specially below. + # If we eventually move toward wrapping them in a ForwardRef in __class_getitem__ in the future, + # we may be able to remove this special case. + param_names = [param if isinstance(param, str) else _repr.display_as_type(param) for param in params] + params_component = ', '.join(param_names) + return f'{cls.__name__}[{params_component}]' + + def model_post_init(self, context: Any, /) -> None: + """Override this method to perform additional initialization after `__init__` and `model_construct`. + This is useful if you want to do some validation that requires the entire model to be initialized. + """ + pass + + @classmethod + def model_rebuild( + cls, + *, + force: bool = False, + raise_errors: bool = True, + _parent_namespace_depth: int = 2, + _types_namespace: MappingNamespace | None = None, + ) -> bool | None: + """Try to rebuild the pydantic-core schema for the model. + + This may be necessary when one of the annotations is a ForwardRef which could not be resolved during + the initial attempt to build the schema, and automatic rebuilding fails. + + Args: + force: Whether to force the rebuilding of the model schema, defaults to `False`. + raise_errors: Whether to raise errors, defaults to `True`. + _parent_namespace_depth: The depth level of the parent namespace, defaults to 2. + _types_namespace: The types namespace, defaults to `None`. + + Returns: + Returns `None` if the schema is already "complete" and rebuilding was not required. + If rebuilding _was_ required, returns `True` if rebuilding was successful, otherwise `False`. + """ + if not force and cls.__pydantic_complete__: + return None + + for attr in ('__pydantic_core_schema__', '__pydantic_validator__', '__pydantic_serializer__'): + if attr in cls.__dict__ and not isinstance(getattr(cls, attr), _mock_val_ser.MockValSer): + # Deleting the validator/serializer is necessary as otherwise they can get reused in + # pydantic-core. We do so only if they aren't mock instances, otherwise — as `model_rebuild()` + # isn't thread-safe — concurrent model instantiations can lead to the parent validator being used. + # Same applies for the core schema that can be reused in schema generation. + delattr(cls, attr) + + cls.__pydantic_complete__ = False + + if _types_namespace is not None: + rebuild_ns = _types_namespace + elif _parent_namespace_depth > 0: + rebuild_ns = _typing_extra.parent_frame_namespace(parent_depth=_parent_namespace_depth, force=True) or {} + else: + rebuild_ns = {} + + parent_ns = _model_construction.unpack_lenient_weakvaluedict(cls.__pydantic_parent_namespace__) or {} + + ns_resolver = _namespace_utils.NsResolver( + parent_namespace={**rebuild_ns, **parent_ns}, + ) + + if not cls.__pydantic_fields_complete__: + typevars_map = _generics.get_model_typevars_map(cls) + try: + cls.__pydantic_fields__ = _fields.rebuild_model_fields( + cls, + ns_resolver=ns_resolver, + typevars_map=typevars_map, + ) + except NameError as e: + exc = PydanticUndefinedAnnotation.from_name_error(e) + _mock_val_ser.set_model_mocks(cls, f'`{exc.name}`') + if raise_errors: + raise exc from e + + if not raise_errors and not cls.__pydantic_fields_complete__: + # No need to continue with schema gen, it is guaranteed to fail + return False + + assert cls.__pydantic_fields_complete__ + + return _model_construction.complete_model_class( + cls, + _config.ConfigWrapper(cls.model_config, check=False), + raise_errors=raise_errors, + ns_resolver=ns_resolver, + ) + + @classmethod + def model_validate( + cls, + obj: Any, + *, + strict: bool | None = None, + from_attributes: bool | None = None, + context: Any | None = None, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> Self: + """Validate a pydantic model instance. + + Args: + obj: The object to validate. + strict: Whether to enforce types strictly. + from_attributes: Whether to extract data from object attributes. + context: Additional context to pass to the validator. + by_alias: Whether to use the field's alias when validating against the provided input data. + by_name: Whether to use the field's name when validating against the provided input data. + + Raises: + ValidationError: If the object could not be validated. + + Returns: + The validated model instance. + """ + # `__tracebackhide__` tells pytest and some other tools to omit this function from tracebacks + __tracebackhide__ = True + + if by_alias is False and by_name is not True: + raise PydanticUserError( + 'At least one of `by_alias` or `by_name` must be set to True.', + code='validate-by-alias-and-name-false', + ) + + return cls.__pydantic_validator__.validate_python( + obj, strict=strict, from_attributes=from_attributes, context=context, by_alias=by_alias, by_name=by_name + ) + + @classmethod + def model_validate_json( + cls, + json_data: str | bytes | bytearray, + *, + strict: bool | None = None, + context: Any | None = None, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> Self: + """!!! abstract "Usage Documentation" + [JSON Parsing](../concepts/json.md#json-parsing) + + Validate the given JSON data against the Pydantic model. + + Args: + json_data: The JSON data to validate. + strict: Whether to enforce types strictly. + context: Extra variables to pass to the validator. + by_alias: Whether to use the field's alias when validating against the provided input data. + by_name: Whether to use the field's name when validating against the provided input data. + + Returns: + The validated Pydantic model. + + Raises: + ValidationError: If `json_data` is not a JSON string or the object could not be validated. + """ + # `__tracebackhide__` tells pytest and some other tools to omit this function from tracebacks + __tracebackhide__ = True + + if by_alias is False and by_name is not True: + raise PydanticUserError( + 'At least one of `by_alias` or `by_name` must be set to True.', + code='validate-by-alias-and-name-false', + ) + + return cls.__pydantic_validator__.validate_json( + json_data, strict=strict, context=context, by_alias=by_alias, by_name=by_name + ) + + @classmethod + def model_validate_strings( + cls, + obj: Any, + *, + strict: bool | None = None, + context: Any | None = None, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> Self: + """Validate the given object with string data against the Pydantic model. + + Args: + obj: The object containing string data to validate. + strict: Whether to enforce types strictly. + context: Extra variables to pass to the validator. + by_alias: Whether to use the field's alias when validating against the provided input data. + by_name: Whether to use the field's name when validating against the provided input data. + + Returns: + The validated Pydantic model. + """ + # `__tracebackhide__` tells pytest and some other tools to omit this function from tracebacks + __tracebackhide__ = True + + if by_alias is False and by_name is not True: + raise PydanticUserError( + 'At least one of `by_alias` or `by_name` must be set to True.', + code='validate-by-alias-and-name-false', + ) + + return cls.__pydantic_validator__.validate_strings( + obj, strict=strict, context=context, by_alias=by_alias, by_name=by_name + ) + + @classmethod + def __get_pydantic_core_schema__(cls, source: type[BaseModel], handler: GetCoreSchemaHandler, /) -> CoreSchema: + # This warning is only emitted when calling `super().__get_pydantic_core_schema__` from a model subclass. + # In the generate schema logic, this method (`BaseModel.__get_pydantic_core_schema__`) is special cased to + # *not* be called if not overridden. + warnings.warn( + 'The `__get_pydantic_core_schema__` method of the `BaseModel` class is deprecated. If you are calling ' + '`super().__get_pydantic_core_schema__` when overriding the method on a Pydantic model, consider using ' + '`handler(source)` instead. However, note that overriding this method on models can lead to unexpected ' + 'side effects.', + PydanticDeprecatedSince211, + stacklevel=2, + ) + # Logic copied over from `GenerateSchema._model_schema`: + schema = cls.__dict__.get('__pydantic_core_schema__') + if schema is not None and not isinstance(schema, _mock_val_ser.MockCoreSchema): + return cls.__pydantic_core_schema__ + + return handler(source) + + @classmethod + def __get_pydantic_json_schema__( + cls, + core_schema: CoreSchema, + handler: GetJsonSchemaHandler, + /, + ) -> JsonSchemaValue: + """Hook into generating the model's JSON schema. + + Args: + core_schema: A `pydantic-core` CoreSchema. + You can ignore this argument and call the handler with a new CoreSchema, + wrap this CoreSchema (`{'type': 'nullable', 'schema': current_schema}`), + or just call the handler with the original schema. + handler: Call into Pydantic's internal JSON schema generation. + This will raise a `pydantic.errors.PydanticInvalidForJsonSchema` if JSON schema + generation fails. + Since this gets called by `BaseModel.model_json_schema` you can override the + `schema_generator` argument to that function to change JSON schema generation globally + for a type. + + Returns: + A JSON schema, as a Python object. + """ + return handler(core_schema) + + @classmethod + def __pydantic_init_subclass__(cls, **kwargs: Any) -> None: + """This is intended to behave just like `__init_subclass__`, but is called by `ModelMetaclass` + only after the class is actually fully initialized. In particular, attributes like `model_fields` will + be present when this is called. + + This is necessary because `__init_subclass__` will always be called by `type.__new__`, + and it would require a prohibitively large refactor to the `ModelMetaclass` to ensure that + `type.__new__` was called in such a manner that the class would already be sufficiently initialized. + + This will receive the same `kwargs` that would be passed to the standard `__init_subclass__`, namely, + any kwargs passed to the class definition that aren't used internally by pydantic. + + Args: + **kwargs: Any keyword arguments passed to the class definition that aren't used internally + by pydantic. + """ + pass + + def __class_getitem__( + cls, typevar_values: type[Any] | tuple[type[Any], ...] + ) -> type[BaseModel] | _forward_ref.PydanticRecursiveRef: + cached = _generics.get_cached_generic_type_early(cls, typevar_values) + if cached is not None: + return cached + + if cls is BaseModel: + raise TypeError('Type parameters should be placed on typing.Generic, not BaseModel') + if not hasattr(cls, '__parameters__'): + raise TypeError(f'{cls} cannot be parametrized because it does not inherit from typing.Generic') + if not cls.__pydantic_generic_metadata__['parameters'] and typing.Generic not in cls.__bases__: + raise TypeError(f'{cls} is not a generic class') + + if not isinstance(typevar_values, tuple): + typevar_values = (typevar_values,) + + # For a model `class Model[T, U, V = int](BaseModel): ...` parametrized with `(str, bool)`, + # this gives us `{T: str, U: bool, V: int}`: + typevars_map = _generics.map_generic_model_arguments(cls, typevar_values) + # We also update the provided args to use defaults values (`(str, bool)` becomes `(str, bool, int)`): + typevar_values = tuple(v for v in typevars_map.values()) + + if _utils.all_identical(typevars_map.keys(), typevars_map.values()) and typevars_map: + submodel = cls # if arguments are equal to parameters it's the same object + _generics.set_cached_generic_type(cls, typevar_values, submodel) + else: + parent_args = cls.__pydantic_generic_metadata__['args'] + if not parent_args: + args = typevar_values + else: + args = tuple(_generics.replace_types(arg, typevars_map) for arg in parent_args) + + origin = cls.__pydantic_generic_metadata__['origin'] or cls + model_name = origin.model_parametrized_name(args) + params = tuple( + {param: None for param in _generics.iter_contained_typevars(typevars_map.values())} + ) # use dict as ordered set + + with _generics.generic_recursion_self_type(origin, args) as maybe_self_type: + cached = _generics.get_cached_generic_type_late(cls, typevar_values, origin, args) + if cached is not None: + return cached + + if maybe_self_type is not None: + return maybe_self_type + + # Attempt to rebuild the origin in case new types have been defined + try: + # depth 2 gets you above this __class_getitem__ call. + # Note that we explicitly provide the parent ns, otherwise + # `model_rebuild` will use the parent ns no matter if it is the ns of a module. + # We don't want this here, as this has unexpected effects when a model + # is being parametrized during a forward annotation evaluation. + parent_ns = _typing_extra.parent_frame_namespace(parent_depth=2) or {} + origin.model_rebuild(_types_namespace=parent_ns) + except PydanticUndefinedAnnotation: + # It's okay if it fails, it just means there are still undefined types + # that could be evaluated later. + pass + + submodel = _generics.create_generic_submodel(model_name, origin, args, params) + + _generics.set_cached_generic_type(cls, typevar_values, submodel, origin, args) + + return submodel + + def __copy__(self) -> Self: + """Returns a shallow copy of the model.""" + cls = type(self) + m = cls.__new__(cls) + _object_setattr(m, '__dict__', copy(self.__dict__)) + _object_setattr(m, '__pydantic_extra__', copy(self.__pydantic_extra__)) + _object_setattr(m, '__pydantic_fields_set__', copy(self.__pydantic_fields_set__)) + + if not hasattr(self, '__pydantic_private__') or self.__pydantic_private__ is None: + _object_setattr(m, '__pydantic_private__', None) + else: + _object_setattr( + m, + '__pydantic_private__', + {k: v for k, v in self.__pydantic_private__.items() if v is not PydanticUndefined}, + ) + + return m + + def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self: + """Returns a deep copy of the model.""" + cls = type(self) + m = cls.__new__(cls) + _object_setattr(m, '__dict__', deepcopy(self.__dict__, memo=memo)) + _object_setattr(m, '__pydantic_extra__', deepcopy(self.__pydantic_extra__, memo=memo)) + # This next line doesn't need a deepcopy because __pydantic_fields_set__ is a set[str], + # and attempting a deepcopy would be marginally slower. + _object_setattr(m, '__pydantic_fields_set__', copy(self.__pydantic_fields_set__)) + + if not hasattr(self, '__pydantic_private__') or self.__pydantic_private__ is None: + _object_setattr(m, '__pydantic_private__', None) + else: + _object_setattr( + m, + '__pydantic_private__', + deepcopy({k: v for k, v in self.__pydantic_private__.items() if v is not PydanticUndefined}, memo=memo), + ) + + return m + + if not TYPE_CHECKING: + # We put `__getattr__` in a non-TYPE_CHECKING block because otherwise, mypy allows arbitrary attribute access + # The same goes for __setattr__ and __delattr__, see: https://github.com/pydantic/pydantic/issues/8643 + + def __getattr__(self, item: str) -> Any: + private_attributes = object.__getattribute__(self, '__private_attributes__') + if item in private_attributes: + attribute = private_attributes[item] + if hasattr(attribute, '__get__'): + return attribute.__get__(self, type(self)) # type: ignore + + try: + # Note: self.__pydantic_private__ cannot be None if self.__private_attributes__ has items + return self.__pydantic_private__[item] # type: ignore + except KeyError as exc: + raise AttributeError(f'{type(self).__name__!r} object has no attribute {item!r}') from exc + else: + # `__pydantic_extra__` can fail to be set if the model is not yet fully initialized. + # See `BaseModel.__repr_args__` for more details + try: + pydantic_extra = object.__getattribute__(self, '__pydantic_extra__') + except AttributeError: + pydantic_extra = None + + if pydantic_extra: + try: + return pydantic_extra[item] + except KeyError as exc: + raise AttributeError(f'{type(self).__name__!r} object has no attribute {item!r}') from exc + else: + if hasattr(self.__class__, item): + return super().__getattribute__(item) # Raises AttributeError if appropriate + else: + # this is the current error + raise AttributeError(f'{type(self).__name__!r} object has no attribute {item!r}') + + def __setattr__(self, name: str, value: Any) -> None: + if (setattr_handler := self.__pydantic_setattr_handlers__.get(name)) is not None: + setattr_handler(self, name, value) + # if None is returned from _setattr_handler, the attribute was set directly + elif (setattr_handler := self._setattr_handler(name, value)) is not None: + setattr_handler(self, name, value) # call here to not memo on possibly unknown fields + self.__pydantic_setattr_handlers__[name] = setattr_handler # memoize the handler for faster access + + def _setattr_handler(self, name: str, value: Any) -> Callable[[BaseModel, str, Any], None] | None: + """Get a handler for setting an attribute on the model instance. + + Returns: + A handler for setting an attribute on the model instance. Used for memoization of the handler. + Memoizing the handlers leads to a dramatic performance improvement in `__setattr__` + Returns `None` when memoization is not safe, then the attribute is set directly. + """ + cls = self.__class__ + if name in cls.__class_vars__: + raise AttributeError( + f'{name!r} is a ClassVar of `{cls.__name__}` and cannot be set on an instance. ' + f'If you want to set a value on the class, use `{cls.__name__}.{name} = value`.' + ) + elif not _fields.is_valid_field_name(name): + if (attribute := cls.__private_attributes__.get(name)) is not None: + if hasattr(attribute, '__set__'): + return lambda model, _name, val: attribute.__set__(model, val) + else: + return _SIMPLE_SETATTR_HANDLERS['private'] + else: + _object_setattr(self, name, value) + return None # Can not return memoized handler with possibly freeform attr names + + attr = getattr(cls, name, None) + # NOTE: We currently special case properties and `cached_property`, but we might need + # to generalize this to all data/non-data descriptors at some point. For non-data descriptors + # (such as `cached_property`), it isn't obvious though. `cached_property` caches the value + # to the instance's `__dict__`, but other non-data descriptors might do things differently. + if isinstance(attr, cached_property): + return _SIMPLE_SETATTR_HANDLERS['cached_property'] + + _check_frozen(cls, name, value) + + # We allow properties to be set only on non frozen models for now (to match dataclasses). + # This can be changed if it ever gets requested. + if isinstance(attr, property): + return lambda model, _name, val: attr.__set__(model, val) + elif cls.model_config.get('validate_assignment'): + return _SIMPLE_SETATTR_HANDLERS['validate_assignment'] + elif name not in cls.__pydantic_fields__: + if cls.model_config.get('extra') != 'allow': + # TODO - matching error + raise ValueError(f'"{cls.__name__}" object has no field "{name}"') + elif attr is None: + # attribute does not exist, so put it in extra + self.__pydantic_extra__[name] = value + return None # Can not return memoized handler with possibly freeform attr names + else: + # attribute _does_ exist, and was not in extra, so update it + return _SIMPLE_SETATTR_HANDLERS['extra_known'] + else: + return _SIMPLE_SETATTR_HANDLERS['model_field'] + + def __delattr__(self, item: str) -> Any: + cls = self.__class__ + + if item in self.__private_attributes__: + attribute = self.__private_attributes__[item] + if hasattr(attribute, '__delete__'): + attribute.__delete__(self) # type: ignore + return + + try: + # Note: self.__pydantic_private__ cannot be None if self.__private_attributes__ has items + del self.__pydantic_private__[item] # type: ignore + return + except KeyError as exc: + raise AttributeError(f'{cls.__name__!r} object has no attribute {item!r}') from exc + + # Allow cached properties to be deleted (even if the class is frozen): + attr = getattr(cls, item, None) + if isinstance(attr, cached_property): + return object.__delattr__(self, item) + + _check_frozen(cls, name=item, value=None) + + if item in self.__pydantic_fields__: + object.__delattr__(self, item) + elif self.__pydantic_extra__ is not None and item in self.__pydantic_extra__: + del self.__pydantic_extra__[item] + else: + try: + object.__delattr__(self, item) + except AttributeError: + raise AttributeError(f'{type(self).__name__!r} object has no attribute {item!r}') + + # Because we make use of `@dataclass_transform()`, `__replace__` is already synthesized by + # type checkers, so we define the implementation in this `if not TYPE_CHECKING:` block: + def __replace__(self, **changes: Any) -> Self: + return self.model_copy(update=changes) + + def __getstate__(self) -> dict[Any, Any]: + private = self.__pydantic_private__ + if private: + private = {k: v for k, v in private.items() if v is not PydanticUndefined} + return { + '__dict__': self.__dict__, + '__pydantic_extra__': self.__pydantic_extra__, + '__pydantic_fields_set__': self.__pydantic_fields_set__, + '__pydantic_private__': private, + } + + def __setstate__(self, state: dict[Any, Any]) -> None: + _object_setattr(self, '__pydantic_fields_set__', state.get('__pydantic_fields_set__', {})) + _object_setattr(self, '__pydantic_extra__', state.get('__pydantic_extra__', {})) + _object_setattr(self, '__pydantic_private__', state.get('__pydantic_private__', {})) + _object_setattr(self, '__dict__', state.get('__dict__', {})) + + if not TYPE_CHECKING: + + def __eq__(self, other: Any) -> bool: + if isinstance(other, BaseModel): + # When comparing instances of generic types for equality, as long as all field values are equal, + # only require their generic origin types to be equal, rather than exact type equality. + # This prevents headaches like MyGeneric(x=1) != MyGeneric[Any](x=1). + self_type = self.__pydantic_generic_metadata__['origin'] or self.__class__ + other_type = other.__pydantic_generic_metadata__['origin'] or other.__class__ + + # Perform common checks first + if not ( + self_type == other_type + and getattr(self, '__pydantic_private__', None) == getattr(other, '__pydantic_private__', None) + and self.__pydantic_extra__ == other.__pydantic_extra__ + ): + return False + + # We only want to compare pydantic fields but ignoring fields is costly. + # We'll perform a fast check first, and fallback only when needed + # See GH-7444 and GH-7825 for rationale and a performance benchmark + + # First, do the fast (and sometimes faulty) __dict__ comparison + if self.__dict__ == other.__dict__: + # If the check above passes, then pydantic fields are equal, we can return early + return True + + # We don't want to trigger unnecessary costly filtering of __dict__ on all unequal objects, so we return + # early if there are no keys to ignore (we would just return False later on anyway) + model_fields = type(self).__pydantic_fields__.keys() + if self.__dict__.keys() <= model_fields and other.__dict__.keys() <= model_fields: + return False + + # If we reach here, there are non-pydantic-fields keys, mapped to unequal values, that we need to ignore + # Resort to costly filtering of the __dict__ objects + # We use operator.itemgetter because it is much faster than dict comprehensions + # NOTE: Contrary to standard python class and instances, when the Model class has a default value for an + # attribute and the model instance doesn't have a corresponding attribute, accessing the missing attribute + # raises an error in BaseModel.__getattr__ instead of returning the class attribute + # So we can use operator.itemgetter() instead of operator.attrgetter() + getter = operator.itemgetter(*model_fields) if model_fields else lambda _: _utils._SENTINEL + try: + return getter(self.__dict__) == getter(other.__dict__) + except KeyError: + # In rare cases (such as when using the deprecated BaseModel.copy() method), + # the __dict__ may not contain all model fields, which is how we can get here. + # getter(self.__dict__) is much faster than any 'safe' method that accounts + # for missing keys, and wrapping it in a `try` doesn't slow things down much + # in the common case. + self_fields_proxy = _utils.SafeGetItemProxy(self.__dict__) + other_fields_proxy = _utils.SafeGetItemProxy(other.__dict__) + return getter(self_fields_proxy) == getter(other_fields_proxy) + + # other instance is not a BaseModel + else: + return NotImplemented # delegate to the other item in the comparison + + if TYPE_CHECKING: + # We put `__init_subclass__` in a TYPE_CHECKING block because, even though we want the type-checking benefits + # described in the signature of `__init_subclass__` below, we don't want to modify the default behavior of + # subclass initialization. + + def __init_subclass__(cls, **kwargs: Unpack[ConfigDict]): + """This signature is included purely to help type-checkers check arguments to class declaration, which + provides a way to conveniently set model_config key/value pairs. + + ```python + from pydantic import BaseModel + + class MyModel(BaseModel, extra='allow'): ... + ``` + + However, this may be deceiving, since the _actual_ calls to `__init_subclass__` will not receive any + of the config arguments, and will only receive any keyword arguments passed during class initialization + that are _not_ expected keys in ConfigDict. (This is due to the way `ModelMetaclass.__new__` works.) + + Args: + **kwargs: Keyword arguments passed to the class definition, which set model_config + + Note: + You may want to override `__pydantic_init_subclass__` instead, which behaves similarly but is called + *after* the class is fully initialized. + """ + + def __iter__(self) -> TupleGenerator: + """So `dict(model)` works.""" + yield from [(k, v) for (k, v) in self.__dict__.items() if not k.startswith('_')] + extra = self.__pydantic_extra__ + if extra: + yield from extra.items() + + def __repr__(self) -> str: + return f'{self.__repr_name__()}({self.__repr_str__(", ")})' + + def __repr_args__(self) -> _repr.ReprArgs: + # Eagerly create the repr of computed fields, as this may trigger access of cached properties and as such + # modify the instance's `__dict__`. If we don't do it now, it could happen when iterating over the `__dict__` + # below if the instance happens to be referenced in a field, and would modify the `__dict__` size *during* iteration. + computed_fields_repr_args = [ + (k, getattr(self, k)) for k, v in self.__pydantic_computed_fields__.items() if v.repr + ] + + for k, v in self.__dict__.items(): + field = self.__pydantic_fields__.get(k) + if field and field.repr: + if v is not self: + yield k, v + else: + yield k, self.__repr_recursion__(v) + # `__pydantic_extra__` can fail to be set if the model is not yet fully initialized. + # This can happen if a `ValidationError` is raised during initialization and the instance's + # repr is generated as part of the exception handling. Therefore, we use `getattr` here + # with a fallback, even though the type hints indicate the attribute will always be present. + try: + pydantic_extra = object.__getattribute__(self, '__pydantic_extra__') + except AttributeError: + pydantic_extra = None + + if pydantic_extra is not None: + yield from ((k, v) for k, v in pydantic_extra.items()) + yield from computed_fields_repr_args + + # take logic from `_repr.Representation` without the side effects of inheritance, see #5740 + __repr_name__ = _repr.Representation.__repr_name__ + __repr_recursion__ = _repr.Representation.__repr_recursion__ + __repr_str__ = _repr.Representation.__repr_str__ + __pretty__ = _repr.Representation.__pretty__ + __rich_repr__ = _repr.Representation.__rich_repr__ + + def __str__(self) -> str: + return self.__repr_str__(' ') + + # ##### Deprecated methods from v1 ##### + @property + @typing_extensions.deprecated( + 'The `__fields__` attribute is deprecated, use `model_fields` instead.', category=None + ) + def __fields__(self) -> dict[str, FieldInfo]: + warnings.warn( + 'The `__fields__` attribute is deprecated, use `model_fields` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + return getattr(type(self), '__pydantic_fields__', {}) + + @property + @typing_extensions.deprecated( + 'The `__fields_set__` attribute is deprecated, use `model_fields_set` instead.', + category=None, + ) + def __fields_set__(self) -> set[str]: + warnings.warn( + 'The `__fields_set__` attribute is deprecated, use `model_fields_set` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + return self.__pydantic_fields_set__ + + @typing_extensions.deprecated('The `dict` method is deprecated; use `model_dump` instead.', category=None) + def dict( # noqa: D102 + self, + *, + include: IncEx | None = None, + exclude: IncEx | None = None, + by_alias: bool = False, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + ) -> Dict[str, Any]: # noqa UP006 + warnings.warn( + 'The `dict` method is deprecated; use `model_dump` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + return self.model_dump( + include=include, + exclude=exclude, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + ) + + @typing_extensions.deprecated('The `json` method is deprecated; use `model_dump_json` instead.', category=None) + def json( # noqa: D102 + self, + *, + include: IncEx | None = None, + exclude: IncEx | None = None, + by_alias: bool = False, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + encoder: Callable[[Any], Any] | None = PydanticUndefined, # type: ignore[assignment] + models_as_dict: bool = PydanticUndefined, # type: ignore[assignment] + **dumps_kwargs: Any, + ) -> str: + warnings.warn( + 'The `json` method is deprecated; use `model_dump_json` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + if encoder is not PydanticUndefined: + raise TypeError('The `encoder` argument is no longer supported; use field serializers instead.') + if models_as_dict is not PydanticUndefined: + raise TypeError('The `models_as_dict` argument is no longer supported; use a model serializer instead.') + if dumps_kwargs: + raise TypeError('`dumps_kwargs` keyword arguments are no longer supported.') + return self.model_dump_json( + include=include, + exclude=exclude, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + ) + + @classmethod + @typing_extensions.deprecated('The `parse_obj` method is deprecated; use `model_validate` instead.', category=None) + def parse_obj(cls, obj: Any) -> Self: # noqa: D102 + warnings.warn( + 'The `parse_obj` method is deprecated; use `model_validate` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + return cls.model_validate(obj) + + @classmethod + @typing_extensions.deprecated( + 'The `parse_raw` method is deprecated; if your data is JSON use `model_validate_json`, ' + 'otherwise load the data then use `model_validate` instead.', + category=None, + ) + def parse_raw( # noqa: D102 + cls, + b: str | bytes, + *, + content_type: str | None = None, + encoding: str = 'utf8', + proto: DeprecatedParseProtocol | None = None, + allow_pickle: bool = False, + ) -> Self: # pragma: no cover + warnings.warn( + 'The `parse_raw` method is deprecated; if your data is JSON use `model_validate_json`, ' + 'otherwise load the data then use `model_validate` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + from .deprecated import parse + + try: + obj = parse.load_str_bytes( + b, + proto=proto, + content_type=content_type, + encoding=encoding, + allow_pickle=allow_pickle, + ) + except (ValueError, TypeError) as exc: + import json + + # try to match V1 + if isinstance(exc, UnicodeDecodeError): + type_str = 'value_error.unicodedecode' + elif isinstance(exc, json.JSONDecodeError): + type_str = 'value_error.jsondecode' + elif isinstance(exc, ValueError): + type_str = 'value_error' + else: + type_str = 'type_error' + + # ctx is missing here, but since we've added `input` to the error, we're not pretending it's the same + error: pydantic_core.InitErrorDetails = { + # The type: ignore on the next line is to ignore the requirement of LiteralString + 'type': pydantic_core.PydanticCustomError(type_str, str(exc)), # type: ignore + 'loc': ('__root__',), + 'input': b, + } + raise pydantic_core.ValidationError.from_exception_data(cls.__name__, [error]) + return cls.model_validate(obj) + + @classmethod + @typing_extensions.deprecated( + 'The `parse_file` method is deprecated; load the data from file, then if your data is JSON ' + 'use `model_validate_json`, otherwise `model_validate` instead.', + category=None, + ) + def parse_file( # noqa: D102 + cls, + path: str | Path, + *, + content_type: str | None = None, + encoding: str = 'utf8', + proto: DeprecatedParseProtocol | None = None, + allow_pickle: bool = False, + ) -> Self: + warnings.warn( + 'The `parse_file` method is deprecated; load the data from file, then if your data is JSON ' + 'use `model_validate_json`, otherwise `model_validate` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + from .deprecated import parse + + obj = parse.load_file( + path, + proto=proto, + content_type=content_type, + encoding=encoding, + allow_pickle=allow_pickle, + ) + return cls.parse_obj(obj) + + @classmethod + @typing_extensions.deprecated( + 'The `from_orm` method is deprecated; set ' + "`model_config['from_attributes']=True` and use `model_validate` instead.", + category=None, + ) + def from_orm(cls, obj: Any) -> Self: # noqa: D102 + warnings.warn( + 'The `from_orm` method is deprecated; set ' + "`model_config['from_attributes']=True` and use `model_validate` instead.", + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + if not cls.model_config.get('from_attributes', None): + raise PydanticUserError( + 'You must set the config attribute `from_attributes=True` to use from_orm', code=None + ) + return cls.model_validate(obj) + + @classmethod + @typing_extensions.deprecated('The `construct` method is deprecated; use `model_construct` instead.', category=None) + def construct(cls, _fields_set: set[str] | None = None, **values: Any) -> Self: # noqa: D102 + warnings.warn( + 'The `construct` method is deprecated; use `model_construct` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + return cls.model_construct(_fields_set=_fields_set, **values) + + @typing_extensions.deprecated( + 'The `copy` method is deprecated; use `model_copy` instead. ' + 'See the docstring of `BaseModel.copy` for details about how to handle `include` and `exclude`.', + category=None, + ) + def copy( + self, + *, + include: AbstractSetIntStr | MappingIntStrAny | None = None, + exclude: AbstractSetIntStr | MappingIntStrAny | None = None, + update: Dict[str, Any] | None = None, # noqa UP006 + deep: bool = False, + ) -> Self: # pragma: no cover + """Returns a copy of the model. + + !!! warning "Deprecated" + This method is now deprecated; use `model_copy` instead. + + If you need `include` or `exclude`, use: + + ```python {test="skip" lint="skip"} + data = self.model_dump(include=include, exclude=exclude, round_trip=True) + data = {**data, **(update or {})} + copied = self.model_validate(data) + ``` + + Args: + include: Optional set or mapping specifying which fields to include in the copied model. + exclude: Optional set or mapping specifying which fields to exclude in the copied model. + update: Optional dictionary of field-value pairs to override field values in the copied model. + deep: If True, the values of fields that are Pydantic models will be deep-copied. + + Returns: + A copy of the model with included, excluded and updated fields as specified. + """ + warnings.warn( + 'The `copy` method is deprecated; use `model_copy` instead. ' + 'See the docstring of `BaseModel.copy` for details about how to handle `include` and `exclude`.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + from .deprecated import copy_internals + + values = dict( + copy_internals._iter( + self, to_dict=False, by_alias=False, include=include, exclude=exclude, exclude_unset=False + ), + **(update or {}), + ) + if self.__pydantic_private__ is None: + private = None + else: + private = {k: v for k, v in self.__pydantic_private__.items() if v is not PydanticUndefined} + + if self.__pydantic_extra__ is None: + extra: dict[str, Any] | None = None + else: + extra = self.__pydantic_extra__.copy() + for k in list(self.__pydantic_extra__): + if k not in values: # k was in the exclude + extra.pop(k) + for k in list(values): + if k in self.__pydantic_extra__: # k must have come from extra + extra[k] = values.pop(k) + + # new `__pydantic_fields_set__` can have unset optional fields with a set value in `update` kwarg + if update: + fields_set = self.__pydantic_fields_set__ | update.keys() + else: + fields_set = set(self.__pydantic_fields_set__) + + # removing excluded fields from `__pydantic_fields_set__` + if exclude: + fields_set -= set(exclude) + + return copy_internals._copy_and_set_values(self, values, fields_set, extra, private, deep=deep) + + @classmethod + @typing_extensions.deprecated('The `schema` method is deprecated; use `model_json_schema` instead.', category=None) + def schema( # noqa: D102 + cls, by_alias: bool = True, ref_template: str = DEFAULT_REF_TEMPLATE + ) -> Dict[str, Any]: # noqa UP006 + warnings.warn( + 'The `schema` method is deprecated; use `model_json_schema` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + return cls.model_json_schema(by_alias=by_alias, ref_template=ref_template) + + @classmethod + @typing_extensions.deprecated( + 'The `schema_json` method is deprecated; use `model_json_schema` and json.dumps instead.', + category=None, + ) + def schema_json( # noqa: D102 + cls, *, by_alias: bool = True, ref_template: str = DEFAULT_REF_TEMPLATE, **dumps_kwargs: Any + ) -> str: # pragma: no cover + warnings.warn( + 'The `schema_json` method is deprecated; use `model_json_schema` and json.dumps instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + import json + + from .deprecated.json import pydantic_encoder + + return json.dumps( + cls.model_json_schema(by_alias=by_alias, ref_template=ref_template), + default=pydantic_encoder, + **dumps_kwargs, + ) + + @classmethod + @typing_extensions.deprecated('The `validate` method is deprecated; use `model_validate` instead.', category=None) + def validate(cls, value: Any) -> Self: # noqa: D102 + warnings.warn( + 'The `validate` method is deprecated; use `model_validate` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + return cls.model_validate(value) + + @classmethod + @typing_extensions.deprecated( + 'The `update_forward_refs` method is deprecated; use `model_rebuild` instead.', + category=None, + ) + def update_forward_refs(cls, **localns: Any) -> None: # noqa: D102 + warnings.warn( + 'The `update_forward_refs` method is deprecated; use `model_rebuild` instead.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + if localns: # pragma: no cover + raise TypeError('`localns` arguments are not longer accepted.') + cls.model_rebuild(force=True) + + @typing_extensions.deprecated( + 'The private method `_iter` will be removed and should no longer be used.', category=None + ) + def _iter(self, *args: Any, **kwargs: Any) -> Any: + warnings.warn( + 'The private method `_iter` will be removed and should no longer be used.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + from .deprecated import copy_internals + + return copy_internals._iter(self, *args, **kwargs) + + @typing_extensions.deprecated( + 'The private method `_copy_and_set_values` will be removed and should no longer be used.', + category=None, + ) + def _copy_and_set_values(self, *args: Any, **kwargs: Any) -> Any: + warnings.warn( + 'The private method `_copy_and_set_values` will be removed and should no longer be used.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + from .deprecated import copy_internals + + return copy_internals._copy_and_set_values(self, *args, **kwargs) + + @classmethod + @typing_extensions.deprecated( + 'The private method `_get_value` will be removed and should no longer be used.', + category=None, + ) + def _get_value(cls, *args: Any, **kwargs: Any) -> Any: + warnings.warn( + 'The private method `_get_value` will be removed and should no longer be used.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + from .deprecated import copy_internals + + return copy_internals._get_value(cls, *args, **kwargs) + + @typing_extensions.deprecated( + 'The private method `_calculate_keys` will be removed and should no longer be used.', + category=None, + ) + def _calculate_keys(self, *args: Any, **kwargs: Any) -> Any: + warnings.warn( + 'The private method `_calculate_keys` will be removed and should no longer be used.', + category=PydanticDeprecatedSince20, + stacklevel=2, + ) + from .deprecated import copy_internals + + return copy_internals._calculate_keys(self, *args, **kwargs) + + +ModelT = TypeVar('ModelT', bound=BaseModel) + + +@overload +def create_model( + model_name: str, + /, + *, + __config__: ConfigDict | None = None, + __doc__: str | None = None, + __base__: None = None, + __module__: str = __name__, + __validators__: dict[str, Callable[..., Any]] | None = None, + __cls_kwargs__: dict[str, Any] | None = None, + **field_definitions: Any | tuple[str, Any], +) -> type[BaseModel]: ... + + +@overload +def create_model( + model_name: str, + /, + *, + __config__: ConfigDict | None = None, + __doc__: str | None = None, + __base__: type[ModelT] | tuple[type[ModelT], ...], + __module__: str = __name__, + __validators__: dict[str, Callable[..., Any]] | None = None, + __cls_kwargs__: dict[str, Any] | None = None, + **field_definitions: Any | tuple[str, Any], +) -> type[ModelT]: ... + + +def create_model( # noqa: C901 + model_name: str, + /, + *, + __config__: ConfigDict | None = None, + __doc__: str | None = None, + __base__: type[ModelT] | tuple[type[ModelT], ...] | None = None, + __module__: str | None = None, + __validators__: dict[str, Callable[..., Any]] | None = None, + __cls_kwargs__: dict[str, Any] | None = None, + # TODO PEP 747: replace `Any` by the TypeForm: + **field_definitions: Any | tuple[str, Any], +) -> type[ModelT]: + """!!! abstract "Usage Documentation" + [Dynamic Model Creation](../concepts/models.md#dynamic-model-creation) + + Dynamically creates and returns a new Pydantic model, in other words, `create_model` dynamically creates a + subclass of [`BaseModel`][pydantic.BaseModel]. + + Args: + model_name: The name of the newly created model. + __config__: The configuration of the new model. + __doc__: The docstring of the new model. + __base__: The base class or classes for the new model. + __module__: The name of the module that the model belongs to; + if `None`, the value is taken from `sys._getframe(1)` + __validators__: A dictionary of methods that validate fields. The keys are the names of the validation methods to + be added to the model, and the values are the validation methods themselves. You can read more about functional + validators [here](https://docs.pydantic.dev/2.9/concepts/validators/#field-validators). + __cls_kwargs__: A dictionary of keyword arguments for class creation, such as `metaclass`. + **field_definitions: Field definitions of the new model. Either: + + - a single element, representing the type annotation of the field. + - a two-tuple, the first element being the type and the second element the assigned value + (either a default or the [`Field()`][pydantic.Field] function). + + Returns: + The new [model][pydantic.BaseModel]. + + Raises: + PydanticUserError: If `__base__` and `__config__` are both passed. + """ + if __base__ is None: + __base__ = (cast('type[ModelT]', BaseModel),) + elif not isinstance(__base__, tuple): + __base__ = (__base__,) + + __cls_kwargs__ = __cls_kwargs__ or {} + + fields: dict[str, Any] = {} + annotations: dict[str, Any] = {} + + for f_name, f_def in field_definitions.items(): + if isinstance(f_def, tuple): + if len(f_def) != 2: + raise PydanticUserError( + f'Field definition for {f_name!r} should a single element representing the type or a two-tuple, the first element ' + 'being the type and the second element the assigned value (either a default or the `Field()` function).', + code='create-model-field-definitions', + ) + + annotations[f_name] = f_def[0] + fields[f_name] = f_def[1] + else: + annotations[f_name] = f_def + + if __module__ is None: + f = sys._getframe(1) + __module__ = f.f_globals['__name__'] + + namespace: dict[str, Any] = {'__annotations__': annotations, '__module__': __module__} + if __doc__: + namespace.update({'__doc__': __doc__}) + if __validators__: + namespace.update(__validators__) + namespace.update(fields) + if __config__: + namespace['model_config'] = __config__ + resolved_bases = types.resolve_bases(__base__) + meta, ns, kwds = types.prepare_class(model_name, resolved_bases, kwds=__cls_kwargs__) + if resolved_bases is not __base__: + ns['__orig_bases__'] = __base__ + namespace.update(ns) + + return meta( + model_name, + resolved_bases, + namespace, + __pydantic_reset_parent_namespace__=False, + _create_model_module=__module__, + **kwds, + ) + + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/mypy.py b/lib/python3.12/site-packages/pydantic/mypy.py new file mode 100644 index 0000000000000000000000000000000000000000..1ac1b8761927362950940018a18791f407937f95 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/mypy.py @@ -0,0 +1,1380 @@ +"""This module includes classes and functions designed specifically for use with the mypy plugin.""" + +from __future__ import annotations + +import sys +from collections.abc import Iterator +from configparser import ConfigParser +from typing import Any, Callable + +from mypy.errorcodes import ErrorCode +from mypy.expandtype import expand_type, expand_type_by_instance +from mypy.nodes import ( + ARG_NAMED, + ARG_NAMED_OPT, + ARG_OPT, + ARG_POS, + ARG_STAR2, + INVARIANT, + MDEF, + Argument, + AssignmentStmt, + Block, + CallExpr, + ClassDef, + Context, + Decorator, + DictExpr, + EllipsisExpr, + Expression, + FuncDef, + IfStmt, + JsonDict, + MemberExpr, + NameExpr, + PassStmt, + PlaceholderNode, + RefExpr, + Statement, + StrExpr, + SymbolTableNode, + TempNode, + TypeAlias, + TypeInfo, + Var, +) +from mypy.options import Options +from mypy.plugin import ( + CheckerPluginInterface, + ClassDefContext, + MethodContext, + Plugin, + ReportConfigContext, + SemanticAnalyzerPluginInterface, +) +from mypy.plugins.common import ( + deserialize_and_fixup_type, +) +from mypy.semanal import set_callable_name +from mypy.server.trigger import make_wildcard_trigger +from mypy.state import state +from mypy.type_visitor import TypeTranslator +from mypy.typeops import map_type_from_supertype +from mypy.types import ( + AnyType, + CallableType, + Instance, + NoneType, + Type, + TypeOfAny, + TypeType, + TypeVarType, + UnionType, + get_proper_type, +) +from mypy.typevars import fill_typevars +from mypy.util import get_unique_redefinition_name +from mypy.version import __version__ as mypy_version + +from pydantic._internal import _fields +from pydantic.version import parse_mypy_version + +CONFIGFILE_KEY = 'pydantic-mypy' +METADATA_KEY = 'pydantic-mypy-metadata' +BASEMODEL_FULLNAME = 'pydantic.main.BaseModel' +BASESETTINGS_FULLNAME = 'pydantic_settings.main.BaseSettings' +ROOT_MODEL_FULLNAME = 'pydantic.root_model.RootModel' +MODEL_METACLASS_FULLNAME = 'pydantic._internal._model_construction.ModelMetaclass' +FIELD_FULLNAME = 'pydantic.fields.Field' +DATACLASS_FULLNAME = 'pydantic.dataclasses.dataclass' +MODEL_VALIDATOR_FULLNAME = 'pydantic.functional_validators.model_validator' +DECORATOR_FULLNAMES = { + 'pydantic.functional_validators.field_validator', + 'pydantic.functional_validators.model_validator', + 'pydantic.functional_serializers.serializer', + 'pydantic.functional_serializers.model_serializer', + 'pydantic.deprecated.class_validators.validator', + 'pydantic.deprecated.class_validators.root_validator', +} +IMPLICIT_CLASSMETHOD_DECORATOR_FULLNAMES = DECORATOR_FULLNAMES - {'pydantic.functional_serializers.model_serializer'} + + +MYPY_VERSION_TUPLE = parse_mypy_version(mypy_version) +BUILTINS_NAME = 'builtins' + +# Increment version if plugin changes and mypy caches should be invalidated +__version__ = 2 + + +def plugin(version: str) -> type[Plugin]: + """`version` is the mypy version string. + + We might want to use this to print a warning if the mypy version being used is + newer, or especially older, than we expect (or need). + + Args: + version: The mypy version string. + + Return: + The Pydantic mypy plugin type. + """ + return PydanticPlugin + + +class PydanticPlugin(Plugin): + """The Pydantic mypy plugin.""" + + def __init__(self, options: Options) -> None: + self.plugin_config = PydanticPluginConfig(options) + self._plugin_data = self.plugin_config.to_data() + super().__init__(options) + + def get_base_class_hook(self, fullname: str) -> Callable[[ClassDefContext], None] | None: + """Update Pydantic model class.""" + sym = self.lookup_fully_qualified(fullname) + if sym and isinstance(sym.node, TypeInfo): # pragma: no branch + # No branching may occur if the mypy cache has not been cleared + if sym.node.has_base(BASEMODEL_FULLNAME): + return self._pydantic_model_class_maker_callback + return None + + def get_metaclass_hook(self, fullname: str) -> Callable[[ClassDefContext], None] | None: + """Update Pydantic `ModelMetaclass` definition.""" + if fullname == MODEL_METACLASS_FULLNAME: + return self._pydantic_model_metaclass_marker_callback + return None + + def get_method_hook(self, fullname: str) -> Callable[[MethodContext], Type] | None: + """Adjust return type of `from_orm` method call.""" + if fullname.endswith('.from_orm'): + return from_attributes_callback + return None + + def report_config_data(self, ctx: ReportConfigContext) -> dict[str, Any]: + """Return all plugin config data. + + Used by mypy to determine if cache needs to be discarded. + """ + return self._plugin_data + + def _pydantic_model_class_maker_callback(self, ctx: ClassDefContext) -> None: + transformer = PydanticModelTransformer(ctx.cls, ctx.reason, ctx.api, self.plugin_config) + transformer.transform() + + def _pydantic_model_metaclass_marker_callback(self, ctx: ClassDefContext) -> None: + """Reset dataclass_transform_spec attribute of ModelMetaclass. + + Let the plugin handle it. This behavior can be disabled + if 'debug_dataclass_transform' is set to True', for testing purposes. + """ + if self.plugin_config.debug_dataclass_transform: + return + info_metaclass = ctx.cls.info.declared_metaclass + assert info_metaclass, "callback not passed from 'get_metaclass_hook'" + if getattr(info_metaclass.type, 'dataclass_transform_spec', None): + info_metaclass.type.dataclass_transform_spec = None + + +class PydanticPluginConfig: + """A Pydantic mypy plugin config holder. + + Attributes: + init_forbid_extra: Whether to add a `**kwargs` at the end of the generated `__init__` signature. + init_typed: Whether to annotate fields in the generated `__init__`. + warn_required_dynamic_aliases: Whether to raise required dynamic aliases error. + debug_dataclass_transform: Whether to not reset `dataclass_transform_spec` attribute + of `ModelMetaclass` for testing purposes. + """ + + __slots__ = ( + 'init_forbid_extra', + 'init_typed', + 'warn_required_dynamic_aliases', + 'debug_dataclass_transform', + ) + init_forbid_extra: bool + init_typed: bool + warn_required_dynamic_aliases: bool + debug_dataclass_transform: bool # undocumented + + def __init__(self, options: Options) -> None: + if options.config_file is None: # pragma: no cover + return + + toml_config = parse_toml(options.config_file) + if toml_config is not None: + config = toml_config.get('tool', {}).get('pydantic-mypy', {}) + for key in self.__slots__: + setting = config.get(key, False) + if not isinstance(setting, bool): + raise ValueError(f'Configuration value must be a boolean for key: {key}') + setattr(self, key, setting) + else: + plugin_config = ConfigParser() + plugin_config.read(options.config_file) + for key in self.__slots__: + setting = plugin_config.getboolean(CONFIGFILE_KEY, key, fallback=False) + setattr(self, key, setting) + + def to_data(self) -> dict[str, Any]: + """Returns a dict of config names to their values.""" + return {key: getattr(self, key) for key in self.__slots__} + + +def from_attributes_callback(ctx: MethodContext) -> Type: + """Raise an error if from_attributes is not enabled.""" + model_type: Instance + ctx_type = ctx.type + if isinstance(ctx_type, TypeType): + ctx_type = ctx_type.item + if isinstance(ctx_type, CallableType) and isinstance(ctx_type.ret_type, Instance): + model_type = ctx_type.ret_type # called on the class + elif isinstance(ctx_type, Instance): + model_type = ctx_type # called on an instance (unusual, but still valid) + else: # pragma: no cover + detail = f'ctx.type: {ctx_type} (of type {ctx_type.__class__.__name__})' + error_unexpected_behavior(detail, ctx.api, ctx.context) + return ctx.default_return_type + pydantic_metadata = model_type.type.metadata.get(METADATA_KEY) + if pydantic_metadata is None: + return ctx.default_return_type + if not model_type.type.has_base(BASEMODEL_FULLNAME): + # not a Pydantic v2 model + return ctx.default_return_type + from_attributes = pydantic_metadata.get('config', {}).get('from_attributes') + if from_attributes is not True: + error_from_attributes(model_type.type.name, ctx.api, ctx.context) + return ctx.default_return_type + + +class PydanticModelField: + """Based on mypy.plugins.dataclasses.DataclassAttribute.""" + + def __init__( + self, + name: str, + alias: str | None, + is_frozen: bool, + has_dynamic_alias: bool, + has_default: bool, + strict: bool | None, + line: int, + column: int, + type: Type | None, + info: TypeInfo, + ): + self.name = name + self.alias = alias + self.is_frozen = is_frozen + self.has_dynamic_alias = has_dynamic_alias + self.has_default = has_default + self.strict = strict + self.line = line + self.column = column + self.type = type + self.info = info + + def to_argument( + self, + current_info: TypeInfo, + typed: bool, + model_strict: bool, + force_optional: bool, + use_alias: bool, + api: SemanticAnalyzerPluginInterface, + force_typevars_invariant: bool, + is_root_model_root: bool, + ) -> Argument: + """Based on mypy.plugins.dataclasses.DataclassAttribute.to_argument.""" + variable = self.to_var(current_info, api, use_alias, force_typevars_invariant) + + strict = model_strict if self.strict is None else self.strict + if typed or strict: + type_annotation = self.expand_type(current_info, api, include_root_type=True) + else: + type_annotation = AnyType(TypeOfAny.explicit) + + return Argument( + variable=variable, + type_annotation=type_annotation, + initializer=None, + kind=ARG_OPT + if is_root_model_root + else (ARG_NAMED_OPT if force_optional or self.has_default else ARG_NAMED), + ) + + def expand_type( + self, + current_info: TypeInfo, + api: SemanticAnalyzerPluginInterface, + force_typevars_invariant: bool = False, + include_root_type: bool = False, + ) -> Type | None: + """Based on mypy.plugins.dataclasses.DataclassAttribute.expand_type.""" + if force_typevars_invariant: + # In some cases, mypy will emit an error "Cannot use a covariant type variable as a parameter" + # To prevent that, we add an option to replace typevars with invariant ones while building certain + # method signatures (in particular, `__init__`). There may be a better way to do this, if this causes + # us problems in the future, we should look into why the dataclasses plugin doesn't have this issue. + if isinstance(self.type, TypeVarType): + modified_type = self.type.copy_modified() + modified_type.variance = INVARIANT + self.type = modified_type + + if self.type is not None and self.info.self_type is not None: + # In general, it is not safe to call `expand_type()` during semantic analysis, + # however this plugin is called very late, so all types should be fully ready. + # Also, it is tricky to avoid eager expansion of Self types here (e.g. because + # we serialize attributes). + with state.strict_optional_set(api.options.strict_optional): + filled_with_typevars = fill_typevars(current_info) + # Cannot be TupleType as current_info represents a Pydantic model: + assert isinstance(filled_with_typevars, Instance) + if force_typevars_invariant: + for arg in filled_with_typevars.args: + if isinstance(arg, TypeVarType): + arg.variance = INVARIANT + + expanded_type = expand_type(self.type, {self.info.self_type.id: filled_with_typevars}) + if include_root_type and isinstance(expanded_type, Instance) and is_root_model(expanded_type.type): + # When a root model is used as a field, Pydantic allows both an instance of the root model + # as well as instances of the `root` field type: + root_type = expanded_type.type['root'].type + if root_type is None: + # Happens if the hint for 'root' has unsolved forward references + return expanded_type + expanded_root_type = expand_type_by_instance(root_type, expanded_type) + expanded_type = UnionType([expanded_type, expanded_root_type]) + return expanded_type + return self.type + + def to_var( + self, + current_info: TypeInfo, + api: SemanticAnalyzerPluginInterface, + use_alias: bool, + force_typevars_invariant: bool = False, + ) -> Var: + """Based on mypy.plugins.dataclasses.DataclassAttribute.to_var.""" + if use_alias and self.alias is not None: + name = self.alias + else: + name = self.name + + return Var(name, self.expand_type(current_info, api, force_typevars_invariant)) + + def serialize(self) -> JsonDict: + """Based on mypy.plugins.dataclasses.DataclassAttribute.serialize.""" + assert self.type + return { + 'name': self.name, + 'alias': self.alias, + 'is_frozen': self.is_frozen, + 'has_dynamic_alias': self.has_dynamic_alias, + 'has_default': self.has_default, + 'strict': self.strict, + 'line': self.line, + 'column': self.column, + 'type': self.type.serialize(), + } + + @classmethod + def deserialize(cls, info: TypeInfo, data: JsonDict, api: SemanticAnalyzerPluginInterface) -> PydanticModelField: + """Based on mypy.plugins.dataclasses.DataclassAttribute.deserialize.""" + data = data.copy() + typ = deserialize_and_fixup_type(data.pop('type'), api) + return cls(type=typ, info=info, **data) + + def expand_typevar_from_subtype(self, sub_type: TypeInfo, api: SemanticAnalyzerPluginInterface) -> None: + """Expands type vars in the context of a subtype when an attribute is inherited + from a generic super type. + """ + if self.type is not None: + with state.strict_optional_set(api.options.strict_optional): + self.type = map_type_from_supertype(self.type, sub_type, self.info) + + +class PydanticModelClassVar: + """Based on mypy.plugins.dataclasses.DataclassAttribute. + + ClassVars are ignored by subclasses. + + Attributes: + name: the ClassVar name + """ + + def __init__(self, name): + self.name = name + + @classmethod + def deserialize(cls, data: JsonDict) -> PydanticModelClassVar: + """Based on mypy.plugins.dataclasses.DataclassAttribute.deserialize.""" + data = data.copy() + return cls(**data) + + def serialize(self) -> JsonDict: + """Based on mypy.plugins.dataclasses.DataclassAttribute.serialize.""" + return { + 'name': self.name, + } + + +class PydanticModelTransformer: + """Transform the BaseModel subclass according to the plugin settings. + + Attributes: + tracked_config_fields: A set of field configs that the plugin has to track their value. + """ + + tracked_config_fields: set[str] = { + 'extra', + 'frozen', + 'from_attributes', + 'populate_by_name', + 'validate_by_alias', + 'validate_by_name', + 'alias_generator', + 'strict', + } + + def __init__( + self, + cls: ClassDef, + reason: Expression | Statement, + api: SemanticAnalyzerPluginInterface, + plugin_config: PydanticPluginConfig, + ) -> None: + self._cls = cls + self._reason = reason + self._api = api + + self.plugin_config = plugin_config + + def transform(self) -> bool: + """Configures the BaseModel subclass according to the plugin settings. + + In particular: + + * determines the model config and fields, + * adds a fields-aware signature for the initializer and construct methods + * freezes the class if frozen = True + * stores the fields, config, and if the class is settings in the mypy metadata for access by subclasses + """ + info = self._cls.info + is_a_root_model = is_root_model(info) + config = self.collect_config() + fields, class_vars = self.collect_fields_and_class_vars(config, is_a_root_model) + if fields is None or class_vars is None: + # Some definitions are not ready. We need another pass. + return False + for field in fields: + if field.type is None: + return False + + is_settings = info.has_base(BASESETTINGS_FULLNAME) + self.add_initializer(fields, config, is_settings, is_a_root_model) + self.add_model_construct_method(fields, config, is_settings, is_a_root_model) + self.set_frozen(fields, self._api, frozen=config.frozen is True) + + self.adjust_decorator_signatures() + + info.metadata[METADATA_KEY] = { + 'fields': {field.name: field.serialize() for field in fields}, + 'class_vars': {class_var.name: class_var.serialize() for class_var in class_vars}, + 'config': config.get_values_dict(), + } + + return True + + def adjust_decorator_signatures(self) -> None: + """When we decorate a function `f` with `pydantic.validator(...)`, `pydantic.field_validator` + or `pydantic.serializer(...)`, mypy sees `f` as a regular method taking a `self` instance, + even though pydantic internally wraps `f` with `classmethod` if necessary. + + Teach mypy this by marking any function whose outermost decorator is a `validator()`, + `field_validator()` or `serializer()` call as a `classmethod`. + """ + for sym in self._cls.info.names.values(): + if isinstance(sym.node, Decorator): + first_dec = sym.node.original_decorators[0] + if ( + isinstance(first_dec, CallExpr) + and isinstance(first_dec.callee, NameExpr) + and first_dec.callee.fullname in IMPLICIT_CLASSMETHOD_DECORATOR_FULLNAMES + # @model_validator(mode="after") is an exception, it expects a regular method + and not ( + first_dec.callee.fullname == MODEL_VALIDATOR_FULLNAME + and any( + first_dec.arg_names[i] == 'mode' and isinstance(arg, StrExpr) and arg.value == 'after' + for i, arg in enumerate(first_dec.args) + ) + ) + ): + # TODO: Only do this if the first argument of the decorated function is `cls` + sym.node.func.is_class = True + + def collect_config(self) -> ModelConfigData: # noqa: C901 (ignore complexity) + """Collects the values of the config attributes that are used by the plugin, accounting for parent classes.""" + cls = self._cls + config = ModelConfigData() + + has_config_kwargs = False + has_config_from_namespace = False + + # Handle `class MyModel(BaseModel, =, ...):` + for name, expr in cls.keywords.items(): + config_data = self.get_config_update(name, expr) + if config_data: + has_config_kwargs = True + config.update(config_data) + + # Handle `model_config` + stmt: Statement | None = None + for stmt in cls.defs.body: + if not isinstance(stmt, (AssignmentStmt, ClassDef)): + continue + + if isinstance(stmt, AssignmentStmt): + lhs = stmt.lvalues[0] + if not isinstance(lhs, NameExpr) or lhs.name != 'model_config': + continue + + if isinstance(stmt.rvalue, CallExpr): # calls to `dict` or `ConfigDict` + for arg_name, arg in zip(stmt.rvalue.arg_names, stmt.rvalue.args): + if arg_name is None: + continue + config.update(self.get_config_update(arg_name, arg, lax_extra=True)) + elif isinstance(stmt.rvalue, DictExpr): # dict literals + for key_expr, value_expr in stmt.rvalue.items: + if not isinstance(key_expr, StrExpr): + continue + config.update(self.get_config_update(key_expr.value, value_expr)) + + elif isinstance(stmt, ClassDef): + if stmt.name != 'Config': # 'deprecated' Config-class + continue + for substmt in stmt.defs.body: + if not isinstance(substmt, AssignmentStmt): + continue + lhs = substmt.lvalues[0] + if not isinstance(lhs, NameExpr): + continue + config.update(self.get_config_update(lhs.name, substmt.rvalue)) + + if has_config_kwargs: + self._api.fail( + 'Specifying config in two places is ambiguous, use either Config attribute or class kwargs', + cls, + ) + break + + has_config_from_namespace = True + + if has_config_kwargs or has_config_from_namespace: + if ( + stmt + and config.has_alias_generator + and not (config.validate_by_name or config.populate_by_name) + and self.plugin_config.warn_required_dynamic_aliases + ): + error_required_dynamic_aliases(self._api, stmt) + + for info in cls.info.mro[1:]: # 0 is the current class + if METADATA_KEY not in info.metadata: + continue + + # Each class depends on the set of fields in its ancestors + self._api.add_plugin_dependency(make_wildcard_trigger(info.fullname)) + for name, value in info.metadata[METADATA_KEY]['config'].items(): + config.setdefault(name, value) + return config + + def collect_fields_and_class_vars( + self, model_config: ModelConfigData, is_root_model: bool + ) -> tuple[list[PydanticModelField] | None, list[PydanticModelClassVar] | None]: + """Collects the fields for the model, accounting for parent classes.""" + cls = self._cls + + # First, collect fields and ClassVars belonging to any class in the MRO, ignoring duplicates. + # + # We iterate through the MRO in reverse because attrs defined in the parent must appear + # earlier in the attributes list than attrs defined in the child. See: + # https://docs.python.org/3/library/dataclasses.html#inheritance + # + # However, we also want fields defined in the subtype to override ones defined + # in the parent. We can implement this via a dict without disrupting the attr order + # because dicts preserve insertion order in Python 3.7+. + found_fields: dict[str, PydanticModelField] = {} + found_class_vars: dict[str, PydanticModelClassVar] = {} + for info in reversed(cls.info.mro[1:-1]): # 0 is the current class, -2 is BaseModel, -1 is object + # if BASEMODEL_METADATA_TAG_KEY in info.metadata and BASEMODEL_METADATA_KEY not in info.metadata: + # # We haven't processed the base class yet. Need another pass. + # return None, None + if METADATA_KEY not in info.metadata: + continue + + # Each class depends on the set of attributes in its dataclass ancestors. + self._api.add_plugin_dependency(make_wildcard_trigger(info.fullname)) + + for name, data in info.metadata[METADATA_KEY]['fields'].items(): + field = PydanticModelField.deserialize(info, data, self._api) + # (The following comment comes directly from the dataclasses plugin) + # TODO: We shouldn't be performing type operations during the main + # semantic analysis pass, since some TypeInfo attributes might + # still be in flux. This should be performed in a later phase. + field.expand_typevar_from_subtype(cls.info, self._api) + found_fields[name] = field + + sym_node = cls.info.names.get(name) + if sym_node and sym_node.node and not isinstance(sym_node.node, Var): + self._api.fail( + 'BaseModel field may only be overridden by another field', + sym_node.node, + ) + # Collect ClassVars + for name, data in info.metadata[METADATA_KEY]['class_vars'].items(): + found_class_vars[name] = PydanticModelClassVar.deserialize(data) + + # Second, collect fields and ClassVars belonging to the current class. + current_field_names: set[str] = set() + current_class_vars_names: set[str] = set() + for stmt in self._get_assignment_statements_from_block(cls.defs): + maybe_field = self.collect_field_or_class_var_from_stmt(stmt, model_config, found_class_vars) + if maybe_field is None: + continue + + lhs = stmt.lvalues[0] + assert isinstance(lhs, NameExpr) # collect_field_or_class_var_from_stmt guarantees this + if isinstance(maybe_field, PydanticModelField): + if is_root_model and lhs.name != 'root': + error_extra_fields_on_root_model(self._api, stmt) + else: + current_field_names.add(lhs.name) + found_fields[lhs.name] = maybe_field + elif isinstance(maybe_field, PydanticModelClassVar): + current_class_vars_names.add(lhs.name) + found_class_vars[lhs.name] = maybe_field + + return list(found_fields.values()), list(found_class_vars.values()) + + def _get_assignment_statements_from_if_statement(self, stmt: IfStmt) -> Iterator[AssignmentStmt]: + for body in stmt.body: + if not body.is_unreachable: + yield from self._get_assignment_statements_from_block(body) + if stmt.else_body is not None and not stmt.else_body.is_unreachable: + yield from self._get_assignment_statements_from_block(stmt.else_body) + + def _get_assignment_statements_from_block(self, block: Block) -> Iterator[AssignmentStmt]: + for stmt in block.body: + if isinstance(stmt, AssignmentStmt): + yield stmt + elif isinstance(stmt, IfStmt): + yield from self._get_assignment_statements_from_if_statement(stmt) + + def collect_field_or_class_var_from_stmt( # noqa C901 + self, stmt: AssignmentStmt, model_config: ModelConfigData, class_vars: dict[str, PydanticModelClassVar] + ) -> PydanticModelField | PydanticModelClassVar | None: + """Get pydantic model field from statement. + + Args: + stmt: The statement. + model_config: Configuration settings for the model. + class_vars: ClassVars already known to be defined on the model. + + Returns: + A pydantic model field if it could find the field in statement. Otherwise, `None`. + """ + cls = self._cls + + lhs = stmt.lvalues[0] + if not isinstance(lhs, NameExpr) or not _fields.is_valid_field_name(lhs.name) or lhs.name == 'model_config': + return None + + if not stmt.new_syntax: + if ( + isinstance(stmt.rvalue, CallExpr) + and isinstance(stmt.rvalue.callee, CallExpr) + and isinstance(stmt.rvalue.callee.callee, NameExpr) + and stmt.rvalue.callee.callee.fullname in DECORATOR_FULLNAMES + ): + # This is a (possibly-reused) validator or serializer, not a field + # In particular, it looks something like: my_validator = validator('my_field')(f) + # Eventually, we may want to attempt to respect model_config['ignored_types'] + return None + + if lhs.name in class_vars: + # Class vars are not fields and are not required to be annotated + return None + + # The assignment does not have an annotation, and it's not anything else we recognize + error_untyped_fields(self._api, stmt) + return None + + lhs = stmt.lvalues[0] + if not isinstance(lhs, NameExpr): + return None + + if not _fields.is_valid_field_name(lhs.name) or lhs.name == 'model_config': + return None + + sym = cls.info.names.get(lhs.name) + if sym is None: # pragma: no cover + # This is likely due to a star import (see the dataclasses plugin for a more detailed explanation) + # This is the same logic used in the dataclasses plugin + return None + + node = sym.node + if isinstance(node, PlaceholderNode): # pragma: no cover + # See the PlaceholderNode docstring for more detail about how this can occur + # Basically, it is an edge case when dealing with complex import logic + + # The dataclasses plugin now asserts this cannot happen, but I'd rather not error if it does.. + return None + + if isinstance(node, TypeAlias): + self._api.fail( + 'Type aliases inside BaseModel definitions are not supported at runtime', + node, + ) + # Skip processing this node. This doesn't match the runtime behaviour, + # but the only alternative would be to modify the SymbolTable, + # and it's a little hairy to do that in a plugin. + return None + + if not isinstance(node, Var): # pragma: no cover + # Don't know if this edge case still happens with the `is_valid_field` check above + # but better safe than sorry + + # The dataclasses plugin now asserts this cannot happen, but I'd rather not error if it does.. + return None + + # x: ClassVar[int] is not a field + if node.is_classvar: + return PydanticModelClassVar(lhs.name) + + # x: InitVar[int] is not supported in BaseModel + node_type = get_proper_type(node.type) + if isinstance(node_type, Instance) and node_type.type.fullname == 'dataclasses.InitVar': + self._api.fail( + 'InitVar is not supported in BaseModel', + node, + ) + + has_default = self.get_has_default(stmt) + strict = self.get_strict(stmt) + + if sym.type is None and node.is_final and node.is_inferred: + # This follows the logic from the dataclasses plugin. The following comment is taken verbatim: + # + # This is a special case, assignment like x: Final = 42 is classified + # annotated above, but mypy strips the `Final` turning it into x = 42. + # We do not support inferred types in dataclasses, so we can try inferring + # type for simple literals, and otherwise require an explicit type + # argument for Final[...]. + typ = self._api.analyze_simple_literal_type(stmt.rvalue, is_final=True) + if typ: + node.type = typ + else: + self._api.fail( + 'Need type argument for Final[...] with non-literal default in BaseModel', + stmt, + ) + node.type = AnyType(TypeOfAny.from_error) + + if node.is_final and has_default: + # TODO this path should be removed (see https://github.com/pydantic/pydantic/issues/11119) + return PydanticModelClassVar(lhs.name) + + alias, has_dynamic_alias = self.get_alias_info(stmt) + if ( + has_dynamic_alias + and not (model_config.validate_by_name or model_config.populate_by_name) + and self.plugin_config.warn_required_dynamic_aliases + ): + error_required_dynamic_aliases(self._api, stmt) + is_frozen = self.is_field_frozen(stmt) + + init_type = self._infer_dataclass_attr_init_type(sym, lhs.name, stmt) + return PydanticModelField( + name=lhs.name, + has_dynamic_alias=has_dynamic_alias, + has_default=has_default, + strict=strict, + alias=alias, + is_frozen=is_frozen, + line=stmt.line, + column=stmt.column, + type=init_type, + info=cls.info, + ) + + def _infer_dataclass_attr_init_type(self, sym: SymbolTableNode, name: str, context: Context) -> Type | None: + """Infer __init__ argument type for an attribute. + + In particular, possibly use the signature of __set__. + """ + default = sym.type + if sym.implicit: + return default + t = get_proper_type(sym.type) + + # Perform a simple-minded inference from the signature of __set__, if present. + # We can't use mypy.checkmember here, since this plugin runs before type checking. + # We only support some basic scanerios here, which is hopefully sufficient for + # the vast majority of use cases. + if not isinstance(t, Instance): + return default + setter = t.type.get('__set__') + if setter: + if isinstance(setter.node, FuncDef): + super_info = t.type.get_containing_type_info('__set__') + assert super_info + if setter.type: + setter_type = get_proper_type(map_type_from_supertype(setter.type, t.type, super_info)) + else: + return AnyType(TypeOfAny.unannotated) + if isinstance(setter_type, CallableType) and setter_type.arg_kinds == [ + ARG_POS, + ARG_POS, + ARG_POS, + ]: + return expand_type_by_instance(setter_type.arg_types[2], t) + else: + self._api.fail(f'Unsupported signature for "__set__" in "{t.type.name}"', context) + else: + self._api.fail(f'Unsupported "__set__" in "{t.type.name}"', context) + + return default + + def add_initializer( + self, fields: list[PydanticModelField], config: ModelConfigData, is_settings: bool, is_root_model: bool + ) -> None: + """Adds a fields-aware `__init__` method to the class. + + The added `__init__` will be annotated with types vs. all `Any` depending on the plugin settings. + """ + if '__init__' in self._cls.info.names and not self._cls.info.names['__init__'].plugin_generated: + return # Don't generate an __init__ if one already exists + + typed = self.plugin_config.init_typed + model_strict = bool(config.strict) + use_alias = not (config.validate_by_name or config.populate_by_name) and config.validate_by_alias is not False + requires_dynamic_aliases = bool(config.has_alias_generator and not config.validate_by_name) + args = self.get_field_arguments( + fields, + typed=typed, + model_strict=model_strict, + requires_dynamic_aliases=requires_dynamic_aliases, + use_alias=use_alias, + is_settings=is_settings, + is_root_model=is_root_model, + force_typevars_invariant=True, + ) + + if is_settings: + base_settings_node = self._api.lookup_fully_qualified(BASESETTINGS_FULLNAME).node + assert isinstance(base_settings_node, TypeInfo) + if '__init__' in base_settings_node.names: + base_settings_init_node = base_settings_node.names['__init__'].node + assert isinstance(base_settings_init_node, FuncDef) + if base_settings_init_node is not None and base_settings_init_node.type is not None: + func_type = base_settings_init_node.type + assert isinstance(func_type, CallableType) + for arg_idx, arg_name in enumerate(func_type.arg_names): + if arg_name is None or arg_name.startswith('__') or not arg_name.startswith('_'): + continue + analyzed_variable_type = self._api.anal_type(func_type.arg_types[arg_idx]) + if analyzed_variable_type is not None and arg_name == '_cli_settings_source': + # _cli_settings_source is defined as CliSettingsSource[Any], and as such + # the Any causes issues with --disallow-any-explicit. As a workaround, change + # the Any type (as if CliSettingsSource was left unparameterized): + analyzed_variable_type = analyzed_variable_type.accept( + ChangeExplicitTypeOfAny(TypeOfAny.from_omitted_generics) + ) + variable = Var(arg_name, analyzed_variable_type) + args.append(Argument(variable, analyzed_variable_type, None, ARG_OPT)) + + if not self.should_init_forbid_extra(fields, config): + var = Var('kwargs') + args.append(Argument(var, AnyType(TypeOfAny.explicit), None, ARG_STAR2)) + + add_method(self._api, self._cls, '__init__', args=args, return_type=NoneType()) + + def add_model_construct_method( + self, + fields: list[PydanticModelField], + config: ModelConfigData, + is_settings: bool, + is_root_model: bool, + ) -> None: + """Adds a fully typed `model_construct` classmethod to the class. + + Similar to the fields-aware __init__ method, but always uses the field names (not aliases), + and does not treat settings fields as optional. + """ + set_str = self._api.named_type(f'{BUILTINS_NAME}.set', [self._api.named_type(f'{BUILTINS_NAME}.str')]) + optional_set_str = UnionType([set_str, NoneType()]) + fields_set_argument = Argument(Var('_fields_set', optional_set_str), optional_set_str, None, ARG_OPT) + with state.strict_optional_set(self._api.options.strict_optional): + args = self.get_field_arguments( + fields, + typed=True, + model_strict=bool(config.strict), + requires_dynamic_aliases=False, + use_alias=False, + is_settings=is_settings, + is_root_model=is_root_model, + ) + if not self.should_init_forbid_extra(fields, config): + var = Var('kwargs') + args.append(Argument(var, AnyType(TypeOfAny.explicit), None, ARG_STAR2)) + + args = args + [fields_set_argument] if is_root_model else [fields_set_argument] + args + + add_method( + self._api, + self._cls, + 'model_construct', + args=args, + return_type=fill_typevars(self._cls.info), + is_classmethod=True, + ) + + def set_frozen(self, fields: list[PydanticModelField], api: SemanticAnalyzerPluginInterface, frozen: bool) -> None: + """Marks all fields as properties so that attempts to set them trigger mypy errors. + + This is the same approach used by the attrs and dataclasses plugins. + """ + info = self._cls.info + for field in fields: + sym_node = info.names.get(field.name) + if sym_node is not None: + var = sym_node.node + if isinstance(var, Var): + var.is_property = frozen or field.is_frozen + elif isinstance(var, PlaceholderNode) and not self._api.final_iteration: + # See https://github.com/pydantic/pydantic/issues/5191 to hit this branch for test coverage + self._api.defer() + else: # pragma: no cover + # I don't know whether it's possible to hit this branch, but I've added it for safety + try: + var_str = str(var) + except TypeError: + # This happens for PlaceholderNode; perhaps it will happen for other types in the future.. + var_str = repr(var) + detail = f'sym_node.node: {var_str} (of type {var.__class__})' + error_unexpected_behavior(detail, self._api, self._cls) + else: + var = field.to_var(info, api, use_alias=False) + var.info = info + var.is_property = frozen + var._fullname = info.fullname + '.' + var.name + info.names[var.name] = SymbolTableNode(MDEF, var) + + def get_config_update(self, name: str, arg: Expression, lax_extra: bool = False) -> ModelConfigData | None: + """Determines the config update due to a single kwarg in the ConfigDict definition. + + Warns if a tracked config attribute is set to a value the plugin doesn't know how to interpret (e.g., an int) + """ + if name not in self.tracked_config_fields: + return None + if name == 'extra': + if isinstance(arg, StrExpr): + forbid_extra = arg.value == 'forbid' + elif isinstance(arg, MemberExpr): + forbid_extra = arg.name == 'forbid' + else: + if not lax_extra: + # Only emit an error for other types of `arg` (e.g., `NameExpr`, `ConditionalExpr`, etc.) when + # reading from a config class, etc. If a ConfigDict is used, then we don't want to emit an error + # because you'll get type checking from the ConfigDict itself. + # + # It would be nice if we could introspect the types better otherwise, but I don't know what the API + # is to evaluate an expr into its type and then check if that type is compatible with the expected + # type. Note that you can still get proper type checking via: `model_config = ConfigDict(...)`, just + # if you don't use an explicit string, the plugin won't be able to infer whether extra is forbidden. + error_invalid_config_value(name, self._api, arg) + return None + return ModelConfigData(forbid_extra=forbid_extra) + if name == 'alias_generator': + has_alias_generator = True + if isinstance(arg, NameExpr) and arg.fullname == 'builtins.None': + has_alias_generator = False + return ModelConfigData(has_alias_generator=has_alias_generator) + if isinstance(arg, NameExpr) and arg.fullname in ('builtins.True', 'builtins.False'): + return ModelConfigData(**{name: arg.fullname == 'builtins.True'}) + error_invalid_config_value(name, self._api, arg) + return None + + @staticmethod + def get_has_default(stmt: AssignmentStmt) -> bool: + """Returns a boolean indicating whether the field defined in `stmt` is a required field.""" + expr = stmt.rvalue + if isinstance(expr, TempNode): + # TempNode means annotation-only, so has no default + return False + if isinstance(expr, CallExpr) and isinstance(expr.callee, RefExpr) and expr.callee.fullname == FIELD_FULLNAME: + # The "default value" is a call to `Field`; at this point, the field has a default if and only if: + # * there is a positional argument that is not `...` + # * there is a keyword argument named "default" that is not `...` + # * there is a "default_factory" that is not `None` + for arg, name in zip(expr.args, expr.arg_names): + # If name is None, then this arg is the default because it is the only positional argument. + if name is None or name == 'default': + return arg.__class__ is not EllipsisExpr + if name == 'default_factory': + return not (isinstance(arg, NameExpr) and arg.fullname == 'builtins.None') + return False + # Has no default if the "default value" is Ellipsis (i.e., `field_name: Annotation = ...`) + return not isinstance(expr, EllipsisExpr) + + @staticmethod + def get_strict(stmt: AssignmentStmt) -> bool | None: + """Returns a the `strict` value of a field if defined, otherwise `None`.""" + expr = stmt.rvalue + if isinstance(expr, CallExpr) and isinstance(expr.callee, RefExpr) and expr.callee.fullname == FIELD_FULLNAME: + for arg, name in zip(expr.args, expr.arg_names): + if name != 'strict': + continue + if isinstance(arg, NameExpr): + if arg.fullname == 'builtins.True': + return True + elif arg.fullname == 'builtins.False': + return False + return None + return None + + @staticmethod + def get_alias_info(stmt: AssignmentStmt) -> tuple[str | None, bool]: + """Returns a pair (alias, has_dynamic_alias), extracted from the declaration of the field defined in `stmt`. + + `has_dynamic_alias` is True if and only if an alias is provided, but not as a string literal. + If `has_dynamic_alias` is True, `alias` will be None. + """ + expr = stmt.rvalue + if isinstance(expr, TempNode): + # TempNode means annotation-only + return None, False + + if not ( + isinstance(expr, CallExpr) and isinstance(expr.callee, RefExpr) and expr.callee.fullname == FIELD_FULLNAME + ): + # Assigned value is not a call to pydantic.fields.Field + return None, False + + if 'validation_alias' in expr.arg_names: + arg = expr.args[expr.arg_names.index('validation_alias')] + elif 'alias' in expr.arg_names: + arg = expr.args[expr.arg_names.index('alias')] + else: + return None, False + + if isinstance(arg, StrExpr): + return arg.value, False + else: + return None, True + + @staticmethod + def is_field_frozen(stmt: AssignmentStmt) -> bool: + """Returns whether the field is frozen, extracted from the declaration of the field defined in `stmt`. + + Note that this is only whether the field was declared to be frozen in a ` = Field(frozen=True)` + sense; this does not determine whether the field is frozen because the entire model is frozen; that is + handled separately. + """ + expr = stmt.rvalue + if isinstance(expr, TempNode): + # TempNode means annotation-only + return False + + if not ( + isinstance(expr, CallExpr) and isinstance(expr.callee, RefExpr) and expr.callee.fullname == FIELD_FULLNAME + ): + # Assigned value is not a call to pydantic.fields.Field + return False + + for i, arg_name in enumerate(expr.arg_names): + if arg_name == 'frozen': + arg = expr.args[i] + return isinstance(arg, NameExpr) and arg.fullname == 'builtins.True' + return False + + def get_field_arguments( + self, + fields: list[PydanticModelField], + typed: bool, + model_strict: bool, + use_alias: bool, + requires_dynamic_aliases: bool, + is_settings: bool, + is_root_model: bool, + force_typevars_invariant: bool = False, + ) -> list[Argument]: + """Helper function used during the construction of the `__init__` and `model_construct` method signatures. + + Returns a list of mypy Argument instances for use in the generated signatures. + """ + info = self._cls.info + arguments = [ + field.to_argument( + info, + typed=typed, + model_strict=model_strict, + force_optional=requires_dynamic_aliases or is_settings, + use_alias=use_alias, + api=self._api, + force_typevars_invariant=force_typevars_invariant, + is_root_model_root=is_root_model and field.name == 'root', + ) + for field in fields + if not (use_alias and field.has_dynamic_alias) + ] + return arguments + + def should_init_forbid_extra(self, fields: list[PydanticModelField], config: ModelConfigData) -> bool: + """Indicates whether the generated `__init__` should get a `**kwargs` at the end of its signature. + + We disallow arbitrary kwargs if the extra config setting is "forbid", or if the plugin config says to, + *unless* a required dynamic alias is present (since then we can't determine a valid signature). + """ + if not (config.validate_by_name or config.populate_by_name): + if self.is_dynamic_alias_present(fields, bool(config.has_alias_generator)): + return False + if config.forbid_extra: + return True + return self.plugin_config.init_forbid_extra + + @staticmethod + def is_dynamic_alias_present(fields: list[PydanticModelField], has_alias_generator: bool) -> bool: + """Returns whether any fields on the model have a "dynamic alias", i.e., an alias that cannot be + determined during static analysis. + """ + for field in fields: + if field.has_dynamic_alias: + return True + if has_alias_generator: + for field in fields: + if field.alias is None: + return True + return False + + +class ChangeExplicitTypeOfAny(TypeTranslator): + """A type translator used to change type of Any's, if explicit.""" + + def __init__(self, type_of_any: int) -> None: + self._type_of_any = type_of_any + super().__init__() + + def visit_any(self, t: AnyType) -> Type: # noqa: D102 + if t.type_of_any == TypeOfAny.explicit: + return t.copy_modified(type_of_any=self._type_of_any) + else: + return t + + +class ModelConfigData: + """Pydantic mypy plugin model config class.""" + + def __init__( + self, + forbid_extra: bool | None = None, + frozen: bool | None = None, + from_attributes: bool | None = None, + populate_by_name: bool | None = None, + validate_by_alias: bool | None = None, + validate_by_name: bool | None = None, + has_alias_generator: bool | None = None, + strict: bool | None = None, + ): + self.forbid_extra = forbid_extra + self.frozen = frozen + self.from_attributes = from_attributes + self.populate_by_name = populate_by_name + self.validate_by_alias = validate_by_alias + self.validate_by_name = validate_by_name + self.has_alias_generator = has_alias_generator + self.strict = strict + + def get_values_dict(self) -> dict[str, Any]: + """Returns a dict of Pydantic model config names to their values. + + It includes the config if config value is not `None`. + """ + return {k: v for k, v in self.__dict__.items() if v is not None} + + def update(self, config: ModelConfigData | None) -> None: + """Update Pydantic model config values.""" + if config is None: + return + for k, v in config.get_values_dict().items(): + setattr(self, k, v) + + def setdefault(self, key: str, value: Any) -> None: + """Set default value for Pydantic model config if config value is `None`.""" + if getattr(self, key) is None: + setattr(self, key, value) + + +def is_root_model(info: TypeInfo) -> bool: + """Return whether the type info is a root model subclass (or the `RootModel` class itself).""" + return info.has_base(ROOT_MODEL_FULLNAME) + + +ERROR_ORM = ErrorCode('pydantic-orm', 'Invalid from_attributes call', 'Pydantic') +ERROR_CONFIG = ErrorCode('pydantic-config', 'Invalid config value', 'Pydantic') +ERROR_ALIAS = ErrorCode('pydantic-alias', 'Dynamic alias disallowed', 'Pydantic') +ERROR_UNEXPECTED = ErrorCode('pydantic-unexpected', 'Unexpected behavior', 'Pydantic') +ERROR_UNTYPED = ErrorCode('pydantic-field', 'Untyped field disallowed', 'Pydantic') +ERROR_FIELD_DEFAULTS = ErrorCode('pydantic-field', 'Invalid Field defaults', 'Pydantic') +ERROR_EXTRA_FIELD_ROOT_MODEL = ErrorCode('pydantic-field', 'Extra field on RootModel subclass', 'Pydantic') + + +def error_from_attributes(model_name: str, api: CheckerPluginInterface, context: Context) -> None: + """Emits an error when the model does not have `from_attributes=True`.""" + api.fail(f'"{model_name}" does not have from_attributes=True', context, code=ERROR_ORM) + + +def error_invalid_config_value(name: str, api: SemanticAnalyzerPluginInterface, context: Context) -> None: + """Emits an error when the config value is invalid.""" + api.fail(f'Invalid value for "Config.{name}"', context, code=ERROR_CONFIG) + + +def error_required_dynamic_aliases(api: SemanticAnalyzerPluginInterface, context: Context) -> None: + """Emits required dynamic aliases error. + + This will be called when `warn_required_dynamic_aliases=True`. + """ + api.fail('Required dynamic aliases disallowed', context, code=ERROR_ALIAS) + + +def error_unexpected_behavior( + detail: str, api: CheckerPluginInterface | SemanticAnalyzerPluginInterface, context: Context +) -> None: # pragma: no cover + """Emits unexpected behavior error.""" + # Can't think of a good way to test this, but I confirmed it renders as desired by adding to a non-error path + link = 'https://github.com/pydantic/pydantic/issues/new/choose' + full_message = f'The pydantic mypy plugin ran into unexpected behavior: {detail}\n' + full_message += f'Please consider reporting this bug at {link} so we can try to fix it!' + api.fail(full_message, context, code=ERROR_UNEXPECTED) + + +def error_untyped_fields(api: SemanticAnalyzerPluginInterface, context: Context) -> None: + """Emits an error when there is an untyped field in the model.""" + api.fail('Untyped fields disallowed', context, code=ERROR_UNTYPED) + + +def error_extra_fields_on_root_model(api: CheckerPluginInterface, context: Context) -> None: + """Emits an error when there is more than just a root field defined for a subclass of RootModel.""" + api.fail('Only `root` is allowed as a field of a `RootModel`', context, code=ERROR_EXTRA_FIELD_ROOT_MODEL) + + +def add_method( + api: SemanticAnalyzerPluginInterface | CheckerPluginInterface, + cls: ClassDef, + name: str, + args: list[Argument], + return_type: Type, + self_type: Type | None = None, + tvar_def: TypeVarType | None = None, + is_classmethod: bool = False, +) -> None: + """Very closely related to `mypy.plugins.common.add_method_to_class`, with a few pydantic-specific changes.""" + info = cls.info + + # First remove any previously generated methods with the same name + # to avoid clashes and problems in the semantic analyzer. + if name in info.names: + sym = info.names[name] + if sym.plugin_generated and isinstance(sym.node, FuncDef): + cls.defs.body.remove(sym.node) # pragma: no cover + + if isinstance(api, SemanticAnalyzerPluginInterface): + function_type = api.named_type('builtins.function') + else: + function_type = api.named_generic_type('builtins.function', []) + + if is_classmethod: + self_type = self_type or TypeType(fill_typevars(info)) + first = [Argument(Var('_cls'), self_type, None, ARG_POS, True)] + else: + self_type = self_type or fill_typevars(info) + # `self` is positional *ONLY* here, but this can't be expressed + # fully in the mypy internal API. ARG_POS is the closest we can get. + # Using ARG_POS will, however, give mypy errors if a `self` field + # is present on a model: + # + # Name "self" already defined (possibly by an import) [no-redef] + # + # As a workaround, we give this argument a name that will + # never conflict. By its positional nature, this name will not + # be used or exposed to users. + first = [Argument(Var('__pydantic_self__'), self_type, None, ARG_POS)] + args = first + args + + arg_types, arg_names, arg_kinds = [], [], [] + for arg in args: + assert arg.type_annotation, 'All arguments must be fully typed.' + arg_types.append(arg.type_annotation) + arg_names.append(arg.variable.name) + arg_kinds.append(arg.kind) + + signature = CallableType( + arg_types, arg_kinds, arg_names, return_type, function_type, variables=[tvar_def] if tvar_def else None + ) + + func = FuncDef(name, args, Block([PassStmt()])) + func.info = info + func.type = set_callable_name(signature, func) + func.is_class = is_classmethod + func._fullname = info.fullname + '.' + name + func.line = info.line + + # NOTE: we would like the plugin generated node to dominate, but we still + # need to keep any existing definitions so they get semantically analyzed. + if name in info.names: + # Get a nice unique name instead. + r_name = get_unique_redefinition_name(name, info.names) + info.names[r_name] = info.names[name] + + # Add decorator for is_classmethod + # The dataclasses plugin claims this is unnecessary for classmethods, but not including it results in a + # signature incompatible with the superclass, which causes mypy errors to occur for every subclass of BaseModel. + if is_classmethod: + func.is_decorated = True + v = Var(name, func.type) + v.info = info + v._fullname = func._fullname + v.is_classmethod = True + dec = Decorator(func, [NameExpr('classmethod')], v) + dec.line = info.line + sym = SymbolTableNode(MDEF, dec) + else: + sym = SymbolTableNode(MDEF, func) + sym.plugin_generated = True + info.names[name] = sym + + info.defn.defs.body.append(func) + + +def parse_toml(config_file: str) -> dict[str, Any] | None: + """Returns a dict of config keys to values. + + It reads configs from toml file and returns `None` if the file is not a toml file. + """ + if not config_file.endswith('.toml'): + return None + + if sys.version_info >= (3, 11): + import tomllib as toml_ + else: + try: + import tomli as toml_ + except ImportError: # pragma: no cover + import warnings + + warnings.warn('No TOML parser installed, cannot read configuration from `pyproject.toml`.') + return None + + with open(config_file, 'rb') as rf: + return toml_.load(rf) diff --git a/lib/python3.12/site-packages/pydantic/networks.py b/lib/python3.12/site-packages/pydantic/networks.py new file mode 100644 index 0000000000000000000000000000000000000000..222157873da628cee3046709dfd3a5825a30769b --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/networks.py @@ -0,0 +1,1312 @@ +"""The networks module contains types for common network-related fields.""" + +from __future__ import annotations as _annotations + +import dataclasses as _dataclasses +import re +from dataclasses import fields +from functools import lru_cache +from importlib.metadata import version +from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network +from typing import TYPE_CHECKING, Annotated, Any, ClassVar + +from pydantic_core import ( + MultiHostHost, + PydanticCustomError, + PydanticSerializationUnexpectedValue, + SchemaSerializer, + core_schema, +) +from pydantic_core import MultiHostUrl as _CoreMultiHostUrl +from pydantic_core import Url as _CoreUrl +from typing_extensions import Self, TypeAlias + +from pydantic.errors import PydanticUserError + +from ._internal import _repr, _schema_generation_shared +from ._migration import getattr_migration +from .annotated_handlers import GetCoreSchemaHandler +from .json_schema import JsonSchemaValue +from .type_adapter import TypeAdapter + +if TYPE_CHECKING: + import email_validator + + NetworkType: TypeAlias = 'str | bytes | int | tuple[str | bytes | int, str | int]' + +else: + email_validator = None + + +__all__ = [ + 'AnyUrl', + 'AnyHttpUrl', + 'FileUrl', + 'FtpUrl', + 'HttpUrl', + 'WebsocketUrl', + 'AnyWebsocketUrl', + 'UrlConstraints', + 'EmailStr', + 'NameEmail', + 'IPvAnyAddress', + 'IPvAnyInterface', + 'IPvAnyNetwork', + 'PostgresDsn', + 'CockroachDsn', + 'AmqpDsn', + 'RedisDsn', + 'MongoDsn', + 'KafkaDsn', + 'NatsDsn', + 'validate_email', + 'MySQLDsn', + 'MariaDBDsn', + 'ClickHouseDsn', + 'SnowflakeDsn', +] + + +@_dataclasses.dataclass +class UrlConstraints: + """Url constraints. + + Attributes: + max_length: The maximum length of the url. Defaults to `None`. + allowed_schemes: The allowed schemes. Defaults to `None`. + host_required: Whether the host is required. Defaults to `None`. + default_host: The default host. Defaults to `None`. + default_port: The default port. Defaults to `None`. + default_path: The default path. Defaults to `None`. + """ + + max_length: int | None = None + allowed_schemes: list[str] | None = None + host_required: bool | None = None + default_host: str | None = None + default_port: int | None = None + default_path: str | None = None + + def __hash__(self) -> int: + return hash( + ( + self.max_length, + tuple(self.allowed_schemes) if self.allowed_schemes is not None else None, + self.host_required, + self.default_host, + self.default_port, + self.default_path, + ) + ) + + @property + def defined_constraints(self) -> dict[str, Any]: + """Fetch a key / value mapping of constraints to values that are not None. Used for core schema updates.""" + return {field.name: value for field in fields(self) if (value := getattr(self, field.name)) is not None} + + def __get_pydantic_core_schema__(self, source: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + schema = handler(source) + + # for function-wrap schemas, url constraints is applied to the inner schema + # because when we generate schemas for urls, we wrap a core_schema.url_schema() with a function-wrap schema + # that helps with validation on initialization, see _BaseUrl and _BaseMultiHostUrl below. + schema_to_mutate = schema['schema'] if schema['type'] == 'function-wrap' else schema + if annotated_type := schema_to_mutate['type'] not in ('url', 'multi-host-url'): + raise PydanticUserError( + f"'UrlConstraints' cannot annotate '{annotated_type}'.", code='invalid-annotated-type' + ) + for constraint_key, constraint_value in self.defined_constraints.items(): + schema_to_mutate[constraint_key] = constraint_value + return schema + + +class _BaseUrl: + _constraints: ClassVar[UrlConstraints] = UrlConstraints() + _url: _CoreUrl + + def __init__(self, url: str | _CoreUrl | _BaseUrl) -> None: + self._url = _build_type_adapter(self.__class__).validate_python(url)._url + + @property + def scheme(self) -> str: + """The scheme part of the URL. + + e.g. `https` in `https://user:pass@host:port/path?query#fragment` + """ + return self._url.scheme + + @property + def username(self) -> str | None: + """The username part of the URL, or `None`. + + e.g. `user` in `https://user:pass@host:port/path?query#fragment` + """ + return self._url.username + + @property + def password(self) -> str | None: + """The password part of the URL, or `None`. + + e.g. `pass` in `https://user:pass@host:port/path?query#fragment` + """ + return self._url.password + + @property + def host(self) -> str | None: + """The host part of the URL, or `None`. + + If the URL must be punycode encoded, this is the encoded host, e.g if the input URL is `https://£££.com`, + `host` will be `xn--9aaa.com` + """ + return self._url.host + + def unicode_host(self) -> str | None: + """The host part of the URL as a unicode string, or `None`. + + e.g. `host` in `https://user:pass@host:port/path?query#fragment` + + If the URL must be punycode encoded, this is the decoded host, e.g if the input URL is `https://£££.com`, + `unicode_host()` will be `£££.com` + """ + return self._url.unicode_host() + + @property + def port(self) -> int | None: + """The port part of the URL, or `None`. + + e.g. `port` in `https://user:pass@host:port/path?query#fragment` + """ + return self._url.port + + @property + def path(self) -> str | None: + """The path part of the URL, or `None`. + + e.g. `/path` in `https://user:pass@host:port/path?query#fragment` + """ + return self._url.path + + @property + def query(self) -> str | None: + """The query part of the URL, or `None`. + + e.g. `query` in `https://user:pass@host:port/path?query#fragment` + """ + return self._url.query + + def query_params(self) -> list[tuple[str, str]]: + """The query part of the URL as a list of key-value pairs. + + e.g. `[('foo', 'bar')]` in `https://user:pass@host:port/path?foo=bar#fragment` + """ + return self._url.query_params() + + @property + def fragment(self) -> str | None: + """The fragment part of the URL, or `None`. + + e.g. `fragment` in `https://user:pass@host:port/path?query#fragment` + """ + return self._url.fragment + + def unicode_string(self) -> str: + """The URL as a unicode string, unlike `__str__()` this will not punycode encode the host. + + If the URL must be punycode encoded, this is the decoded string, e.g if the input URL is `https://£££.com`, + `unicode_string()` will be `https://£££.com` + """ + return self._url.unicode_string() + + def encoded_string(self) -> str: + """The URL's encoded string representation via __str__(). + + This returns the punycode-encoded host version of the URL as a string. + """ + return str(self) + + def __str__(self) -> str: + """The URL as a string, this will punycode encode the host if required.""" + return str(self._url) + + def __repr__(self) -> str: + return f'{self.__class__.__name__}({str(self._url)!r})' + + def __deepcopy__(self, memo: dict) -> Self: + return self.__class__(self._url) + + def __eq__(self, other: Any) -> bool: + return self.__class__ is other.__class__ and self._url == other._url + + def __lt__(self, other: Any) -> bool: + return self.__class__ is other.__class__ and self._url < other._url + + def __gt__(self, other: Any) -> bool: + return self.__class__ is other.__class__ and self._url > other._url + + def __le__(self, other: Any) -> bool: + return self.__class__ is other.__class__ and self._url <= other._url + + def __ge__(self, other: Any) -> bool: + return self.__class__ is other.__class__ and self._url >= other._url + + def __hash__(self) -> int: + return hash(self._url) + + def __len__(self) -> int: + return len(str(self._url)) + + @classmethod + def build( + cls, + *, + scheme: str, + username: str | None = None, + password: str | None = None, + host: str, + port: int | None = None, + path: str | None = None, + query: str | None = None, + fragment: str | None = None, + ) -> Self: + """Build a new `Url` instance from its component parts. + + Args: + scheme: The scheme part of the URL. + username: The username part of the URL, or omit for no username. + password: The password part of the URL, or omit for no password. + host: The host part of the URL. + port: The port part of the URL, or omit for no port. + path: The path part of the URL, or omit for no path. + query: The query part of the URL, or omit for no query. + fragment: The fragment part of the URL, or omit for no fragment. + + Returns: + An instance of URL + """ + return cls( + _CoreUrl.build( + scheme=scheme, + username=username, + password=password, + host=host, + port=port, + path=path, + query=query, + fragment=fragment, + ) + ) + + @classmethod + def serialize_url(cls, url: Any, info: core_schema.SerializationInfo) -> str | Self: + if not isinstance(url, cls): + raise PydanticSerializationUnexpectedValue( + f"Expected `{cls}` but got `{type(url)}` with value `'{url}'` - serialized value may not be as expected." + ) + if info.mode == 'json': + return str(url) + return url + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[_BaseUrl], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + def wrap_val(v, h): + if isinstance(v, source): + return v + if isinstance(v, _BaseUrl): + v = str(v) + core_url = h(v) + instance = source.__new__(source) + instance._url = core_url + return instance + + return core_schema.no_info_wrap_validator_function( + wrap_val, + schema=core_schema.url_schema(**cls._constraints.defined_constraints), + serialization=core_schema.plain_serializer_function_ser_schema( + cls.serialize_url, info_arg=True, when_used='always' + ), + ) + + @classmethod + def __get_pydantic_json_schema__( + cls, core_schema: core_schema.CoreSchema, handler: _schema_generation_shared.GetJsonSchemaHandler + ) -> JsonSchemaValue: + # we use the url schema for json schema generation, but we might have to extract it from + # the function-wrap schema we use as a tool for validation on initialization + inner_schema = core_schema['schema'] if core_schema['type'] == 'function-wrap' else core_schema + return handler(inner_schema) + + __pydantic_serializer__ = SchemaSerializer(core_schema.any_schema(serialization=core_schema.to_string_ser_schema())) + + +class _BaseMultiHostUrl: + _constraints: ClassVar[UrlConstraints] = UrlConstraints() + _url: _CoreMultiHostUrl + + def __init__(self, url: str | _CoreMultiHostUrl | _BaseMultiHostUrl) -> None: + self._url = _build_type_adapter(self.__class__).validate_python(url)._url + + @property + def scheme(self) -> str: + """The scheme part of the URL. + + e.g. `https` in `https://foo.com,bar.com/path?query#fragment` + """ + return self._url.scheme + + @property + def path(self) -> str | None: + """The path part of the URL, or `None`. + + e.g. `/path` in `https://foo.com,bar.com/path?query#fragment` + """ + return self._url.path + + @property + def query(self) -> str | None: + """The query part of the URL, or `None`. + + e.g. `query` in `https://foo.com,bar.com/path?query#fragment` + """ + return self._url.query + + def query_params(self) -> list[tuple[str, str]]: + """The query part of the URL as a list of key-value pairs. + + e.g. `[('foo', 'bar')]` in `https://foo.com,bar.com/path?foo=bar#fragment` + """ + return self._url.query_params() + + @property + def fragment(self) -> str | None: + """The fragment part of the URL, or `None`. + + e.g. `fragment` in `https://foo.com,bar.com/path?query#fragment` + """ + return self._url.fragment + + def hosts(self) -> list[MultiHostHost]: + '''The hosts of the `MultiHostUrl` as [`MultiHostHost`][pydantic_core.MultiHostHost] typed dicts. + + ```python + from pydantic_core import MultiHostUrl + + mhu = MultiHostUrl('https://foo.com:123,foo:bar@bar.com/path') + print(mhu.hosts()) + """ + [ + {'username': None, 'password': None, 'host': 'foo.com', 'port': 123}, + {'username': 'foo', 'password': 'bar', 'host': 'bar.com', 'port': 443} + ] + ``` + Returns: + A list of dicts, each representing a host. + ''' + return self._url.hosts() + + def encoded_string(self) -> str: + """The URL's encoded string representation via __str__(). + + This returns the punycode-encoded host version of the URL as a string. + """ + return str(self) + + def unicode_string(self) -> str: + """The URL as a unicode string, unlike `__str__()` this will not punycode encode the hosts.""" + return self._url.unicode_string() + + def __str__(self) -> str: + """The URL as a string, this will punycode encode the host if required.""" + return str(self._url) + + def __repr__(self) -> str: + return f'{self.__class__.__name__}({str(self._url)!r})' + + def __deepcopy__(self, memo: dict) -> Self: + return self.__class__(self._url) + + def __eq__(self, other: Any) -> bool: + return self.__class__ is other.__class__ and self._url == other._url + + def __hash__(self) -> int: + return hash(self._url) + + def __len__(self) -> int: + return len(str(self._url)) + + @classmethod + def build( + cls, + *, + scheme: str, + hosts: list[MultiHostHost] | None = None, + username: str | None = None, + password: str | None = None, + host: str | None = None, + port: int | None = None, + path: str | None = None, + query: str | None = None, + fragment: str | None = None, + ) -> Self: + """Build a new `MultiHostUrl` instance from its component parts. + + This method takes either `hosts` - a list of `MultiHostHost` typed dicts, or the individual components + `username`, `password`, `host` and `port`. + + Args: + scheme: The scheme part of the URL. + hosts: Multiple hosts to build the URL from. + username: The username part of the URL. + password: The password part of the URL. + host: The host part of the URL. + port: The port part of the URL. + path: The path part of the URL. + query: The query part of the URL, or omit for no query. + fragment: The fragment part of the URL, or omit for no fragment. + + Returns: + An instance of `MultiHostUrl` + """ + return cls( + _CoreMultiHostUrl.build( + scheme=scheme, + hosts=hosts, + username=username, + password=password, + host=host, + port=port, + path=path, + query=query, + fragment=fragment, + ) + ) + + @classmethod + def serialize_url(cls, url: Any, info: core_schema.SerializationInfo) -> str | Self: + if not isinstance(url, cls): + raise PydanticSerializationUnexpectedValue( + f"Expected `{cls}` but got `{type(url)}` with value `'{url}'` - serialized value may not be as expected." + ) + if info.mode == 'json': + return str(url) + return url + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[_BaseMultiHostUrl], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + def wrap_val(v, h): + if isinstance(v, source): + return v + if isinstance(v, _BaseMultiHostUrl): + v = str(v) + core_url = h(v) + instance = source.__new__(source) + instance._url = core_url + return instance + + return core_schema.no_info_wrap_validator_function( + wrap_val, + schema=core_schema.multi_host_url_schema(**cls._constraints.defined_constraints), + serialization=core_schema.plain_serializer_function_ser_schema( + cls.serialize_url, info_arg=True, when_used='always' + ), + ) + + @classmethod + def __get_pydantic_json_schema__( + cls, core_schema: core_schema.CoreSchema, handler: _schema_generation_shared.GetJsonSchemaHandler + ) -> JsonSchemaValue: + # we use the url schema for json schema generation, but we might have to extract it from + # the function-wrap schema we use as a tool for validation on initialization + inner_schema = core_schema['schema'] if core_schema['type'] == 'function-wrap' else core_schema + return handler(inner_schema) + + __pydantic_serializer__ = SchemaSerializer(core_schema.any_schema(serialization=core_schema.to_string_ser_schema())) + + +@lru_cache +def _build_type_adapter(cls: type[_BaseUrl | _BaseMultiHostUrl]) -> TypeAdapter: + return TypeAdapter(cls) + + +class AnyUrl(_BaseUrl): + """Base type for all URLs. + + * Any scheme allowed + * Top-level domain (TLD) not required + * Host not required + + Assuming an input URL of `http://samuel:pass@example.com:8000/the/path/?query=here#fragment=is;this=bit`, + the types export the following properties: + + - `scheme`: the URL scheme (`http`), always set. + - `host`: the URL host (`example.com`). + - `username`: optional username if included (`samuel`). + - `password`: optional password if included (`pass`). + - `port`: optional port (`8000`). + - `path`: optional path (`/the/path/`). + - `query`: optional URL query (for example, `GET` arguments or "search string", such as `query=here`). + - `fragment`: optional fragment (`fragment=is;this=bit`). + """ + + +# Note: all single host urls inherit from `AnyUrl` to preserve compatibility with pre-v2.10 code +# Where urls were annotated variants of `AnyUrl`, which was an alias to `pydantic_core.Url` + + +class AnyHttpUrl(AnyUrl): + """A type that will accept any http or https URL. + + * TLD not required + * Host not required + """ + + _constraints = UrlConstraints(allowed_schemes=['http', 'https']) + + +class HttpUrl(AnyUrl): + """A type that will accept any http or https URL. + + * TLD not required + * Host not required + * Max length 2083 + + ```python + from pydantic import BaseModel, HttpUrl, ValidationError + + class MyModel(BaseModel): + url: HttpUrl + + m = MyModel(url='http://www.example.com') # (1)! + print(m.url) + #> http://www.example.com/ + + try: + MyModel(url='ftp://invalid.url') + except ValidationError as e: + print(e) + ''' + 1 validation error for MyModel + url + URL scheme should be 'http' or 'https' [type=url_scheme, input_value='ftp://invalid.url', input_type=str] + ''' + + try: + MyModel(url='not a url') + except ValidationError as e: + print(e) + ''' + 1 validation error for MyModel + url + Input should be a valid URL, relative URL without a base [type=url_parsing, input_value='not a url', input_type=str] + ''' + ``` + + 1. Note: mypy would prefer `m = MyModel(url=HttpUrl('http://www.example.com'))`, but Pydantic will convert the string to an HttpUrl instance anyway. + + "International domains" (e.g. a URL where the host or TLD includes non-ascii characters) will be encoded via + [punycode](https://en.wikipedia.org/wiki/Punycode) (see + [this article](https://www.xudongz.com/blog/2017/idn-phishing/) for a good description of why this is important): + + ```python + from pydantic import BaseModel, HttpUrl + + class MyModel(BaseModel): + url: HttpUrl + + m1 = MyModel(url='http://puny£code.com') + print(m1.url) + #> http://xn--punycode-eja.com/ + m2 = MyModel(url='https://www.аррӏе.com/') + print(m2.url) + #> https://www.xn--80ak6aa92e.com/ + m3 = MyModel(url='https://www.example.珠宝/') + print(m3.url) + #> https://www.example.xn--pbt977c/ + ``` + + + !!! warning "Underscores in Hostnames" + In Pydantic, underscores are allowed in all parts of a domain except the TLD. + Technically this might be wrong - in theory the hostname cannot have underscores, but subdomains can. + + To explain this; consider the following two cases: + + - `exam_ple.co.uk`: the hostname is `exam_ple`, which should not be allowed since it contains an underscore. + - `foo_bar.example.com` the hostname is `example`, which should be allowed since the underscore is in the subdomain. + + Without having an exhaustive list of TLDs, it would be impossible to differentiate between these two. Therefore + underscores are allowed, but you can always do further validation in a validator if desired. + + Also, Chrome, Firefox, and Safari all currently accept `http://exam_ple.com` as a URL, so we're in good + (or at least big) company. + """ + + _constraints = UrlConstraints(max_length=2083, allowed_schemes=['http', 'https']) + + +class AnyWebsocketUrl(AnyUrl): + """A type that will accept any ws or wss URL. + + * TLD not required + * Host not required + """ + + _constraints = UrlConstraints(allowed_schemes=['ws', 'wss']) + + +class WebsocketUrl(AnyUrl): + """A type that will accept any ws or wss URL. + + * TLD not required + * Host not required + * Max length 2083 + """ + + _constraints = UrlConstraints(max_length=2083, allowed_schemes=['ws', 'wss']) + + +class FileUrl(AnyUrl): + """A type that will accept any file URL. + + * Host not required + """ + + _constraints = UrlConstraints(allowed_schemes=['file']) + + +class FtpUrl(AnyUrl): + """A type that will accept ftp URL. + + * TLD not required + * Host not required + """ + + _constraints = UrlConstraints(allowed_schemes=['ftp']) + + +class PostgresDsn(_BaseMultiHostUrl): + """A type that will accept any Postgres DSN. + + * User info required + * TLD not required + * Host required + * Supports multiple hosts + + If further validation is required, these properties can be used by validators to enforce specific behaviour: + + ```python + from pydantic import ( + BaseModel, + HttpUrl, + PostgresDsn, + ValidationError, + field_validator, + ) + + class MyModel(BaseModel): + url: HttpUrl + + m = MyModel(url='http://www.example.com') + + # the repr() method for a url will display all properties of the url + print(repr(m.url)) + #> HttpUrl('http://www.example.com/') + print(m.url.scheme) + #> http + print(m.url.host) + #> www.example.com + print(m.url.port) + #> 80 + + class MyDatabaseModel(BaseModel): + db: PostgresDsn + + @field_validator('db') + def check_db_name(cls, v): + assert v.path and len(v.path) > 1, 'database must be provided' + return v + + m = MyDatabaseModel(db='postgres://user:pass@localhost:5432/foobar') + print(m.db) + #> postgres://user:pass@localhost:5432/foobar + + try: + MyDatabaseModel(db='postgres://user:pass@localhost:5432') + except ValidationError as e: + print(e) + ''' + 1 validation error for MyDatabaseModel + db + Assertion failed, database must be provided + assert (None) + + where None = PostgresDsn('postgres://user:pass@localhost:5432').path [type=assertion_error, input_value='postgres://user:pass@localhost:5432', input_type=str] + ''' + ``` + """ + + _constraints = UrlConstraints( + host_required=True, + allowed_schemes=[ + 'postgres', + 'postgresql', + 'postgresql+asyncpg', + 'postgresql+pg8000', + 'postgresql+psycopg', + 'postgresql+psycopg2', + 'postgresql+psycopg2cffi', + 'postgresql+py-postgresql', + 'postgresql+pygresql', + ], + ) + + @property + def host(self) -> str: + """The required URL host.""" + return self._url.host # pyright: ignore[reportAttributeAccessIssue] + + +class CockroachDsn(AnyUrl): + """A type that will accept any Cockroach DSN. + + * User info required + * TLD not required + * Host required + """ + + _constraints = UrlConstraints( + host_required=True, + allowed_schemes=[ + 'cockroachdb', + 'cockroachdb+psycopg2', + 'cockroachdb+asyncpg', + ], + ) + + @property + def host(self) -> str: + """The required URL host.""" + return self._url.host # pyright: ignore[reportReturnType] + + +class AmqpDsn(AnyUrl): + """A type that will accept any AMQP DSN. + + * User info required + * TLD not required + * Host not required + """ + + _constraints = UrlConstraints(allowed_schemes=['amqp', 'amqps']) + + +class RedisDsn(AnyUrl): + """A type that will accept any Redis DSN. + + * User info required + * TLD not required + * Host required (e.g., `rediss://:pass@localhost`) + """ + + _constraints = UrlConstraints( + allowed_schemes=['redis', 'rediss'], + default_host='localhost', + default_port=6379, + default_path='/0', + host_required=True, + ) + + @property + def host(self) -> str: + """The required URL host.""" + return self._url.host # pyright: ignore[reportReturnType] + + +class MongoDsn(_BaseMultiHostUrl): + """A type that will accept any MongoDB DSN. + + * User info not required + * Database name not required + * Port not required + * User info may be passed without user part (e.g., `mongodb://mongodb0.example.com:27017`). + """ + + _constraints = UrlConstraints(allowed_schemes=['mongodb', 'mongodb+srv'], default_port=27017) + + +class KafkaDsn(AnyUrl): + """A type that will accept any Kafka DSN. + + * User info required + * TLD not required + * Host not required + """ + + _constraints = UrlConstraints(allowed_schemes=['kafka'], default_host='localhost', default_port=9092) + + +class NatsDsn(_BaseMultiHostUrl): + """A type that will accept any NATS DSN. + + NATS is a connective technology built for the ever increasingly hyper-connected world. + It is a single technology that enables applications to securely communicate across + any combination of cloud vendors, on-premise, edge, web and mobile, and devices. + More: https://nats.io + """ + + _constraints = UrlConstraints( + allowed_schemes=['nats', 'tls', 'ws', 'wss'], default_host='localhost', default_port=4222 + ) + + +class MySQLDsn(AnyUrl): + """A type that will accept any MySQL DSN. + + * User info required + * TLD not required + * Host not required + """ + + _constraints = UrlConstraints( + allowed_schemes=[ + 'mysql', + 'mysql+mysqlconnector', + 'mysql+aiomysql', + 'mysql+asyncmy', + 'mysql+mysqldb', + 'mysql+pymysql', + 'mysql+cymysql', + 'mysql+pyodbc', + ], + default_port=3306, + host_required=True, + ) + + +class MariaDBDsn(AnyUrl): + """A type that will accept any MariaDB DSN. + + * User info required + * TLD not required + * Host not required + """ + + _constraints = UrlConstraints( + allowed_schemes=['mariadb', 'mariadb+mariadbconnector', 'mariadb+pymysql'], + default_port=3306, + ) + + +class ClickHouseDsn(AnyUrl): + """A type that will accept any ClickHouse DSN. + + * User info required + * TLD not required + * Host not required + """ + + _constraints = UrlConstraints( + allowed_schemes=[ + 'clickhouse+native', + 'clickhouse+asynch', + 'clickhouse+http', + 'clickhouse', + 'clickhouses', + 'clickhousedb', + ], + default_host='localhost', + default_port=9000, + ) + + +class SnowflakeDsn(AnyUrl): + """A type that will accept any Snowflake DSN. + + * User info required + * TLD not required + * Host required + """ + + _constraints = UrlConstraints( + allowed_schemes=['snowflake'], + host_required=True, + ) + + @property + def host(self) -> str: + """The required URL host.""" + return self._url.host # pyright: ignore[reportReturnType] + + +def import_email_validator() -> None: + global email_validator + try: + import email_validator + except ImportError as e: + raise ImportError('email-validator is not installed, run `pip install pydantic[email]`') from e + if not version('email-validator').partition('.')[0] == '2': + raise ImportError('email-validator version >= 2.0 required, run pip install -U email-validator') + + +if TYPE_CHECKING: + EmailStr = Annotated[str, ...] +else: + + class EmailStr: + """ + Info: + To use this type, you need to install the optional + [`email-validator`](https://github.com/JoshData/python-email-validator) package: + + ```bash + pip install email-validator + ``` + + Validate email addresses. + + ```python + from pydantic import BaseModel, EmailStr + + class Model(BaseModel): + email: EmailStr + + print(Model(email='contact@mail.com')) + #> email='contact@mail.com' + ``` + """ # noqa: D212 + + @classmethod + def __get_pydantic_core_schema__( + cls, + _source: type[Any], + _handler: GetCoreSchemaHandler, + ) -> core_schema.CoreSchema: + import_email_validator() + return core_schema.no_info_after_validator_function(cls._validate, core_schema.str_schema()) + + @classmethod + def __get_pydantic_json_schema__( + cls, core_schema: core_schema.CoreSchema, handler: _schema_generation_shared.GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = handler(core_schema) + field_schema.update(type='string', format='email') + return field_schema + + @classmethod + def _validate(cls, input_value: str, /) -> str: + return validate_email(input_value)[1] + + +class NameEmail(_repr.Representation): + """ + Info: + To use this type, you need to install the optional + [`email-validator`](https://github.com/JoshData/python-email-validator) package: + + ```bash + pip install email-validator + ``` + + Validate a name and email address combination, as specified by + [RFC 5322](https://datatracker.ietf.org/doc/html/rfc5322#section-3.4). + + The `NameEmail` has two properties: `name` and `email`. + In case the `name` is not provided, it's inferred from the email address. + + ```python + from pydantic import BaseModel, NameEmail + + class User(BaseModel): + email: NameEmail + + user = User(email='Fred Bloggs ') + print(user.email) + #> Fred Bloggs + print(user.email.name) + #> Fred Bloggs + + user = User(email='fred.bloggs@example.com') + print(user.email) + #> fred.bloggs + print(user.email.name) + #> fred.bloggs + ``` + """ # noqa: D212 + + __slots__ = 'name', 'email' + + def __init__(self, name: str, email: str): + self.name = name + self.email = email + + def __eq__(self, other: Any) -> bool: + return isinstance(other, NameEmail) and (self.name, self.email) == (other.name, other.email) + + @classmethod + def __get_pydantic_json_schema__( + cls, core_schema: core_schema.CoreSchema, handler: _schema_generation_shared.GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = handler(core_schema) + field_schema.update(type='string', format='name-email') + return field_schema + + @classmethod + def __get_pydantic_core_schema__( + cls, + _source: type[Any], + _handler: GetCoreSchemaHandler, + ) -> core_schema.CoreSchema: + import_email_validator() + + return core_schema.no_info_after_validator_function( + cls._validate, + core_schema.json_or_python_schema( + json_schema=core_schema.str_schema(), + python_schema=core_schema.union_schema( + [core_schema.is_instance_schema(cls), core_schema.str_schema()], + custom_error_type='name_email_type', + custom_error_message='Input is not a valid NameEmail', + ), + serialization=core_schema.to_string_ser_schema(), + ), + ) + + @classmethod + def _validate(cls, input_value: Self | str, /) -> Self: + if isinstance(input_value, str): + name, email = validate_email(input_value) + return cls(name, email) + else: + return input_value + + def __str__(self) -> str: + if '@' in self.name: + return f'"{self.name}" <{self.email}>' + + return f'{self.name} <{self.email}>' + + +IPvAnyAddressType: TypeAlias = 'IPv4Address | IPv6Address' +IPvAnyInterfaceType: TypeAlias = 'IPv4Interface | IPv6Interface' +IPvAnyNetworkType: TypeAlias = 'IPv4Network | IPv6Network' + +if TYPE_CHECKING: + IPvAnyAddress = IPvAnyAddressType + IPvAnyInterface = IPvAnyInterfaceType + IPvAnyNetwork = IPvAnyNetworkType +else: + + class IPvAnyAddress: + """Validate an IPv4 or IPv6 address. + + ```python + from pydantic import BaseModel + from pydantic.networks import IPvAnyAddress + + class IpModel(BaseModel): + ip: IPvAnyAddress + + print(IpModel(ip='127.0.0.1')) + #> ip=IPv4Address('127.0.0.1') + + try: + IpModel(ip='http://www.example.com') + except ValueError as e: + print(e.errors()) + ''' + [ + { + 'type': 'ip_any_address', + 'loc': ('ip',), + 'msg': 'value is not a valid IPv4 or IPv6 address', + 'input': 'http://www.example.com', + } + ] + ''' + ``` + """ + + __slots__ = () + + def __new__(cls, value: Any) -> IPvAnyAddressType: + """Validate an IPv4 or IPv6 address.""" + try: + return IPv4Address(value) + except ValueError: + pass + + try: + return IPv6Address(value) + except ValueError: + raise PydanticCustomError('ip_any_address', 'value is not a valid IPv4 or IPv6 address') + + @classmethod + def __get_pydantic_json_schema__( + cls, core_schema: core_schema.CoreSchema, handler: _schema_generation_shared.GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = {} + field_schema.update(type='string', format='ipvanyaddress') + return field_schema + + @classmethod + def __get_pydantic_core_schema__( + cls, + _source: type[Any], + _handler: GetCoreSchemaHandler, + ) -> core_schema.CoreSchema: + return core_schema.no_info_plain_validator_function( + cls._validate, serialization=core_schema.to_string_ser_schema() + ) + + @classmethod + def _validate(cls, input_value: Any, /) -> IPvAnyAddressType: + return cls(input_value) # type: ignore[return-value] + + class IPvAnyInterface: + """Validate an IPv4 or IPv6 interface.""" + + __slots__ = () + + def __new__(cls, value: NetworkType) -> IPvAnyInterfaceType: + """Validate an IPv4 or IPv6 interface.""" + try: + return IPv4Interface(value) + except ValueError: + pass + + try: + return IPv6Interface(value) + except ValueError: + raise PydanticCustomError('ip_any_interface', 'value is not a valid IPv4 or IPv6 interface') + + @classmethod + def __get_pydantic_json_schema__( + cls, core_schema: core_schema.CoreSchema, handler: _schema_generation_shared.GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = {} + field_schema.update(type='string', format='ipvanyinterface') + return field_schema + + @classmethod + def __get_pydantic_core_schema__( + cls, + _source: type[Any], + _handler: GetCoreSchemaHandler, + ) -> core_schema.CoreSchema: + return core_schema.no_info_plain_validator_function( + cls._validate, serialization=core_schema.to_string_ser_schema() + ) + + @classmethod + def _validate(cls, input_value: NetworkType, /) -> IPvAnyInterfaceType: + return cls(input_value) # type: ignore[return-value] + + class IPvAnyNetwork: + """Validate an IPv4 or IPv6 network.""" + + __slots__ = () + + def __new__(cls, value: NetworkType) -> IPvAnyNetworkType: + """Validate an IPv4 or IPv6 network.""" + # Assume IP Network is defined with a default value for `strict` argument. + # Define your own class if you want to specify network address check strictness. + try: + return IPv4Network(value) + except ValueError: + pass + + try: + return IPv6Network(value) + except ValueError: + raise PydanticCustomError('ip_any_network', 'value is not a valid IPv4 or IPv6 network') + + @classmethod + def __get_pydantic_json_schema__( + cls, core_schema: core_schema.CoreSchema, handler: _schema_generation_shared.GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = {} + field_schema.update(type='string', format='ipvanynetwork') + return field_schema + + @classmethod + def __get_pydantic_core_schema__( + cls, + _source: type[Any], + _handler: GetCoreSchemaHandler, + ) -> core_schema.CoreSchema: + return core_schema.no_info_plain_validator_function( + cls._validate, serialization=core_schema.to_string_ser_schema() + ) + + @classmethod + def _validate(cls, input_value: NetworkType, /) -> IPvAnyNetworkType: + return cls(input_value) # type: ignore[return-value] + + +def _build_pretty_email_regex() -> re.Pattern[str]: + name_chars = r'[\w!#$%&\'*+\-/=?^_`{|}~]' + unquoted_name_group = rf'((?:{name_chars}+\s+)*{name_chars}+)' + quoted_name_group = r'"((?:[^"]|\")+)"' + email_group = r'<(.+)>' + return re.compile(rf'\s*(?:{unquoted_name_group}|{quoted_name_group})?\s*{email_group}\s*') + + +pretty_email_regex = _build_pretty_email_regex() + +MAX_EMAIL_LENGTH = 2048 +"""Maximum length for an email. +A somewhat arbitrary but very generous number compared to what is allowed by most implementations. +""" + + +def validate_email(value: str) -> tuple[str, str]: + """Email address validation using [email-validator](https://pypi.org/project/email-validator/). + + Returns: + A tuple containing the local part of the email (or the name for "pretty" email addresses) + and the normalized email. + + Raises: + PydanticCustomError: If the email is invalid. + + Note: + Note that: + + * Raw IP address (literal) domain parts are not allowed. + * `"John Doe "` style "pretty" email addresses are processed. + * Spaces are striped from the beginning and end of addresses, but no error is raised. + """ + if email_validator is None: + import_email_validator() + + if len(value) > MAX_EMAIL_LENGTH: + raise PydanticCustomError( + 'value_error', + 'value is not a valid email address: {reason}', + {'reason': f'Length must not exceed {MAX_EMAIL_LENGTH} characters'}, + ) + + m = pretty_email_regex.fullmatch(value) + name: str | None = None + if m: + unquoted_name, quoted_name, value = m.groups() + name = unquoted_name or quoted_name + + email = value.strip() + + try: + parts = email_validator.validate_email(email, check_deliverability=False) + except email_validator.EmailNotValidError as e: + raise PydanticCustomError( + 'value_error', 'value is not a valid email address: {reason}', {'reason': str(e.args[0])} + ) from e + + email = parts.normalized + assert email is not None + name = name or parts.local_part + return name, email + + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/parse.py b/lib/python3.12/site-packages/pydantic/parse.py new file mode 100644 index 0000000000000000000000000000000000000000..68b7f0464ece32edfb955eec0cb24655752716d6 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/parse.py @@ -0,0 +1,5 @@ +"""The `parse` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/plugin/__init__.py b/lib/python3.12/site-packages/pydantic/plugin/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..362030584fcc81bbbdb9be82156ab0adb1119609 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/plugin/__init__.py @@ -0,0 +1,188 @@ +"""!!! abstract "Usage Documentation" + [Build a Plugin](../concepts/plugins.md#build-a-plugin) + +Plugin interface for Pydantic plugins, and related types. +""" + +from __future__ import annotations + +from typing import Any, Callable, Literal, NamedTuple + +from pydantic_core import CoreConfig, CoreSchema, ValidationError +from typing_extensions import Protocol, TypeAlias + +__all__ = ( + 'PydanticPluginProtocol', + 'BaseValidateHandlerProtocol', + 'ValidatePythonHandlerProtocol', + 'ValidateJsonHandlerProtocol', + 'ValidateStringsHandlerProtocol', + 'NewSchemaReturns', + 'SchemaTypePath', + 'SchemaKind', +) + +NewSchemaReturns: TypeAlias = 'tuple[ValidatePythonHandlerProtocol | None, ValidateJsonHandlerProtocol | None, ValidateStringsHandlerProtocol | None]' + + +class SchemaTypePath(NamedTuple): + """Path defining where `schema_type` was defined, or where `TypeAdapter` was called.""" + + module: str + name: str + + +SchemaKind: TypeAlias = Literal['BaseModel', 'TypeAdapter', 'dataclass', 'create_model', 'validate_call'] + + +class PydanticPluginProtocol(Protocol): + """Protocol defining the interface for Pydantic plugins.""" + + def new_schema_validator( + self, + schema: CoreSchema, + schema_type: Any, + schema_type_path: SchemaTypePath, + schema_kind: SchemaKind, + config: CoreConfig | None, + plugin_settings: dict[str, object], + ) -> tuple[ + ValidatePythonHandlerProtocol | None, ValidateJsonHandlerProtocol | None, ValidateStringsHandlerProtocol | None + ]: + """This method is called for each plugin every time a new [`SchemaValidator`][pydantic_core.SchemaValidator] + is created. + + It should return an event handler for each of the three validation methods, or `None` if the plugin does not + implement that method. + + Args: + schema: The schema to validate against. + schema_type: The original type which the schema was created from, e.g. the model class. + schema_type_path: Path defining where `schema_type` was defined, or where `TypeAdapter` was called. + schema_kind: The kind of schema to validate against. + config: The config to use for validation. + plugin_settings: Any plugin settings. + + Returns: + A tuple of optional event handlers for each of the three validation methods - + `validate_python`, `validate_json`, `validate_strings`. + """ + raise NotImplementedError('Pydantic plugins should implement `new_schema_validator`.') + + +class BaseValidateHandlerProtocol(Protocol): + """Base class for plugin callbacks protocols. + + You shouldn't implement this protocol directly, instead use one of the subclasses with adds the correctly + typed `on_error` method. + """ + + on_enter: Callable[..., None] + """`on_enter` is changed to be more specific on all subclasses""" + + def on_success(self, result: Any) -> None: + """Callback to be notified of successful validation. + + Args: + result: The result of the validation. + """ + return + + def on_error(self, error: ValidationError) -> None: + """Callback to be notified of validation errors. + + Args: + error: The validation error. + """ + return + + def on_exception(self, exception: Exception) -> None: + """Callback to be notified of validation exceptions. + + Args: + exception: The exception raised during validation. + """ + return + + +class ValidatePythonHandlerProtocol(BaseValidateHandlerProtocol, Protocol): + """Event handler for `SchemaValidator.validate_python`.""" + + def on_enter( + self, + input: Any, + *, + strict: bool | None = None, + from_attributes: bool | None = None, + context: dict[str, Any] | None = None, + self_instance: Any | None = None, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> None: + """Callback to be notified of validation start, and create an instance of the event handler. + + Args: + input: The input to be validated. + strict: Whether to validate the object in strict mode. + from_attributes: Whether to validate objects as inputs by extracting attributes. + context: The context to use for validation, this is passed to functional validators. + self_instance: An instance of a model to set attributes on from validation, this is used when running + validation from the `__init__` method of a model. + by_alias: Whether to use the field's alias to match the input data to an attribute. + by_name: Whether to use the field's name to match the input data to an attribute. + """ + pass + + +class ValidateJsonHandlerProtocol(BaseValidateHandlerProtocol, Protocol): + """Event handler for `SchemaValidator.validate_json`.""" + + def on_enter( + self, + input: str | bytes | bytearray, + *, + strict: bool | None = None, + context: dict[str, Any] | None = None, + self_instance: Any | None = None, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> None: + """Callback to be notified of validation start, and create an instance of the event handler. + + Args: + input: The JSON data to be validated. + strict: Whether to validate the object in strict mode. + context: The context to use for validation, this is passed to functional validators. + self_instance: An instance of a model to set attributes on from validation, this is used when running + validation from the `__init__` method of a model. + by_alias: Whether to use the field's alias to match the input data to an attribute. + by_name: Whether to use the field's name to match the input data to an attribute. + """ + pass + + +StringInput: TypeAlias = 'dict[str, StringInput]' + + +class ValidateStringsHandlerProtocol(BaseValidateHandlerProtocol, Protocol): + """Event handler for `SchemaValidator.validate_strings`.""" + + def on_enter( + self, + input: StringInput, + *, + strict: bool | None = None, + context: dict[str, Any] | None = None, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> None: + """Callback to be notified of validation start, and create an instance of the event handler. + + Args: + input: The string data to be validated. + strict: Whether to validate the object in strict mode. + context: The context to use for validation, this is passed to functional validators. + by_alias: Whether to use the field's alias to match the input data to an attribute. + by_name: Whether to use the field's name to match the input data to an attribute. + """ + pass diff --git a/lib/python3.12/site-packages/pydantic/plugin/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/plugin/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d8812a4466ffe78873f67a1782f7a0d6cb71ba5 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/plugin/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/plugin/__pycache__/_loader.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/plugin/__pycache__/_loader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..61602df9c65580f43d09f7d0dcbd24864a8c117f Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/plugin/__pycache__/_loader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/plugin/__pycache__/_schema_validator.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/plugin/__pycache__/_schema_validator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fc30388573cec1566d1309bc4a311cc61db7082 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/plugin/__pycache__/_schema_validator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/plugin/_loader.py b/lib/python3.12/site-packages/pydantic/plugin/_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..7d1f0f29cba8e282e19eb031702251a489b72a7d --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/plugin/_loader.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import importlib.metadata as importlib_metadata +import os +import warnings +from collections.abc import Iterable +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + from . import PydanticPluginProtocol + + +PYDANTIC_ENTRY_POINT_GROUP: Final[str] = 'pydantic' + +# cache of plugins +_plugins: dict[str, PydanticPluginProtocol] | None = None +# return no plugins while loading plugins to avoid recursion and errors while import plugins +# this means that if plugins use pydantic +_loading_plugins: bool = False + + +def get_plugins() -> Iterable[PydanticPluginProtocol]: + """Load plugins for Pydantic. + + Inspired by: https://github.com/pytest-dev/pluggy/blob/1.3.0/src/pluggy/_manager.py#L376-L402 + """ + disabled_plugins = os.getenv('PYDANTIC_DISABLE_PLUGINS') + global _plugins, _loading_plugins + if _loading_plugins: + # this happens when plugins themselves use pydantic, we return no plugins + return () + elif disabled_plugins in ('__all__', '1', 'true'): + return () + elif _plugins is None: + _plugins = {} + # set _loading_plugins so any plugins that use pydantic don't themselves use plugins + _loading_plugins = True + try: + for dist in importlib_metadata.distributions(): + for entry_point in dist.entry_points: + if entry_point.group != PYDANTIC_ENTRY_POINT_GROUP: + continue + if entry_point.value in _plugins: + continue + if disabled_plugins is not None and entry_point.name in disabled_plugins.split(','): + continue + try: + _plugins[entry_point.value] = entry_point.load() + except (ImportError, AttributeError) as e: + warnings.warn( + f'{e.__class__.__name__} while loading the `{entry_point.name}` Pydantic plugin, ' + f'this plugin will not be installed.\n\n{e!r}' + ) + finally: + _loading_plugins = False + + return _plugins.values() diff --git a/lib/python3.12/site-packages/pydantic/plugin/_schema_validator.py b/lib/python3.12/site-packages/pydantic/plugin/_schema_validator.py new file mode 100644 index 0000000000000000000000000000000000000000..83f2562b1b117fbeda087edadeb97b314c8b70ce --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/plugin/_schema_validator.py @@ -0,0 +1,140 @@ +"""Pluggable schema validator for pydantic.""" + +from __future__ import annotations + +import functools +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar + +from pydantic_core import CoreConfig, CoreSchema, SchemaValidator, ValidationError +from typing_extensions import ParamSpec + +if TYPE_CHECKING: + from . import BaseValidateHandlerProtocol, PydanticPluginProtocol, SchemaKind, SchemaTypePath + + +P = ParamSpec('P') +R = TypeVar('R') +Event = Literal['on_validate_python', 'on_validate_json', 'on_validate_strings'] +events: list[Event] = list(Event.__args__) # type: ignore + + +def create_schema_validator( + schema: CoreSchema, + schema_type: Any, + schema_type_module: str, + schema_type_name: str, + schema_kind: SchemaKind, + config: CoreConfig | None = None, + plugin_settings: dict[str, Any] | None = None, +) -> SchemaValidator | PluggableSchemaValidator: + """Create a `SchemaValidator` or `PluggableSchemaValidator` if plugins are installed. + + Returns: + If plugins are installed then return `PluggableSchemaValidator`, otherwise return `SchemaValidator`. + """ + from . import SchemaTypePath + from ._loader import get_plugins + + plugins = get_plugins() + if plugins: + return PluggableSchemaValidator( + schema, + schema_type, + SchemaTypePath(schema_type_module, schema_type_name), + schema_kind, + config, + plugins, + plugin_settings or {}, + ) + else: + return SchemaValidator(schema, config) + + +class PluggableSchemaValidator: + """Pluggable schema validator.""" + + __slots__ = '_schema_validator', 'validate_json', 'validate_python', 'validate_strings' + + def __init__( + self, + schema: CoreSchema, + schema_type: Any, + schema_type_path: SchemaTypePath, + schema_kind: SchemaKind, + config: CoreConfig | None, + plugins: Iterable[PydanticPluginProtocol], + plugin_settings: dict[str, Any], + ) -> None: + self._schema_validator = SchemaValidator(schema, config) + + python_event_handlers: list[BaseValidateHandlerProtocol] = [] + json_event_handlers: list[BaseValidateHandlerProtocol] = [] + strings_event_handlers: list[BaseValidateHandlerProtocol] = [] + for plugin in plugins: + try: + p, j, s = plugin.new_schema_validator( + schema, schema_type, schema_type_path, schema_kind, config, plugin_settings + ) + except TypeError as e: # pragma: no cover + raise TypeError(f'Error using plugin `{plugin.__module__}:{plugin.__class__.__name__}`: {e}') from e + if p is not None: + python_event_handlers.append(p) + if j is not None: + json_event_handlers.append(j) + if s is not None: + strings_event_handlers.append(s) + + self.validate_python = build_wrapper(self._schema_validator.validate_python, python_event_handlers) + self.validate_json = build_wrapper(self._schema_validator.validate_json, json_event_handlers) + self.validate_strings = build_wrapper(self._schema_validator.validate_strings, strings_event_handlers) + + def __getattr__(self, name: str) -> Any: + return getattr(self._schema_validator, name) + + +def build_wrapper(func: Callable[P, R], event_handlers: list[BaseValidateHandlerProtocol]) -> Callable[P, R]: + if not event_handlers: + return func + else: + on_enters = tuple(h.on_enter for h in event_handlers if filter_handlers(h, 'on_enter')) + on_successes = tuple(h.on_success for h in event_handlers if filter_handlers(h, 'on_success')) + on_errors = tuple(h.on_error for h in event_handlers if filter_handlers(h, 'on_error')) + on_exceptions = tuple(h.on_exception for h in event_handlers if filter_handlers(h, 'on_exception')) + + @functools.wraps(func) + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: + for on_enter_handler in on_enters: + on_enter_handler(*args, **kwargs) + + try: + result = func(*args, **kwargs) + except ValidationError as error: + for on_error_handler in on_errors: + on_error_handler(error) + raise + except Exception as exception: + for on_exception_handler in on_exceptions: + on_exception_handler(exception) + raise + else: + for on_success_handler in on_successes: + on_success_handler(result) + return result + + return wrapper + + +def filter_handlers(handler_cls: BaseValidateHandlerProtocol, method_name: str) -> bool: + """Filter out handler methods which are not implemented by the plugin directly - e.g. are missing + or are inherited from the protocol. + """ + handler = getattr(handler_cls, method_name, None) + if handler is None: + return False + elif handler.__module__ == 'pydantic.plugin': + # this is the original handler, from the protocol due to runtime inheritance + # we don't want to call it + return False + else: + return True diff --git a/lib/python3.12/site-packages/pydantic/py.typed b/lib/python3.12/site-packages/pydantic/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/pydantic/root_model.py b/lib/python3.12/site-packages/pydantic/root_model.py new file mode 100644 index 0000000000000000000000000000000000000000..8b3ff014edce1fb5974ee70ebb23aa6460f62bf7 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/root_model.py @@ -0,0 +1,157 @@ +"""RootModel class and type definitions.""" + +from __future__ import annotations as _annotations + +import typing +from copy import copy, deepcopy + +from pydantic_core import PydanticUndefined + +from . import PydanticUserError +from ._internal import _model_construction, _repr +from .main import BaseModel, _object_setattr + +if typing.TYPE_CHECKING: + from typing import Any, Literal + + from typing_extensions import Self, dataclass_transform + + from .fields import Field as PydanticModelField + from .fields import PrivateAttr as PydanticModelPrivateAttr + + # dataclass_transform could be applied to RootModel directly, but `ModelMetaclass`'s dataclass_transform + # takes priority (at least with pyright). We trick type checkers into thinking we apply dataclass_transform + # on a new metaclass. + @dataclass_transform(kw_only_default=False, field_specifiers=(PydanticModelField, PydanticModelPrivateAttr)) + class _RootModelMetaclass(_model_construction.ModelMetaclass): ... +else: + _RootModelMetaclass = _model_construction.ModelMetaclass + +__all__ = ('RootModel',) + +RootModelRootType = typing.TypeVar('RootModelRootType') + + +class RootModel(BaseModel, typing.Generic[RootModelRootType], metaclass=_RootModelMetaclass): + """!!! abstract "Usage Documentation" + [`RootModel` and Custom Root Types](../concepts/models.md#rootmodel-and-custom-root-types) + + A Pydantic `BaseModel` for the root object of the model. + + Attributes: + root: The root object of the model. + __pydantic_root_model__: Whether the model is a RootModel. + __pydantic_private__: Private fields in the model. + __pydantic_extra__: Extra fields in the model. + + """ + + __pydantic_root_model__ = True + __pydantic_private__ = None + __pydantic_extra__ = None + + root: RootModelRootType + + def __init_subclass__(cls, **kwargs): + extra = cls.model_config.get('extra') + if extra is not None: + raise PydanticUserError( + "`RootModel` does not support setting `model_config['extra']`", code='root-model-extra' + ) + super().__init_subclass__(**kwargs) + + def __init__(self, /, root: RootModelRootType = PydanticUndefined, **data) -> None: # type: ignore + __tracebackhide__ = True + if data: + if root is not PydanticUndefined: + raise ValueError( + '"RootModel.__init__" accepts either a single positional argument or arbitrary keyword arguments' + ) + root = data # type: ignore + self.__pydantic_validator__.validate_python(root, self_instance=self) + + __init__.__pydantic_base_init__ = True # pyright: ignore[reportFunctionMemberAccess] + + @classmethod + def model_construct(cls, root: RootModelRootType, _fields_set: set[str] | None = None) -> Self: # type: ignore + """Create a new model using the provided root object and update fields set. + + Args: + root: The root object of the model. + _fields_set: The set of fields to be updated. + + Returns: + The new model. + + Raises: + NotImplemented: If the model is not a subclass of `RootModel`. + """ + return super().model_construct(root=root, _fields_set=_fields_set) + + def __getstate__(self) -> dict[Any, Any]: + return { + '__dict__': self.__dict__, + '__pydantic_fields_set__': self.__pydantic_fields_set__, + } + + def __setstate__(self, state: dict[Any, Any]) -> None: + _object_setattr(self, '__pydantic_fields_set__', state['__pydantic_fields_set__']) + _object_setattr(self, '__dict__', state['__dict__']) + + def __copy__(self) -> Self: + """Returns a shallow copy of the model.""" + cls = type(self) + m = cls.__new__(cls) + _object_setattr(m, '__dict__', copy(self.__dict__)) + _object_setattr(m, '__pydantic_fields_set__', copy(self.__pydantic_fields_set__)) + return m + + def __deepcopy__(self, memo: dict[int, Any] | None = None) -> Self: + """Returns a deep copy of the model.""" + cls = type(self) + m = cls.__new__(cls) + _object_setattr(m, '__dict__', deepcopy(self.__dict__, memo=memo)) + # This next line doesn't need a deepcopy because __pydantic_fields_set__ is a set[str], + # and attempting a deepcopy would be marginally slower. + _object_setattr(m, '__pydantic_fields_set__', copy(self.__pydantic_fields_set__)) + return m + + if typing.TYPE_CHECKING: + + def model_dump( # type: ignore + self, + *, + mode: Literal['json', 'python'] | str = 'python', + include: Any = None, + exclude: Any = None, + context: dict[str, Any] | None = None, + by_alias: bool | None = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + round_trip: bool = False, + warnings: bool | Literal['none', 'warn', 'error'] = True, + serialize_as_any: bool = False, + ) -> Any: + """This method is included just to get a more accurate return type for type checkers. + It is included in this `if TYPE_CHECKING:` block since no override is actually necessary. + + See the documentation of `BaseModel.model_dump` for more details about the arguments. + + Generally, this method will have a return type of `RootModelRootType`, assuming that `RootModelRootType` is + not a `BaseModel` subclass. If `RootModelRootType` is a `BaseModel` subclass, then the return + type will likely be `dict[str, Any]`, as `model_dump` calls are recursive. The return type could + even be something different, in the case of a custom serializer. + Thus, `Any` is used here to catch all of these cases. + """ + ... + + def __eq__(self, other: Any) -> bool: + if not isinstance(other, RootModel): + return NotImplemented + return self.__pydantic_fields__['root'].annotation == other.__pydantic_fields__[ + 'root' + ].annotation and super().__eq__(other) + + def __repr_args__(self) -> _repr.ReprArgs: + yield 'root', self.root diff --git a/lib/python3.12/site-packages/pydantic/schema.py b/lib/python3.12/site-packages/pydantic/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..a3245a61afd0f59143f38ee809ee2d784d9198eb --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/schema.py @@ -0,0 +1,5 @@ +"""The `schema` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/tools.py b/lib/python3.12/site-packages/pydantic/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..fdc68c4f4a35e8d8200b3aafcda63d43839da375 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/tools.py @@ -0,0 +1,5 @@ +"""The `tools` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/type_adapter.py b/lib/python3.12/site-packages/pydantic/type_adapter.py new file mode 100644 index 0000000000000000000000000000000000000000..a6cdaba84b0c4312da26a4202757cf7f307827bc --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/type_adapter.py @@ -0,0 +1,727 @@ +"""Type adapter specification.""" + +from __future__ import annotations as _annotations + +import sys +from collections.abc import Callable, Iterable +from dataclasses import is_dataclass +from types import FrameType +from typing import ( + Any, + Generic, + Literal, + TypeVar, + cast, + final, + overload, +) + +from pydantic_core import CoreSchema, SchemaSerializer, SchemaValidator, Some +from typing_extensions import ParamSpec, is_typeddict + +from pydantic.errors import PydanticUserError +from pydantic.main import BaseModel, IncEx + +from ._internal import _config, _generate_schema, _mock_val_ser, _namespace_utils, _repr, _typing_extra, _utils +from .config import ConfigDict +from .errors import PydanticUndefinedAnnotation +from .json_schema import ( + DEFAULT_REF_TEMPLATE, + GenerateJsonSchema, + JsonSchemaKeyT, + JsonSchemaMode, + JsonSchemaValue, +) +from .plugin._schema_validator import PluggableSchemaValidator, create_schema_validator + +T = TypeVar('T') +R = TypeVar('R') +P = ParamSpec('P') +TypeAdapterT = TypeVar('TypeAdapterT', bound='TypeAdapter') + + +def _getattr_no_parents(obj: Any, attribute: str) -> Any: + """Returns the attribute value without attempting to look up attributes from parent types.""" + if hasattr(obj, '__dict__'): + try: + return obj.__dict__[attribute] + except KeyError: + pass + + slots = getattr(obj, '__slots__', None) + if slots is not None and attribute in slots: + return getattr(obj, attribute) + else: + raise AttributeError(attribute) + + +def _type_has_config(type_: Any) -> bool: + """Returns whether the type has config.""" + type_ = _typing_extra.annotated_type(type_) or type_ + try: + return issubclass(type_, BaseModel) or is_dataclass(type_) or is_typeddict(type_) + except TypeError: + # type is not a class + return False + + +@final +class TypeAdapter(Generic[T]): + """!!! abstract "Usage Documentation" + [`TypeAdapter`](../concepts/type_adapter.md) + + Type adapters provide a flexible way to perform validation and serialization based on a Python type. + + A `TypeAdapter` instance exposes some of the functionality from `BaseModel` instance methods + for types that do not have such methods (such as dataclasses, primitive types, and more). + + **Note:** `TypeAdapter` instances are not types, and cannot be used as type annotations for fields. + + Args: + type: The type associated with the `TypeAdapter`. + config: Configuration for the `TypeAdapter`, should be a dictionary conforming to + [`ConfigDict`][pydantic.config.ConfigDict]. + + !!! note + You cannot provide a configuration when instantiating a `TypeAdapter` if the type you're using + has its own config that cannot be overridden (ex: `BaseModel`, `TypedDict`, and `dataclass`). A + [`type-adapter-config-unused`](../errors/usage_errors.md#type-adapter-config-unused) error will + be raised in this case. + _parent_depth: Depth at which to search for the [parent frame][frame-objects]. This frame is used when + resolving forward annotations during schema building, by looking for the globals and locals of this + frame. Defaults to 2, which will result in the frame where the `TypeAdapter` was instantiated. + + !!! note + This parameter is named with an underscore to suggest its private nature and discourage use. + It may be deprecated in a minor version, so we only recommend using it if you're comfortable + with potential change in behavior/support. It's default value is 2 because internally, + the `TypeAdapter` class makes another call to fetch the frame. + module: The module that passes to plugin if provided. + + Attributes: + core_schema: The core schema for the type. + validator: The schema validator for the type. + serializer: The schema serializer for the type. + pydantic_complete: Whether the core schema for the type is successfully built. + + ??? tip "Compatibility with `mypy`" + Depending on the type used, `mypy` might raise an error when instantiating a `TypeAdapter`. As a workaround, you can explicitly + annotate your variable: + + ```py + from typing import Union + + from pydantic import TypeAdapter + + ta: TypeAdapter[Union[str, int]] = TypeAdapter(Union[str, int]) # type: ignore[arg-type] + ``` + + ??? info "Namespace management nuances and implementation details" + + Here, we collect some notes on namespace management, and subtle differences from `BaseModel`: + + `BaseModel` uses its own `__module__` to find out where it was defined + and then looks for symbols to resolve forward references in those globals. + On the other hand, `TypeAdapter` can be initialized with arbitrary objects, + which may not be types and thus do not have a `__module__` available. + So instead we look at the globals in our parent stack frame. + + It is expected that the `ns_resolver` passed to this function will have the correct + namespace for the type we're adapting. See the source code for `TypeAdapter.__init__` + and `TypeAdapter.rebuild` for various ways to construct this namespace. + + This works for the case where this function is called in a module that + has the target of forward references in its scope, but + does not always work for more complex cases. + + For example, take the following: + + ```python {title="a.py"} + IntList = list[int] + OuterDict = dict[str, 'IntList'] + ``` + + ```python {test="skip" title="b.py"} + from a import OuterDict + + from pydantic import TypeAdapter + + IntList = int # replaces the symbol the forward reference is looking for + v = TypeAdapter(OuterDict) + v({'x': 1}) # should fail but doesn't + ``` + + If `OuterDict` were a `BaseModel`, this would work because it would resolve + the forward reference within the `a.py` namespace. + But `TypeAdapter(OuterDict)` can't determine what module `OuterDict` came from. + + In other words, the assumption that _all_ forward references exist in the + module we are being called from is not technically always true. + Although most of the time it is and it works fine for recursive models and such, + `BaseModel`'s behavior isn't perfect either and _can_ break in similar ways, + so there is no right or wrong between the two. + + But at the very least this behavior is _subtly_ different from `BaseModel`'s. + """ + + core_schema: CoreSchema + validator: SchemaValidator | PluggableSchemaValidator + serializer: SchemaSerializer + pydantic_complete: bool + + @overload + def __init__( + self, + type: type[T], + *, + config: ConfigDict | None = ..., + _parent_depth: int = ..., + module: str | None = ..., + ) -> None: ... + + # This second overload is for unsupported special forms (such as Annotated, Union, etc.) + # Currently there is no way to type this correctly + # See https://github.com/python/typing/pull/1618 + @overload + def __init__( + self, + type: Any, + *, + config: ConfigDict | None = ..., + _parent_depth: int = ..., + module: str | None = ..., + ) -> None: ... + + def __init__( + self, + type: Any, + *, + config: ConfigDict | None = None, + _parent_depth: int = 2, + module: str | None = None, + ) -> None: + if _type_has_config(type) and config is not None: + raise PydanticUserError( + 'Cannot use `config` when the type is a BaseModel, dataclass or TypedDict.' + ' These types can have their own config and setting the config via the `config`' + ' parameter to TypeAdapter will not override it, thus the `config` you passed to' + ' TypeAdapter becomes meaningless, which is probably not what you want.', + code='type-adapter-config-unused', + ) + + self._type = type + self._config = config + self._parent_depth = _parent_depth + self.pydantic_complete = False + + parent_frame = self._fetch_parent_frame() + if parent_frame is not None: + globalns = parent_frame.f_globals + # Do not provide a local ns if the type adapter happens to be instantiated at the module level: + localns = parent_frame.f_locals if parent_frame.f_locals is not globalns else {} + else: + globalns = {} + localns = {} + + self._module_name = module or cast(str, globalns.get('__name__', '')) + self._init_core_attrs( + ns_resolver=_namespace_utils.NsResolver( + namespaces_tuple=_namespace_utils.NamespacesTuple(locals=localns, globals=globalns), + parent_namespace=localns, + ), + force=False, + ) + + def _fetch_parent_frame(self) -> FrameType | None: + frame = sys._getframe(self._parent_depth) + if frame.f_globals.get('__name__') == 'typing': + # Because `TypeAdapter` is generic, explicitly parametrizing the class results + # in a `typing._GenericAlias` instance, which proxies instantiation calls to the + # "real" `TypeAdapter` class and thus adding an extra frame to the call. To avoid + # pulling anything from the `typing` module, use the correct frame (the one before): + return frame.f_back + + return frame + + def _init_core_attrs( + self, ns_resolver: _namespace_utils.NsResolver, force: bool, raise_errors: bool = False + ) -> bool: + """Initialize the core schema, validator, and serializer for the type. + + Args: + ns_resolver: The namespace resolver to use when building the core schema for the adapted type. + force: Whether to force the construction of the core schema, validator, and serializer. + If `force` is set to `False` and `_defer_build` is `True`, the core schema, validator, and serializer will be set to mocks. + raise_errors: Whether to raise errors if initializing any of the core attrs fails. + + Returns: + `True` if the core schema, validator, and serializer were successfully initialized, otherwise `False`. + + Raises: + PydanticUndefinedAnnotation: If `PydanticUndefinedAnnotation` occurs in`__get_pydantic_core_schema__` + and `raise_errors=True`. + """ + if not force and self._defer_build: + _mock_val_ser.set_type_adapter_mocks(self) + self.pydantic_complete = False + return False + + try: + self.core_schema = _getattr_no_parents(self._type, '__pydantic_core_schema__') + self.validator = _getattr_no_parents(self._type, '__pydantic_validator__') + self.serializer = _getattr_no_parents(self._type, '__pydantic_serializer__') + + # TODO: we don't go through the rebuild logic here directly because we don't want + # to repeat all of the namespace fetching logic that we've already done + # so we simply skip to the block below that does the actual schema generation + if ( + isinstance(self.core_schema, _mock_val_ser.MockCoreSchema) + or isinstance(self.validator, _mock_val_ser.MockValSer) + or isinstance(self.serializer, _mock_val_ser.MockValSer) + ): + raise AttributeError() + except AttributeError: + config_wrapper = _config.ConfigWrapper(self._config) + + schema_generator = _generate_schema.GenerateSchema(config_wrapper, ns_resolver=ns_resolver) + + try: + core_schema = schema_generator.generate_schema(self._type) + except PydanticUndefinedAnnotation: + if raise_errors: + raise + _mock_val_ser.set_type_adapter_mocks(self) + return False + + try: + self.core_schema = schema_generator.clean_schema(core_schema) + except _generate_schema.InvalidSchemaError: + _mock_val_ser.set_type_adapter_mocks(self) + return False + + core_config = config_wrapper.core_config(None) + + self.validator = create_schema_validator( + schema=self.core_schema, + schema_type=self._type, + schema_type_module=self._module_name, + schema_type_name=str(self._type), + schema_kind='TypeAdapter', + config=core_config, + plugin_settings=config_wrapper.plugin_settings, + ) + self.serializer = SchemaSerializer(self.core_schema, core_config) + + self.pydantic_complete = True + return True + + @property + def _defer_build(self) -> bool: + config = self._config if self._config is not None else self._model_config + if config: + return config.get('defer_build') is True + return False + + @property + def _model_config(self) -> ConfigDict | None: + type_: Any = _typing_extra.annotated_type(self._type) or self._type # Eg FastAPI heavily uses Annotated + if _utils.lenient_issubclass(type_, BaseModel): + return type_.model_config + return getattr(type_, '__pydantic_config__', None) + + def __repr__(self) -> str: + return f'TypeAdapter({_repr.display_as_type(self._type)})' + + def rebuild( + self, + *, + force: bool = False, + raise_errors: bool = True, + _parent_namespace_depth: int = 2, + _types_namespace: _namespace_utils.MappingNamespace | None = None, + ) -> bool | None: + """Try to rebuild the pydantic-core schema for the adapter's type. + + This may be necessary when one of the annotations is a ForwardRef which could not be resolved during + the initial attempt to build the schema, and automatic rebuilding fails. + + Args: + force: Whether to force the rebuilding of the type adapter's schema, defaults to `False`. + raise_errors: Whether to raise errors, defaults to `True`. + _parent_namespace_depth: Depth at which to search for the [parent frame][frame-objects]. This + frame is used when resolving forward annotations during schema rebuilding, by looking for + the locals of this frame. Defaults to 2, which will result in the frame where the method + was called. + _types_namespace: An explicit types namespace to use, instead of using the local namespace + from the parent frame. Defaults to `None`. + + Returns: + Returns `None` if the schema is already "complete" and rebuilding was not required. + If rebuilding _was_ required, returns `True` if rebuilding was successful, otherwise `False`. + """ + if not force and self.pydantic_complete: + return None + + if _types_namespace is not None: + rebuild_ns = _types_namespace + elif _parent_namespace_depth > 0: + rebuild_ns = _typing_extra.parent_frame_namespace(parent_depth=_parent_namespace_depth, force=True) or {} + else: + rebuild_ns = {} + + # we have to manually fetch globals here because there's no type on the stack of the NsResolver + # and so we skip the globalns = get_module_ns_of(typ) call that would normally happen + globalns = sys._getframe(max(_parent_namespace_depth - 1, 1)).f_globals + ns_resolver = _namespace_utils.NsResolver( + namespaces_tuple=_namespace_utils.NamespacesTuple(locals=rebuild_ns, globals=globalns), + parent_namespace=rebuild_ns, + ) + return self._init_core_attrs(ns_resolver=ns_resolver, force=True, raise_errors=raise_errors) + + def validate_python( + self, + object: Any, + /, + *, + strict: bool | None = None, + from_attributes: bool | None = None, + context: dict[str, Any] | None = None, + experimental_allow_partial: bool | Literal['off', 'on', 'trailing-strings'] = False, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> T: + """Validate a Python object against the model. + + Args: + object: The Python object to validate against the model. + strict: Whether to strictly check types. + from_attributes: Whether to extract data from object attributes. + context: Additional context to pass to the validator. + experimental_allow_partial: **Experimental** whether to enable + [partial validation](../concepts/experimental.md#partial-validation), e.g. to process streams. + * False / 'off': Default behavior, no partial validation. + * True / 'on': Enable partial validation. + * 'trailing-strings': Enable partial validation and allow trailing strings in the input. + by_alias: Whether to use the field's alias when validating against the provided input data. + by_name: Whether to use the field's name when validating against the provided input data. + + !!! note + When using `TypeAdapter` with a Pydantic `dataclass`, the use of the `from_attributes` + argument is not supported. + + Returns: + The validated object. + """ + if by_alias is False and by_name is not True: + raise PydanticUserError( + 'At least one of `by_alias` or `by_name` must be set to True.', + code='validate-by-alias-and-name-false', + ) + + return self.validator.validate_python( + object, + strict=strict, + from_attributes=from_attributes, + context=context, + allow_partial=experimental_allow_partial, + by_alias=by_alias, + by_name=by_name, + ) + + def validate_json( + self, + data: str | bytes | bytearray, + /, + *, + strict: bool | None = None, + context: dict[str, Any] | None = None, + experimental_allow_partial: bool | Literal['off', 'on', 'trailing-strings'] = False, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> T: + """!!! abstract "Usage Documentation" + [JSON Parsing](../concepts/json.md#json-parsing) + + Validate a JSON string or bytes against the model. + + Args: + data: The JSON data to validate against the model. + strict: Whether to strictly check types. + context: Additional context to use during validation. + experimental_allow_partial: **Experimental** whether to enable + [partial validation](../concepts/experimental.md#partial-validation), e.g. to process streams. + * False / 'off': Default behavior, no partial validation. + * True / 'on': Enable partial validation. + * 'trailing-strings': Enable partial validation and allow trailing strings in the input. + by_alias: Whether to use the field's alias when validating against the provided input data. + by_name: Whether to use the field's name when validating against the provided input data. + + Returns: + The validated object. + """ + if by_alias is False and by_name is not True: + raise PydanticUserError( + 'At least one of `by_alias` or `by_name` must be set to True.', + code='validate-by-alias-and-name-false', + ) + + return self.validator.validate_json( + data, + strict=strict, + context=context, + allow_partial=experimental_allow_partial, + by_alias=by_alias, + by_name=by_name, + ) + + def validate_strings( + self, + obj: Any, + /, + *, + strict: bool | None = None, + context: dict[str, Any] | None = None, + experimental_allow_partial: bool | Literal['off', 'on', 'trailing-strings'] = False, + by_alias: bool | None = None, + by_name: bool | None = None, + ) -> T: + """Validate object contains string data against the model. + + Args: + obj: The object contains string data to validate. + strict: Whether to strictly check types. + context: Additional context to use during validation. + experimental_allow_partial: **Experimental** whether to enable + [partial validation](../concepts/experimental.md#partial-validation), e.g. to process streams. + * False / 'off': Default behavior, no partial validation. + * True / 'on': Enable partial validation. + * 'trailing-strings': Enable partial validation and allow trailing strings in the input. + by_alias: Whether to use the field's alias when validating against the provided input data. + by_name: Whether to use the field's name when validating against the provided input data. + + Returns: + The validated object. + """ + if by_alias is False and by_name is not True: + raise PydanticUserError( + 'At least one of `by_alias` or `by_name` must be set to True.', + code='validate-by-alias-and-name-false', + ) + + return self.validator.validate_strings( + obj, + strict=strict, + context=context, + allow_partial=experimental_allow_partial, + by_alias=by_alias, + by_name=by_name, + ) + + def get_default_value(self, *, strict: bool | None = None, context: dict[str, Any] | None = None) -> Some[T] | None: + """Get the default value for the wrapped type. + + Args: + strict: Whether to strictly check types. + context: Additional context to pass to the validator. + + Returns: + The default value wrapped in a `Some` if there is one or None if not. + """ + return self.validator.get_default_value(strict=strict, context=context) + + def dump_python( + self, + instance: T, + /, + *, + mode: Literal['json', 'python'] = 'python', + include: IncEx | None = None, + exclude: IncEx | None = None, + by_alias: bool | None = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + round_trip: bool = False, + warnings: bool | Literal['none', 'warn', 'error'] = True, + fallback: Callable[[Any], Any] | None = None, + serialize_as_any: bool = False, + context: dict[str, Any] | None = None, + ) -> Any: + """Dump an instance of the adapted type to a Python object. + + Args: + instance: The Python object to serialize. + mode: The output format. + include: Fields to include in the output. + exclude: Fields to exclude from the output. + by_alias: Whether to use alias names for field names. + exclude_unset: Whether to exclude unset fields. + exclude_defaults: Whether to exclude fields with default values. + exclude_none: Whether to exclude fields with None values. + round_trip: Whether to output the serialized data in a way that is compatible with deserialization. + warnings: How to handle serialization errors. False/"none" ignores them, True/"warn" logs errors, + "error" raises a [`PydanticSerializationError`][pydantic_core.PydanticSerializationError]. + fallback: A function to call when an unknown value is encountered. If not provided, + a [`PydanticSerializationError`][pydantic_core.PydanticSerializationError] error is raised. + serialize_as_any: Whether to serialize fields with duck-typing serialization behavior. + context: Additional context to pass to the serializer. + + Returns: + The serialized object. + """ + return self.serializer.to_python( + instance, + mode=mode, + by_alias=by_alias, + include=include, + exclude=exclude, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + round_trip=round_trip, + warnings=warnings, + fallback=fallback, + serialize_as_any=serialize_as_any, + context=context, + ) + + def dump_json( + self, + instance: T, + /, + *, + indent: int | None = None, + include: IncEx | None = None, + exclude: IncEx | None = None, + by_alias: bool | None = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + round_trip: bool = False, + warnings: bool | Literal['none', 'warn', 'error'] = True, + fallback: Callable[[Any], Any] | None = None, + serialize_as_any: bool = False, + context: dict[str, Any] | None = None, + ) -> bytes: + """!!! abstract "Usage Documentation" + [JSON Serialization](../concepts/json.md#json-serialization) + + Serialize an instance of the adapted type to JSON. + + Args: + instance: The instance to be serialized. + indent: Number of spaces for JSON indentation. + include: Fields to include. + exclude: Fields to exclude. + by_alias: Whether to use alias names for field names. + exclude_unset: Whether to exclude unset fields. + exclude_defaults: Whether to exclude fields with default values. + exclude_none: Whether to exclude fields with a value of `None`. + round_trip: Whether to serialize and deserialize the instance to ensure round-tripping. + warnings: How to handle serialization errors. False/"none" ignores them, True/"warn" logs errors, + "error" raises a [`PydanticSerializationError`][pydantic_core.PydanticSerializationError]. + fallback: A function to call when an unknown value is encountered. If not provided, + a [`PydanticSerializationError`][pydantic_core.PydanticSerializationError] error is raised. + serialize_as_any: Whether to serialize fields with duck-typing serialization behavior. + context: Additional context to pass to the serializer. + + Returns: + The JSON representation of the given instance as bytes. + """ + return self.serializer.to_json( + instance, + indent=indent, + include=include, + exclude=exclude, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + round_trip=round_trip, + warnings=warnings, + fallback=fallback, + serialize_as_any=serialize_as_any, + context=context, + ) + + def json_schema( + self, + *, + by_alias: bool = True, + ref_template: str = DEFAULT_REF_TEMPLATE, + schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, + mode: JsonSchemaMode = 'validation', + ) -> dict[str, Any]: + """Generate a JSON schema for the adapted type. + + Args: + by_alias: Whether to use alias names for field names. + ref_template: The format string used for generating $ref strings. + schema_generator: The generator class used for creating the schema. + mode: The mode to use for schema generation. + + Returns: + The JSON schema for the model as a dictionary. + """ + schema_generator_instance = schema_generator(by_alias=by_alias, ref_template=ref_template) + if isinstance(self.core_schema, _mock_val_ser.MockCoreSchema): + self.core_schema.rebuild() + assert not isinstance(self.core_schema, _mock_val_ser.MockCoreSchema), 'this is a bug! please report it' + return schema_generator_instance.generate(self.core_schema, mode=mode) + + @staticmethod + def json_schemas( + inputs: Iterable[tuple[JsonSchemaKeyT, JsonSchemaMode, TypeAdapter[Any]]], + /, + *, + by_alias: bool = True, + title: str | None = None, + description: str | None = None, + ref_template: str = DEFAULT_REF_TEMPLATE, + schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema, + ) -> tuple[dict[tuple[JsonSchemaKeyT, JsonSchemaMode], JsonSchemaValue], JsonSchemaValue]: + """Generate a JSON schema including definitions from multiple type adapters. + + Args: + inputs: Inputs to schema generation. The first two items will form the keys of the (first) + output mapping; the type adapters will provide the core schemas that get converted into + definitions in the output JSON schema. + by_alias: Whether to use alias names. + title: The title for the schema. + description: The description for the schema. + ref_template: The format string used for generating $ref strings. + schema_generator: The generator class used for creating the schema. + + Returns: + A tuple where: + + - The first element is a dictionary whose keys are tuples of JSON schema key type and JSON mode, and + whose values are the JSON schema corresponding to that pair of inputs. (These schemas may have + JsonRef references to definitions that are defined in the second returned element.) + - The second element is a JSON schema containing all definitions referenced in the first returned + element, along with the optional title and description keys. + + """ + schema_generator_instance = schema_generator(by_alias=by_alias, ref_template=ref_template) + + inputs_ = [] + for key, mode, adapter in inputs: + # This is the same pattern we follow for model json schemas - we attempt a core schema rebuild if we detect a mock + if isinstance(adapter.core_schema, _mock_val_ser.MockCoreSchema): + adapter.core_schema.rebuild() + assert not isinstance(adapter.core_schema, _mock_val_ser.MockCoreSchema), ( + 'this is a bug! please report it' + ) + inputs_.append((key, mode, adapter.core_schema)) + + json_schemas_map, definitions = schema_generator_instance.generate_definitions(inputs_) + + json_schema: dict[str, Any] = {} + if definitions: + json_schema['$defs'] = definitions + if title: + json_schema['title'] = title + if description: + json_schema['description'] = description + + return json_schemas_map, json_schema diff --git a/lib/python3.12/site-packages/pydantic/types.py b/lib/python3.12/site-packages/pydantic/types.py new file mode 100644 index 0000000000000000000000000000000000000000..b5c4fd6060b42e1b552ef4ab117e313e0d6477fc --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/types.py @@ -0,0 +1,3285 @@ +"""The types module contains custom types used by pydantic.""" + +from __future__ import annotations as _annotations + +import base64 +import dataclasses as _dataclasses +import re +from collections.abc import Hashable, Iterator +from datetime import date, datetime +from decimal import Decimal +from enum import Enum +from pathlib import Path +from re import Pattern +from types import ModuleType +from typing import ( + TYPE_CHECKING, + Annotated, + Any, + Callable, + ClassVar, + Generic, + Literal, + TypeVar, + Union, + cast, +) +from uuid import UUID + +import annotated_types +from annotated_types import BaseMetadata, MaxLen, MinLen +from pydantic_core import CoreSchema, PydanticCustomError, SchemaSerializer, core_schema +from typing_extensions import Protocol, TypeAlias, TypeAliasType, deprecated, get_args, get_origin +from typing_inspection.introspection import is_union_origin + +from ._internal import _fields, _internal_dataclass, _utils, _validators +from ._migration import getattr_migration +from .annotated_handlers import GetCoreSchemaHandler, GetJsonSchemaHandler +from .errors import PydanticUserError +from .json_schema import JsonSchemaValue +from .warnings import PydanticDeprecatedSince20 + +if TYPE_CHECKING: + from ._internal._core_metadata import CoreMetadata + +__all__ = ( + 'Strict', + 'StrictStr', + 'SocketPath', + 'conbytes', + 'conlist', + 'conset', + 'confrozenset', + 'constr', + 'ImportString', + 'conint', + 'PositiveInt', + 'NegativeInt', + 'NonNegativeInt', + 'NonPositiveInt', + 'confloat', + 'PositiveFloat', + 'NegativeFloat', + 'NonNegativeFloat', + 'NonPositiveFloat', + 'FiniteFloat', + 'condecimal', + 'UUID1', + 'UUID3', + 'UUID4', + 'UUID5', + 'UUID6', + 'UUID7', + 'UUID8', + 'FilePath', + 'DirectoryPath', + 'NewPath', + 'Json', + 'Secret', + 'SecretStr', + 'SecretBytes', + 'StrictBool', + 'StrictBytes', + 'StrictInt', + 'StrictFloat', + 'PaymentCardNumber', + 'ByteSize', + 'PastDate', + 'FutureDate', + 'PastDatetime', + 'FutureDatetime', + 'condate', + 'AwareDatetime', + 'NaiveDatetime', + 'AllowInfNan', + 'EncoderProtocol', + 'EncodedBytes', + 'EncodedStr', + 'Base64Encoder', + 'Base64Bytes', + 'Base64Str', + 'Base64UrlBytes', + 'Base64UrlStr', + 'GetPydanticSchema', + 'StringConstraints', + 'Tag', + 'Discriminator', + 'JsonValue', + 'OnErrorOmit', + 'FailFast', +) + + +T = TypeVar('T') + + +@_dataclasses.dataclass +class Strict(_fields.PydanticMetadata, BaseMetadata): + """!!! abstract "Usage Documentation" + [Strict Mode with `Annotated` `Strict`](../concepts/strict_mode.md#strict-mode-with-annotated-strict) + + A field metadata class to indicate that a field should be validated in strict mode. + Use this class as an annotation via [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated), as seen below. + + Attributes: + strict: Whether to validate the field in strict mode. + + Example: + ```python + from typing import Annotated + + from pydantic.types import Strict + + StrictBool = Annotated[bool, Strict()] + ``` + """ + + strict: bool = True + + def __hash__(self) -> int: + return hash(self.strict) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ BOOLEAN TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +StrictBool = Annotated[bool, Strict()] +"""A boolean that must be either ``True`` or ``False``.""" + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ INTEGER TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +def conint( + *, + strict: bool | None = None, + gt: int | None = None, + ge: int | None = None, + lt: int | None = None, + le: int | None = None, + multiple_of: int | None = None, +) -> type[int]: + """ + !!! warning "Discouraged" + This function is **discouraged** in favor of using + [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated) with + [`Field`][pydantic.fields.Field] instead. + + This function will be **deprecated** in Pydantic 3.0. + + The reason is that `conint` returns a type, which doesn't play well with static analysis tools. + + === ":x: Don't do this" + ```python + from pydantic import BaseModel, conint + + class Foo(BaseModel): + bar: conint(strict=True, gt=0) + ``` + + === ":white_check_mark: Do this" + ```python + from typing import Annotated + + from pydantic import BaseModel, Field + + class Foo(BaseModel): + bar: Annotated[int, Field(strict=True, gt=0)] + ``` + + A wrapper around `int` that allows for additional constraints. + + Args: + strict: Whether to validate the integer in strict mode. Defaults to `None`. + gt: The value must be greater than this. + ge: The value must be greater than or equal to this. + lt: The value must be less than this. + le: The value must be less than or equal to this. + multiple_of: The value must be a multiple of this. + + Returns: + The wrapped integer type. + + ```python + from pydantic import BaseModel, ValidationError, conint + + class ConstrainedExample(BaseModel): + constrained_int: conint(gt=1) + + m = ConstrainedExample(constrained_int=2) + print(repr(m)) + #> ConstrainedExample(constrained_int=2) + + try: + ConstrainedExample(constrained_int=0) + except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'greater_than', + 'loc': ('constrained_int',), + 'msg': 'Input should be greater than 1', + 'input': 0, + 'ctx': {'gt': 1}, + 'url': 'https://errors.pydantic.dev/2/v/greater_than', + } + ] + ''' + ``` + + """ # noqa: D212 + return Annotated[ # pyright: ignore[reportReturnType] + int, + Strict(strict) if strict is not None else None, + annotated_types.Interval(gt=gt, ge=ge, lt=lt, le=le), + annotated_types.MultipleOf(multiple_of) if multiple_of is not None else None, + ] + + +PositiveInt = Annotated[int, annotated_types.Gt(0)] +"""An integer that must be greater than zero. + +```python +from pydantic import BaseModel, PositiveInt, ValidationError + +class Model(BaseModel): + positive_int: PositiveInt + +m = Model(positive_int=1) +print(repr(m)) +#> Model(positive_int=1) + +try: + Model(positive_int=-1) +except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'greater_than', + 'loc': ('positive_int',), + 'msg': 'Input should be greater than 0', + 'input': -1, + 'ctx': {'gt': 0}, + 'url': 'https://errors.pydantic.dev/2/v/greater_than', + } + ] + ''' +``` +""" +NegativeInt = Annotated[int, annotated_types.Lt(0)] +"""An integer that must be less than zero. + +```python +from pydantic import BaseModel, NegativeInt, ValidationError + +class Model(BaseModel): + negative_int: NegativeInt + +m = Model(negative_int=-1) +print(repr(m)) +#> Model(negative_int=-1) + +try: + Model(negative_int=1) +except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'less_than', + 'loc': ('negative_int',), + 'msg': 'Input should be less than 0', + 'input': 1, + 'ctx': {'lt': 0}, + 'url': 'https://errors.pydantic.dev/2/v/less_than', + } + ] + ''' +``` +""" +NonPositiveInt = Annotated[int, annotated_types.Le(0)] +"""An integer that must be less than or equal to zero. + +```python +from pydantic import BaseModel, NonPositiveInt, ValidationError + +class Model(BaseModel): + non_positive_int: NonPositiveInt + +m = Model(non_positive_int=0) +print(repr(m)) +#> Model(non_positive_int=0) + +try: + Model(non_positive_int=1) +except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'less_than_equal', + 'loc': ('non_positive_int',), + 'msg': 'Input should be less than or equal to 0', + 'input': 1, + 'ctx': {'le': 0}, + 'url': 'https://errors.pydantic.dev/2/v/less_than_equal', + } + ] + ''' +``` +""" +NonNegativeInt = Annotated[int, annotated_types.Ge(0)] +"""An integer that must be greater than or equal to zero. + +```python +from pydantic import BaseModel, NonNegativeInt, ValidationError + +class Model(BaseModel): + non_negative_int: NonNegativeInt + +m = Model(non_negative_int=0) +print(repr(m)) +#> Model(non_negative_int=0) + +try: + Model(non_negative_int=-1) +except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'greater_than_equal', + 'loc': ('non_negative_int',), + 'msg': 'Input should be greater than or equal to 0', + 'input': -1, + 'ctx': {'ge': 0}, + 'url': 'https://errors.pydantic.dev/2/v/greater_than_equal', + } + ] + ''' +``` +""" +StrictInt = Annotated[int, Strict()] +"""An integer that must be validated in strict mode. + +```python +from pydantic import BaseModel, StrictInt, ValidationError + +class StrictIntModel(BaseModel): + strict_int: StrictInt + +try: + StrictIntModel(strict_int=3.14159) +except ValidationError as e: + print(e) + ''' + 1 validation error for StrictIntModel + strict_int + Input should be a valid integer [type=int_type, input_value=3.14159, input_type=float] + ''' +``` +""" + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ FLOAT TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +@_dataclasses.dataclass +class AllowInfNan(_fields.PydanticMetadata): + """A field metadata class to indicate that a field should allow `-inf`, `inf`, and `nan`. + + Use this class as an annotation via [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated), as seen below. + + Attributes: + allow_inf_nan: Whether to allow `-inf`, `inf`, and `nan`. Defaults to `True`. + + Example: + ```python + from typing import Annotated + + from pydantic.types import AllowInfNan + + LaxFloat = Annotated[float, AllowInfNan()] + ``` + """ + + allow_inf_nan: bool = True + + def __hash__(self) -> int: + return hash(self.allow_inf_nan) + + +def confloat( + *, + strict: bool | None = None, + gt: float | None = None, + ge: float | None = None, + lt: float | None = None, + le: float | None = None, + multiple_of: float | None = None, + allow_inf_nan: bool | None = None, +) -> type[float]: + """ + !!! warning "Discouraged" + This function is **discouraged** in favor of using + [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated) with + [`Field`][pydantic.fields.Field] instead. + + This function will be **deprecated** in Pydantic 3.0. + + The reason is that `confloat` returns a type, which doesn't play well with static analysis tools. + + === ":x: Don't do this" + ```python + from pydantic import BaseModel, confloat + + class Foo(BaseModel): + bar: confloat(strict=True, gt=0) + ``` + + === ":white_check_mark: Do this" + ```python + from typing import Annotated + + from pydantic import BaseModel, Field + + class Foo(BaseModel): + bar: Annotated[float, Field(strict=True, gt=0)] + ``` + + A wrapper around `float` that allows for additional constraints. + + Args: + strict: Whether to validate the float in strict mode. + gt: The value must be greater than this. + ge: The value must be greater than or equal to this. + lt: The value must be less than this. + le: The value must be less than or equal to this. + multiple_of: The value must be a multiple of this. + allow_inf_nan: Whether to allow `-inf`, `inf`, and `nan`. + + Returns: + The wrapped float type. + + ```python + from pydantic import BaseModel, ValidationError, confloat + + class ConstrainedExample(BaseModel): + constrained_float: confloat(gt=1.0) + + m = ConstrainedExample(constrained_float=1.1) + print(repr(m)) + #> ConstrainedExample(constrained_float=1.1) + + try: + ConstrainedExample(constrained_float=0.9) + except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'greater_than', + 'loc': ('constrained_float',), + 'msg': 'Input should be greater than 1', + 'input': 0.9, + 'ctx': {'gt': 1.0}, + 'url': 'https://errors.pydantic.dev/2/v/greater_than', + } + ] + ''' + ``` + """ # noqa: D212 + return Annotated[ # pyright: ignore[reportReturnType] + float, + Strict(strict) if strict is not None else None, + annotated_types.Interval(gt=gt, ge=ge, lt=lt, le=le), + annotated_types.MultipleOf(multiple_of) if multiple_of is not None else None, + AllowInfNan(allow_inf_nan) if allow_inf_nan is not None else None, + ] + + +PositiveFloat = Annotated[float, annotated_types.Gt(0)] +"""A float that must be greater than zero. + +```python +from pydantic import BaseModel, PositiveFloat, ValidationError + +class Model(BaseModel): + positive_float: PositiveFloat + +m = Model(positive_float=1.0) +print(repr(m)) +#> Model(positive_float=1.0) + +try: + Model(positive_float=-1.0) +except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'greater_than', + 'loc': ('positive_float',), + 'msg': 'Input should be greater than 0', + 'input': -1.0, + 'ctx': {'gt': 0.0}, + 'url': 'https://errors.pydantic.dev/2/v/greater_than', + } + ] + ''' +``` +""" +NegativeFloat = Annotated[float, annotated_types.Lt(0)] +"""A float that must be less than zero. + +```python +from pydantic import BaseModel, NegativeFloat, ValidationError + +class Model(BaseModel): + negative_float: NegativeFloat + +m = Model(negative_float=-1.0) +print(repr(m)) +#> Model(negative_float=-1.0) + +try: + Model(negative_float=1.0) +except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'less_than', + 'loc': ('negative_float',), + 'msg': 'Input should be less than 0', + 'input': 1.0, + 'ctx': {'lt': 0.0}, + 'url': 'https://errors.pydantic.dev/2/v/less_than', + } + ] + ''' +``` +""" +NonPositiveFloat = Annotated[float, annotated_types.Le(0)] +"""A float that must be less than or equal to zero. + +```python +from pydantic import BaseModel, NonPositiveFloat, ValidationError + +class Model(BaseModel): + non_positive_float: NonPositiveFloat + +m = Model(non_positive_float=0.0) +print(repr(m)) +#> Model(non_positive_float=0.0) + +try: + Model(non_positive_float=1.0) +except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'less_than_equal', + 'loc': ('non_positive_float',), + 'msg': 'Input should be less than or equal to 0', + 'input': 1.0, + 'ctx': {'le': 0.0}, + 'url': 'https://errors.pydantic.dev/2/v/less_than_equal', + } + ] + ''' +``` +""" +NonNegativeFloat = Annotated[float, annotated_types.Ge(0)] +"""A float that must be greater than or equal to zero. + +```python +from pydantic import BaseModel, NonNegativeFloat, ValidationError + +class Model(BaseModel): + non_negative_float: NonNegativeFloat + +m = Model(non_negative_float=0.0) +print(repr(m)) +#> Model(non_negative_float=0.0) + +try: + Model(non_negative_float=-1.0) +except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'greater_than_equal', + 'loc': ('non_negative_float',), + 'msg': 'Input should be greater than or equal to 0', + 'input': -1.0, + 'ctx': {'ge': 0.0}, + 'url': 'https://errors.pydantic.dev/2/v/greater_than_equal', + } + ] + ''' +``` +""" +StrictFloat = Annotated[float, Strict(True)] +"""A float that must be validated in strict mode. + +```python +from pydantic import BaseModel, StrictFloat, ValidationError + +class StrictFloatModel(BaseModel): + strict_float: StrictFloat + +try: + StrictFloatModel(strict_float='1.0') +except ValidationError as e: + print(e) + ''' + 1 validation error for StrictFloatModel + strict_float + Input should be a valid number [type=float_type, input_value='1.0', input_type=str] + ''' +``` +""" +FiniteFloat = Annotated[float, AllowInfNan(False)] +"""A float that must be finite (not ``-inf``, ``inf``, or ``nan``). + +```python +from pydantic import BaseModel, FiniteFloat + +class Model(BaseModel): + finite: FiniteFloat + +m = Model(finite=1.0) +print(m) +#> finite=1.0 +``` +""" + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ BYTES TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +def conbytes( + *, + min_length: int | None = None, + max_length: int | None = None, + strict: bool | None = None, +) -> type[bytes]: + """A wrapper around `bytes` that allows for additional constraints. + + Args: + min_length: The minimum length of the bytes. + max_length: The maximum length of the bytes. + strict: Whether to validate the bytes in strict mode. + + Returns: + The wrapped bytes type. + """ + return Annotated[ # pyright: ignore[reportReturnType] + bytes, + Strict(strict) if strict is not None else None, + annotated_types.Len(min_length or 0, max_length), + ] + + +StrictBytes = Annotated[bytes, Strict()] +"""A bytes that must be validated in strict mode.""" + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ STRING TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +@_dataclasses.dataclass(frozen=True) +class StringConstraints(annotated_types.GroupedMetadata): + """!!! abstract "Usage Documentation" + [`StringConstraints`](../concepts/fields.md#string-constraints) + + A field metadata class to apply constraints to `str` types. + Use this class as an annotation via [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated), as seen below. + + Attributes: + strip_whitespace: Whether to remove leading and trailing whitespace. + to_upper: Whether to convert the string to uppercase. + to_lower: Whether to convert the string to lowercase. + strict: Whether to validate the string in strict mode. + min_length: The minimum length of the string. + max_length: The maximum length of the string. + pattern: A regex pattern that the string must match. + + Example: + ```python + from typing import Annotated + + from pydantic.types import StringConstraints + + ConstrainedStr = Annotated[str, StringConstraints(min_length=1, max_length=10)] + ``` + """ + + strip_whitespace: bool | None = None + to_upper: bool | None = None + to_lower: bool | None = None + strict: bool | None = None + min_length: int | None = None + max_length: int | None = None + pattern: str | Pattern[str] | None = None + + def __iter__(self) -> Iterator[BaseMetadata]: + if self.min_length is not None: + yield MinLen(self.min_length) + if self.max_length is not None: + yield MaxLen(self.max_length) + if self.strict is not None: + yield Strict(self.strict) + if ( + self.strip_whitespace is not None + or self.pattern is not None + or self.to_lower is not None + or self.to_upper is not None + ): + yield _fields.pydantic_general_metadata( + strip_whitespace=self.strip_whitespace, + to_upper=self.to_upper, + to_lower=self.to_lower, + pattern=self.pattern, + ) + + +def constr( + *, + strip_whitespace: bool | None = None, + to_upper: bool | None = None, + to_lower: bool | None = None, + strict: bool | None = None, + min_length: int | None = None, + max_length: int | None = None, + pattern: str | Pattern[str] | None = None, +) -> type[str]: + """ + !!! warning "Discouraged" + This function is **discouraged** in favor of using + [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated) with + [`StringConstraints`][pydantic.types.StringConstraints] instead. + + This function will be **deprecated** in Pydantic 3.0. + + The reason is that `constr` returns a type, which doesn't play well with static analysis tools. + + === ":x: Don't do this" + ```python + from pydantic import BaseModel, constr + + class Foo(BaseModel): + bar: constr(strip_whitespace=True, to_upper=True, pattern=r'^[A-Z]+$') + ``` + + === ":white_check_mark: Do this" + ```python + from typing import Annotated + + from pydantic import BaseModel, StringConstraints + + class Foo(BaseModel): + bar: Annotated[ + str, + StringConstraints( + strip_whitespace=True, to_upper=True, pattern=r'^[A-Z]+$' + ), + ] + ``` + + A wrapper around `str` that allows for additional constraints. + + ```python + from pydantic import BaseModel, constr + + class Foo(BaseModel): + bar: constr(strip_whitespace=True, to_upper=True) + + foo = Foo(bar=' hello ') + print(foo) + #> bar='HELLO' + ``` + + Args: + strip_whitespace: Whether to remove leading and trailing whitespace. + to_upper: Whether to turn all characters to uppercase. + to_lower: Whether to turn all characters to lowercase. + strict: Whether to validate the string in strict mode. + min_length: The minimum length of the string. + max_length: The maximum length of the string. + pattern: A regex pattern to validate the string against. + + Returns: + The wrapped string type. + """ # noqa: D212 + return Annotated[ # pyright: ignore[reportReturnType] + str, + StringConstraints( + strip_whitespace=strip_whitespace, + to_upper=to_upper, + to_lower=to_lower, + strict=strict, + min_length=min_length, + max_length=max_length, + pattern=pattern, + ), + ] + + +StrictStr = Annotated[str, Strict()] +"""A string that must be validated in strict mode.""" + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~ COLLECTION TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +HashableItemType = TypeVar('HashableItemType', bound=Hashable) + + +def conset( + item_type: type[HashableItemType], *, min_length: int | None = None, max_length: int | None = None +) -> type[set[HashableItemType]]: + """A wrapper around `typing.Set` that allows for additional constraints. + + Args: + item_type: The type of the items in the set. + min_length: The minimum length of the set. + max_length: The maximum length of the set. + + Returns: + The wrapped set type. + """ + return Annotated[set[item_type], annotated_types.Len(min_length or 0, max_length)] # pyright: ignore[reportReturnType] + + +def confrozenset( + item_type: type[HashableItemType], *, min_length: int | None = None, max_length: int | None = None +) -> type[frozenset[HashableItemType]]: + """A wrapper around `typing.FrozenSet` that allows for additional constraints. + + Args: + item_type: The type of the items in the frozenset. + min_length: The minimum length of the frozenset. + max_length: The maximum length of the frozenset. + + Returns: + The wrapped frozenset type. + """ + return Annotated[frozenset[item_type], annotated_types.Len(min_length or 0, max_length)] # pyright: ignore[reportReturnType] + + +AnyItemType = TypeVar('AnyItemType') + + +def conlist( + item_type: type[AnyItemType], + *, + min_length: int | None = None, + max_length: int | None = None, + unique_items: bool | None = None, +) -> type[list[AnyItemType]]: + """A wrapper around [`list`][] that adds validation. + + Args: + item_type: The type of the items in the list. + min_length: The minimum length of the list. Defaults to None. + max_length: The maximum length of the list. Defaults to None. + unique_items: Whether the items in the list must be unique. Defaults to None. + !!! warning Deprecated + The `unique_items` parameter is deprecated, use `Set` instead. + See [this issue](https://github.com/pydantic/pydantic-core/issues/296) for more details. + + Returns: + The wrapped list type. + """ + if unique_items is not None: + raise PydanticUserError( + ( + '`unique_items` is removed, use `Set` instead' + '(this feature is discussed in https://github.com/pydantic/pydantic-core/issues/296)' + ), + code='removed-kwargs', + ) + return Annotated[list[item_type], annotated_types.Len(min_length or 0, max_length)] # pyright: ignore[reportReturnType] + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~ IMPORT STRING TYPE ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +AnyType = TypeVar('AnyType') +if TYPE_CHECKING: + ImportString = Annotated[AnyType, ...] +else: + + class ImportString: + """A type that can be used to import a Python object from a string. + + `ImportString` expects a string and loads the Python object importable at that dotted path. + Attributes of modules may be separated from the module by `:` or `.`, e.g. if `'math:cos'` is provided, + the resulting field value would be the function `cos`. If a `.` is used and both an attribute and submodule + are present at the same path, the module will be preferred. + + On model instantiation, pointers will be evaluated and imported. There is + some nuance to this behavior, demonstrated in the examples below. + + ```python + import math + + from pydantic import BaseModel, Field, ImportString, ValidationError + + class ImportThings(BaseModel): + obj: ImportString + + # A string value will cause an automatic import + my_cos = ImportThings(obj='math.cos') + + # You can use the imported function as you would expect + cos_of_0 = my_cos.obj(0) + assert cos_of_0 == 1 + + # A string whose value cannot be imported will raise an error + try: + ImportThings(obj='foo.bar') + except ValidationError as e: + print(e) + ''' + 1 validation error for ImportThings + obj + Invalid python path: No module named 'foo.bar' [type=import_error, input_value='foo.bar', input_type=str] + ''' + + # Actual python objects can be assigned as well + my_cos = ImportThings(obj=math.cos) + my_cos_2 = ImportThings(obj='math.cos') + my_cos_3 = ImportThings(obj='math:cos') + assert my_cos == my_cos_2 == my_cos_3 + + # You can set default field value either as Python object: + class ImportThingsDefaultPyObj(BaseModel): + obj: ImportString = math.cos + + # or as a string value (but only if used with `validate_default=True`) + class ImportThingsDefaultString(BaseModel): + obj: ImportString = Field(default='math.cos', validate_default=True) + + my_cos_default1 = ImportThingsDefaultPyObj() + my_cos_default2 = ImportThingsDefaultString() + assert my_cos_default1.obj == my_cos_default2.obj == math.cos + + # note: this will not work! + class ImportThingsMissingValidateDefault(BaseModel): + obj: ImportString = 'math.cos' + + my_cos_default3 = ImportThingsMissingValidateDefault() + assert my_cos_default3.obj == 'math.cos' # just string, not evaluated + ``` + + Serializing an `ImportString` type to json is also possible. + + ```python + from pydantic import BaseModel, ImportString + + class ImportThings(BaseModel): + obj: ImportString + + # Create an instance + m = ImportThings(obj='math.cos') + print(m) + #> obj= + print(m.model_dump_json()) + #> {"obj":"math.cos"} + ``` + """ + + @classmethod + def __class_getitem__(cls, item: AnyType) -> AnyType: + return Annotated[item, cls()] + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[Any], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + serializer = core_schema.plain_serializer_function_ser_schema(cls._serialize, when_used='json') + if cls is source: + # Treat bare usage of ImportString (`schema is None`) as the same as ImportString[Any] + return core_schema.no_info_plain_validator_function( + function=_validators.import_string, serialization=serializer + ) + else: + return core_schema.no_info_before_validator_function( + function=_validators.import_string, schema=handler(source), serialization=serializer + ) + + @classmethod + def __get_pydantic_json_schema__(cls, cs: CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue: + return handler(core_schema.str_schema()) + + @staticmethod + def _serialize(v: Any) -> str: + if isinstance(v, ModuleType): + return v.__name__ + elif hasattr(v, '__module__') and hasattr(v, '__name__'): + return f'{v.__module__}.{v.__name__}' + # Handle special cases for sys.XXX streams + # if we see more of these, we should consider a more general solution + elif hasattr(v, 'name'): + if v.name == '': + return 'sys.stdout' + elif v.name == '': + return 'sys.stdin' + elif v.name == '': + return 'sys.stderr' + else: + return v + + def __repr__(self) -> str: + return 'ImportString' + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ DECIMAL TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +def condecimal( + *, + strict: bool | None = None, + gt: int | Decimal | None = None, + ge: int | Decimal | None = None, + lt: int | Decimal | None = None, + le: int | Decimal | None = None, + multiple_of: int | Decimal | None = None, + max_digits: int | None = None, + decimal_places: int | None = None, + allow_inf_nan: bool | None = None, +) -> type[Decimal]: + """ + !!! warning "Discouraged" + This function is **discouraged** in favor of using + [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated) with + [`Field`][pydantic.fields.Field] instead. + + This function will be **deprecated** in Pydantic 3.0. + + The reason is that `condecimal` returns a type, which doesn't play well with static analysis tools. + + === ":x: Don't do this" + ```python + from pydantic import BaseModel, condecimal + + class Foo(BaseModel): + bar: condecimal(strict=True, allow_inf_nan=True) + ``` + + === ":white_check_mark: Do this" + ```python + from decimal import Decimal + from typing import Annotated + + from pydantic import BaseModel, Field + + class Foo(BaseModel): + bar: Annotated[Decimal, Field(strict=True, allow_inf_nan=True)] + ``` + + A wrapper around Decimal that adds validation. + + Args: + strict: Whether to validate the value in strict mode. Defaults to `None`. + gt: The value must be greater than this. Defaults to `None`. + ge: The value must be greater than or equal to this. Defaults to `None`. + lt: The value must be less than this. Defaults to `None`. + le: The value must be less than or equal to this. Defaults to `None`. + multiple_of: The value must be a multiple of this. Defaults to `None`. + max_digits: The maximum number of digits. Defaults to `None`. + decimal_places: The number of decimal places. Defaults to `None`. + allow_inf_nan: Whether to allow infinity and NaN. Defaults to `None`. + + ```python + from decimal import Decimal + + from pydantic import BaseModel, ValidationError, condecimal + + class ConstrainedExample(BaseModel): + constrained_decimal: condecimal(gt=Decimal('1.0')) + + m = ConstrainedExample(constrained_decimal=Decimal('1.1')) + print(repr(m)) + #> ConstrainedExample(constrained_decimal=Decimal('1.1')) + + try: + ConstrainedExample(constrained_decimal=Decimal('0.9')) + except ValidationError as e: + print(e.errors()) + ''' + [ + { + 'type': 'greater_than', + 'loc': ('constrained_decimal',), + 'msg': 'Input should be greater than 1.0', + 'input': Decimal('0.9'), + 'ctx': {'gt': Decimal('1.0')}, + 'url': 'https://errors.pydantic.dev/2/v/greater_than', + } + ] + ''' + ``` + """ # noqa: D212 + return Annotated[ # pyright: ignore[reportReturnType] + Decimal, + Strict(strict) if strict is not None else None, + annotated_types.Interval(gt=gt, ge=ge, lt=lt, le=le), + annotated_types.MultipleOf(multiple_of) if multiple_of is not None else None, + _fields.pydantic_general_metadata(max_digits=max_digits, decimal_places=decimal_places), + AllowInfNan(allow_inf_nan) if allow_inf_nan is not None else None, + ] + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ UUID TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +@_dataclasses.dataclass(**_internal_dataclass.slots_true) +class UuidVersion: + """A field metadata class to indicate a [UUID](https://docs.python.org/3/library/uuid.html) version. + + Use this class as an annotation via [`Annotated`](https://docs.python.org/3/library/typing.html#typing.Annotated), as seen below. + + Attributes: + uuid_version: The version of the UUID. Must be one of 1, 3, 4, 5, or 7. + + Example: + ```python + from typing import Annotated + from uuid import UUID + + from pydantic.types import UuidVersion + + UUID1 = Annotated[UUID, UuidVersion(1)] + ``` + """ + + uuid_version: Literal[1, 3, 4, 5, 6, 7, 8] + + def __get_pydantic_json_schema__( + self, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = handler(core_schema) + field_schema.pop('anyOf', None) # remove the bytes/str union + field_schema.update(type='string', format=f'uuid{self.uuid_version}') + return field_schema + + def __get_pydantic_core_schema__(self, source: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + if isinstance(self, source): + # used directly as a type + return core_schema.uuid_schema(version=self.uuid_version) + else: + # update existing schema with self.uuid_version + schema = handler(source) + _check_annotated_type(schema['type'], 'uuid', self.__class__.__name__) + schema['version'] = self.uuid_version # type: ignore + return schema + + def __hash__(self) -> int: + return hash(type(self.uuid_version)) + + +UUID1 = Annotated[UUID, UuidVersion(1)] +"""A [UUID](https://docs.python.org/3/library/uuid.html) that must be version 1. + +```python +import uuid + +from pydantic import UUID1, BaseModel + +class Model(BaseModel): + uuid1: UUID1 + +Model(uuid1=uuid.uuid1()) +``` +""" +UUID3 = Annotated[UUID, UuidVersion(3)] +"""A [UUID](https://docs.python.org/3/library/uuid.html) that must be version 3. + +```python +import uuid + +from pydantic import UUID3, BaseModel + +class Model(BaseModel): + uuid3: UUID3 + +Model(uuid3=uuid.uuid3(uuid.NAMESPACE_DNS, 'pydantic.org')) +``` +""" +UUID4 = Annotated[UUID, UuidVersion(4)] +"""A [UUID](https://docs.python.org/3/library/uuid.html) that must be version 4. + +```python +import uuid + +from pydantic import UUID4, BaseModel + +class Model(BaseModel): + uuid4: UUID4 + +Model(uuid4=uuid.uuid4()) +``` +""" +UUID5 = Annotated[UUID, UuidVersion(5)] +"""A [UUID](https://docs.python.org/3/library/uuid.html) that must be version 5. + +```python +import uuid + +from pydantic import UUID5, BaseModel + +class Model(BaseModel): + uuid5: UUID5 + +Model(uuid5=uuid.uuid5(uuid.NAMESPACE_DNS, 'pydantic.org')) +``` +""" +UUID6 = Annotated[UUID, UuidVersion(6)] +"""A [UUID](https://docs.python.org/3/library/uuid.html) that must be version 6. + +```python +import uuid + +from pydantic import UUID6, BaseModel + +class Model(BaseModel): + uuid6: UUID6 + +Model(uuid6=uuid.UUID('1efea953-c2d6-6790-aa0a-69db8c87df97')) +``` +""" +UUID7 = Annotated[UUID, UuidVersion(7)] +"""A [UUID](https://docs.python.org/3/library/uuid.html) that must be version 7. + +```python +import uuid + +from pydantic import UUID7, BaseModel + +class Model(BaseModel): + uuid7: UUID7 + +Model(uuid7=uuid.UUID('0194fdcb-1c47-7a09-b52c-561154de0b4a')) +``` +""" +UUID8 = Annotated[UUID, UuidVersion(8)] +"""A [UUID](https://docs.python.org/3/library/uuid.html) that must be version 8. + +```python +import uuid + +from pydantic import UUID8, BaseModel + +class Model(BaseModel): + uuid8: UUID8 + +Model(uuid8=uuid.UUID('81a0b92e-6078-8551-9c81-8ccb666bdab8')) +``` +""" + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ PATH TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +@_dataclasses.dataclass +class PathType: + path_type: Literal['file', 'dir', 'new', 'socket'] + + def __get_pydantic_json_schema__( + self, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = handler(core_schema) + format_conversion = {'file': 'file-path', 'dir': 'directory-path'} + field_schema.update(format=format_conversion.get(self.path_type, 'path'), type='string') + return field_schema + + def __get_pydantic_core_schema__(self, source: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + function_lookup = { + 'file': cast(core_schema.WithInfoValidatorFunction, self.validate_file), + 'dir': cast(core_schema.WithInfoValidatorFunction, self.validate_directory), + 'new': cast(core_schema.WithInfoValidatorFunction, self.validate_new), + 'socket': cast(core_schema.WithInfoValidatorFunction, self.validate_socket), + } + + return core_schema.with_info_after_validator_function( + function_lookup[self.path_type], + handler(source), + ) + + @staticmethod + def validate_file(path: Path, _: core_schema.ValidationInfo) -> Path: + if path.is_file(): + return path + else: + raise PydanticCustomError('path_not_file', 'Path does not point to a file') + + @staticmethod + def validate_socket(path: Path, _: core_schema.ValidationInfo) -> Path: + if path.is_socket(): + return path + else: + raise PydanticCustomError('path_not_socket', 'Path does not point to a socket') + + @staticmethod + def validate_directory(path: Path, _: core_schema.ValidationInfo) -> Path: + if path.is_dir(): + return path + else: + raise PydanticCustomError('path_not_directory', 'Path does not point to a directory') + + @staticmethod + def validate_new(path: Path, _: core_schema.ValidationInfo) -> Path: + if path.exists(): + raise PydanticCustomError('path_exists', 'Path already exists') + elif not path.parent.exists(): + raise PydanticCustomError('parent_does_not_exist', 'Parent directory does not exist') + else: + return path + + def __hash__(self) -> int: + return hash(type(self.path_type)) + + +FilePath = Annotated[Path, PathType('file')] +"""A path that must point to a file. + +```python +from pathlib import Path + +from pydantic import BaseModel, FilePath, ValidationError + +class Model(BaseModel): + f: FilePath + +path = Path('text.txt') +path.touch() +m = Model(f='text.txt') +print(m.model_dump()) +#> {'f': PosixPath('text.txt')} +path.unlink() + +path = Path('directory') +path.mkdir(exist_ok=True) +try: + Model(f='directory') # directory +except ValidationError as e: + print(e) + ''' + 1 validation error for Model + f + Path does not point to a file [type=path_not_file, input_value='directory', input_type=str] + ''' +path.rmdir() + +try: + Model(f='not-exists-file') +except ValidationError as e: + print(e) + ''' + 1 validation error for Model + f + Path does not point to a file [type=path_not_file, input_value='not-exists-file', input_type=str] + ''' +``` +""" +DirectoryPath = Annotated[Path, PathType('dir')] +"""A path that must point to a directory. + +```python +from pathlib import Path + +from pydantic import BaseModel, DirectoryPath, ValidationError + +class Model(BaseModel): + f: DirectoryPath + +path = Path('directory/') +path.mkdir() +m = Model(f='directory/') +print(m.model_dump()) +#> {'f': PosixPath('directory')} +path.rmdir() + +path = Path('file.txt') +path.touch() +try: + Model(f='file.txt') # file +except ValidationError as e: + print(e) + ''' + 1 validation error for Model + f + Path does not point to a directory [type=path_not_directory, input_value='file.txt', input_type=str] + ''' +path.unlink() + +try: + Model(f='not-exists-directory') +except ValidationError as e: + print(e) + ''' + 1 validation error for Model + f + Path does not point to a directory [type=path_not_directory, input_value='not-exists-directory', input_type=str] + ''' +``` +""" +NewPath = Annotated[Path, PathType('new')] +"""A path for a new file or directory that must not already exist. The parent directory must already exist.""" + +SocketPath = Annotated[Path, PathType('socket')] +"""A path to an existing socket file""" + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ JSON TYPE ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +if TYPE_CHECKING: + # Json[list[str]] will be recognized by type checkers as list[str] + Json = Annotated[AnyType, ...] + +else: + + class Json: + """A special type wrapper which loads JSON before parsing. + + You can use the `Json` data type to make Pydantic first load a raw JSON string before + validating the loaded data into the parametrized type: + + ```python + from typing import Any + + from pydantic import BaseModel, Json, ValidationError + + class AnyJsonModel(BaseModel): + json_obj: Json[Any] + + class ConstrainedJsonModel(BaseModel): + json_obj: Json[list[int]] + + print(AnyJsonModel(json_obj='{"b": 1}')) + #> json_obj={'b': 1} + print(ConstrainedJsonModel(json_obj='[1, 2, 3]')) + #> json_obj=[1, 2, 3] + + try: + ConstrainedJsonModel(json_obj=12) + except ValidationError as e: + print(e) + ''' + 1 validation error for ConstrainedJsonModel + json_obj + JSON input should be string, bytes or bytearray [type=json_type, input_value=12, input_type=int] + ''' + + try: + ConstrainedJsonModel(json_obj='[a, b]') + except ValidationError as e: + print(e) + ''' + 1 validation error for ConstrainedJsonModel + json_obj + Invalid JSON: expected value at line 1 column 2 [type=json_invalid, input_value='[a, b]', input_type=str] + ''' + + try: + ConstrainedJsonModel(json_obj='["a", "b"]') + except ValidationError as e: + print(e) + ''' + 2 validation errors for ConstrainedJsonModel + json_obj.0 + Input should be a valid integer, unable to parse string as an integer [type=int_parsing, input_value='a', input_type=str] + json_obj.1 + Input should be a valid integer, unable to parse string as an integer [type=int_parsing, input_value='b', input_type=str] + ''' + ``` + + When you dump the model using `model_dump` or `model_dump_json`, the dumped value will be the result of validation, + not the original JSON string. However, you can use the argument `round_trip=True` to get the original JSON string back: + + ```python + from pydantic import BaseModel, Json + + class ConstrainedJsonModel(BaseModel): + json_obj: Json[list[int]] + + print(ConstrainedJsonModel(json_obj='[1, 2, 3]').model_dump_json()) + #> {"json_obj":[1,2,3]} + print( + ConstrainedJsonModel(json_obj='[1, 2, 3]').model_dump_json(round_trip=True) + ) + #> {"json_obj":"[1,2,3]"} + ``` + """ + + @classmethod + def __class_getitem__(cls, item: AnyType) -> AnyType: + return Annotated[item, cls()] + + @classmethod + def __get_pydantic_core_schema__(cls, source: Any, handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + if cls is source: + return core_schema.json_schema(None) + else: + return core_schema.json_schema(handler(source)) + + def __repr__(self) -> str: + return 'Json' + + def __hash__(self) -> int: + return hash(type(self)) + + def __eq__(self, other: Any) -> bool: + return type(other) is type(self) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ SECRET TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +SecretType = TypeVar('SecretType') + + +class _SecretBase(Generic[SecretType]): + def __init__(self, secret_value: SecretType) -> None: + self._secret_value: SecretType = secret_value + + def get_secret_value(self) -> SecretType: + """Get the secret value. + + Returns: + The secret value. + """ + return self._secret_value + + def __eq__(self, other: Any) -> bool: + return isinstance(other, self.__class__) and self.get_secret_value() == other.get_secret_value() + + def __hash__(self) -> int: + return hash(self.get_secret_value()) + + def __str__(self) -> str: + return str(self._display()) + + def __repr__(self) -> str: + return f'{self.__class__.__name__}({self._display()!r})' + + def _display(self) -> str | bytes: + raise NotImplementedError + + +def _serialize_secret(value: Secret[SecretType], info: core_schema.SerializationInfo) -> str | Secret[SecretType]: + if info.mode == 'json': + return str(value) + else: + return value + + +class Secret(_SecretBase[SecretType]): + """A generic base class used for defining a field with sensitive information that you do not want to be visible in logging or tracebacks. + + You may either directly parametrize `Secret` with a type, or subclass from `Secret` with a parametrized type. The benefit of subclassing + is that you can define a custom `_display` method, which will be used for `repr()` and `str()` methods. The examples below demonstrate both + ways of using `Secret` to create a new secret type. + + 1. Directly parametrizing `Secret` with a type: + + ```python + from pydantic import BaseModel, Secret + + SecretBool = Secret[bool] + + class Model(BaseModel): + secret_bool: SecretBool + + m = Model(secret_bool=True) + print(m.model_dump()) + #> {'secret_bool': Secret('**********')} + + print(m.model_dump_json()) + #> {"secret_bool":"**********"} + + print(m.secret_bool.get_secret_value()) + #> True + ``` + + 2. Subclassing from parametrized `Secret`: + + ```python + from datetime import date + + from pydantic import BaseModel, Secret + + class SecretDate(Secret[date]): + def _display(self) -> str: + return '****/**/**' + + class Model(BaseModel): + secret_date: SecretDate + + m = Model(secret_date=date(2022, 1, 1)) + print(m.model_dump()) + #> {'secret_date': SecretDate('****/**/**')} + + print(m.model_dump_json()) + #> {"secret_date":"****/**/**"} + + print(m.secret_date.get_secret_value()) + #> 2022-01-01 + ``` + + The value returned by the `_display` method will be used for `repr()` and `str()`. + + You can enforce constraints on the underlying type through annotations: + For example: + + ```python + from typing import Annotated + + from pydantic import BaseModel, Field, Secret, ValidationError + + SecretPosInt = Secret[Annotated[int, Field(gt=0, strict=True)]] + + class Model(BaseModel): + sensitive_int: SecretPosInt + + m = Model(sensitive_int=42) + print(m.model_dump()) + #> {'sensitive_int': Secret('**********')} + + try: + m = Model(sensitive_int=-42) # (1)! + except ValidationError as exc_info: + print(exc_info.errors(include_url=False, include_input=False)) + ''' + [ + { + 'type': 'greater_than', + 'loc': ('sensitive_int',), + 'msg': 'Input should be greater than 0', + 'ctx': {'gt': 0}, + } + ] + ''' + + try: + m = Model(sensitive_int='42') # (2)! + except ValidationError as exc_info: + print(exc_info.errors(include_url=False, include_input=False)) + ''' + [ + { + 'type': 'int_type', + 'loc': ('sensitive_int',), + 'msg': 'Input should be a valid integer', + } + ] + ''' + ``` + + 1. The input value is not greater than 0, so it raises a validation error. + 2. The input value is not an integer, so it raises a validation error because the `SecretPosInt` type has strict mode enabled. + """ + + def _display(self) -> str | bytes: + return '**********' if self.get_secret_value() else '' + + @classmethod + def __get_pydantic_core_schema__(cls, source: type[Any], handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + inner_type = None + # if origin_type is Secret, then cls is a GenericAlias, and we can extract the inner type directly + origin_type = get_origin(source) + if origin_type is not None: + inner_type = get_args(source)[0] + # otherwise, we need to get the inner type from the base class + else: + bases = getattr(cls, '__orig_bases__', getattr(cls, '__bases__', [])) + for base in bases: + if get_origin(base) is Secret: + inner_type = get_args(base)[0] + if bases == [] or inner_type is None: + raise TypeError( + f"Can't get secret type from {cls.__name__}. " + 'Please use Secret[], or subclass from Secret[] instead.' + ) + + inner_schema = handler.generate_schema(inner_type) # type: ignore + + def validate_secret_value(value, handler) -> Secret[SecretType]: + if isinstance(value, Secret): + value = value.get_secret_value() + validated_inner = handler(value) + return cls(validated_inner) + + return core_schema.json_or_python_schema( + python_schema=core_schema.no_info_wrap_validator_function( + validate_secret_value, + inner_schema, + ), + json_schema=core_schema.no_info_after_validator_function(lambda x: cls(x), inner_schema), + serialization=core_schema.plain_serializer_function_ser_schema( + _serialize_secret, + info_arg=True, + when_used='always', + ), + ) + + __pydantic_serializer__ = SchemaSerializer( + core_schema.any_schema( + serialization=core_schema.plain_serializer_function_ser_schema( + _serialize_secret, + info_arg=True, + when_used='always', + ) + ) + ) + + +def _secret_display(value: SecretType) -> str: # type: ignore + return '**********' if value else '' + + +def _serialize_secret_field( + value: _SecretField[SecretType], info: core_schema.SerializationInfo +) -> str | _SecretField[SecretType]: + if info.mode == 'json': + # we want the output to always be string without the `b'` prefix for bytes, + # hence we just use `secret_display` + return _secret_display(value.get_secret_value()) + else: + return value + + +class _SecretField(_SecretBase[SecretType]): + _inner_schema: ClassVar[CoreSchema] + _error_kind: ClassVar[str] + + @classmethod + def __get_pydantic_core_schema__(cls, source: type[Any], handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + def get_json_schema(_core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler) -> JsonSchemaValue: + json_schema = handler(cls._inner_schema) + _utils.update_not_none( + json_schema, + type='string', + writeOnly=True, + format='password', + ) + return json_schema + + def get_secret_schema(strict: bool) -> CoreSchema: + inner_schema = {**cls._inner_schema, 'strict': strict} + json_schema = core_schema.no_info_after_validator_function( + source, # construct the type + inner_schema, # pyright: ignore[reportArgumentType] + ) + return core_schema.json_or_python_schema( + python_schema=core_schema.union_schema( + [ + core_schema.is_instance_schema(source), + json_schema, + ], + custom_error_type=cls._error_kind, + ), + json_schema=json_schema, + serialization=core_schema.plain_serializer_function_ser_schema( + _serialize_secret_field, + info_arg=True, + when_used='always', + ), + ) + + return core_schema.lax_or_strict_schema( + lax_schema=get_secret_schema(strict=False), + strict_schema=get_secret_schema(strict=True), + metadata={'pydantic_js_functions': [get_json_schema]}, + ) + + __pydantic_serializer__ = SchemaSerializer( + core_schema.any_schema( + serialization=core_schema.plain_serializer_function_ser_schema( + _serialize_secret_field, + info_arg=True, + when_used='always', + ) + ) + ) + + +class SecretStr(_SecretField[str]): + """A string used for storing sensitive information that you do not want to be visible in logging or tracebacks. + + When the secret value is nonempty, it is displayed as `'**********'` instead of the underlying value in + calls to `repr()` and `str()`. If the value _is_ empty, it is displayed as `''`. + + ```python + from pydantic import BaseModel, SecretStr + + class User(BaseModel): + username: str + password: SecretStr + + user = User(username='scolvin', password='password1') + + print(user) + #> username='scolvin' password=SecretStr('**********') + print(user.password.get_secret_value()) + #> password1 + print((SecretStr('password'), SecretStr(''))) + #> (SecretStr('**********'), SecretStr('')) + ``` + + As seen above, by default, [`SecretStr`][pydantic.types.SecretStr] (and [`SecretBytes`][pydantic.types.SecretBytes]) + will be serialized as `**********` when serializing to json. + + You can use the [`field_serializer`][pydantic.functional_serializers.field_serializer] to dump the + secret as plain-text when serializing to json. + + ```python + from pydantic import BaseModel, SecretBytes, SecretStr, field_serializer + + class Model(BaseModel): + password: SecretStr + password_bytes: SecretBytes + + @field_serializer('password', 'password_bytes', when_used='json') + def dump_secret(self, v): + return v.get_secret_value() + + model = Model(password='IAmSensitive', password_bytes=b'IAmSensitiveBytes') + print(model) + #> password=SecretStr('**********') password_bytes=SecretBytes(b'**********') + print(model.password) + #> ********** + print(model.model_dump()) + ''' + { + 'password': SecretStr('**********'), + 'password_bytes': SecretBytes(b'**********'), + } + ''' + print(model.model_dump_json()) + #> {"password":"IAmSensitive","password_bytes":"IAmSensitiveBytes"} + ``` + """ + + _inner_schema: ClassVar[CoreSchema] = core_schema.str_schema() + _error_kind: ClassVar[str] = 'string_type' + + def __len__(self) -> int: + return len(self._secret_value) + + def _display(self) -> str: + return _secret_display(self._secret_value) + + +class SecretBytes(_SecretField[bytes]): + """A bytes used for storing sensitive information that you do not want to be visible in logging or tracebacks. + + It displays `b'**********'` instead of the string value on `repr()` and `str()` calls. + When the secret value is nonempty, it is displayed as `b'**********'` instead of the underlying value in + calls to `repr()` and `str()`. If the value _is_ empty, it is displayed as `b''`. + + ```python + from pydantic import BaseModel, SecretBytes + + class User(BaseModel): + username: str + password: SecretBytes + + user = User(username='scolvin', password=b'password1') + #> username='scolvin' password=SecretBytes(b'**********') + print(user.password.get_secret_value()) + #> b'password1' + print((SecretBytes(b'password'), SecretBytes(b''))) + #> (SecretBytes(b'**********'), SecretBytes(b'')) + ``` + """ + + _inner_schema: ClassVar[CoreSchema] = core_schema.bytes_schema() + _error_kind: ClassVar[str] = 'bytes_type' + + def __len__(self) -> int: + return len(self._secret_value) + + def _display(self) -> bytes: + return _secret_display(self._secret_value).encode() + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ PAYMENT CARD TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class PaymentCardBrand(str, Enum): + amex = 'American Express' + mastercard = 'Mastercard' + visa = 'Visa' + other = 'other' + + def __str__(self) -> str: + return self.value + + +@deprecated( + 'The `PaymentCardNumber` class is deprecated, use `pydantic_extra_types` instead. ' + 'See https://docs.pydantic.dev/latest/api/pydantic_extra_types_payment/#pydantic_extra_types.payment.PaymentCardNumber.', + category=PydanticDeprecatedSince20, +) +class PaymentCardNumber(str): + """Based on: https://en.wikipedia.org/wiki/Payment_card_number.""" + + strip_whitespace: ClassVar[bool] = True + min_length: ClassVar[int] = 12 + max_length: ClassVar[int] = 19 + bin: str + last4: str + brand: PaymentCardBrand + + def __init__(self, card_number: str): + self.validate_digits(card_number) + + card_number = self.validate_luhn_check_digit(card_number) + + self.bin = card_number[:6] + self.last4 = card_number[-4:] + self.brand = self.validate_brand(card_number) + + @classmethod + def __get_pydantic_core_schema__(cls, source: type[Any], handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + return core_schema.with_info_after_validator_function( + cls.validate, + core_schema.str_schema( + min_length=cls.min_length, max_length=cls.max_length, strip_whitespace=cls.strip_whitespace + ), + ) + + @classmethod + def validate(cls, input_value: str, /, _: core_schema.ValidationInfo) -> PaymentCardNumber: + """Validate the card number and return a `PaymentCardNumber` instance.""" + return cls(input_value) + + @property + def masked(self) -> str: + """Mask all but the last 4 digits of the card number. + + Returns: + A masked card number string. + """ + num_masked = len(self) - 10 # len(bin) + len(last4) == 10 + return f'{self.bin}{"*" * num_masked}{self.last4}' + + @classmethod + def validate_digits(cls, card_number: str) -> None: + """Validate that the card number is all digits.""" + if not card_number.isdigit(): + raise PydanticCustomError('payment_card_number_digits', 'Card number is not all digits') + + @classmethod + def validate_luhn_check_digit(cls, card_number: str) -> str: + """Based on: https://en.wikipedia.org/wiki/Luhn_algorithm.""" + sum_ = int(card_number[-1]) + length = len(card_number) + parity = length % 2 + for i in range(length - 1): + digit = int(card_number[i]) + if i % 2 == parity: + digit *= 2 + if digit > 9: + digit -= 9 + sum_ += digit + valid = sum_ % 10 == 0 + if not valid: + raise PydanticCustomError('payment_card_number_luhn', 'Card number is not luhn valid') + return card_number + + @staticmethod + def validate_brand(card_number: str) -> PaymentCardBrand: + """Validate length based on BIN for major brands: + https://en.wikipedia.org/wiki/Payment_card_number#Issuer_identification_number_(IIN). + """ + if card_number[0] == '4': + brand = PaymentCardBrand.visa + elif 51 <= int(card_number[:2]) <= 55: + brand = PaymentCardBrand.mastercard + elif card_number[:2] in {'34', '37'}: + brand = PaymentCardBrand.amex + else: + brand = PaymentCardBrand.other + + required_length: None | int | str = None + if brand in PaymentCardBrand.mastercard: + required_length = 16 + valid = len(card_number) == required_length + elif brand == PaymentCardBrand.visa: + required_length = '13, 16 or 19' + valid = len(card_number) in {13, 16, 19} + elif brand == PaymentCardBrand.amex: + required_length = 15 + valid = len(card_number) == required_length + else: + valid = True + + if not valid: + raise PydanticCustomError( + 'payment_card_number_brand', + 'Length for a {brand} card must be {required_length}', + {'brand': brand, 'required_length': required_length}, + ) + return brand + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ BYTE SIZE TYPE ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class ByteSize(int): + """Converts a string representing a number of bytes with units (such as `'1KB'` or `'11.5MiB'`) into an integer. + + You can use the `ByteSize` data type to (case-insensitively) convert a string representation of a number of bytes into + an integer, and also to print out human-readable strings representing a number of bytes. + + In conformance with [IEC 80000-13 Standard](https://en.wikipedia.org/wiki/ISO/IEC_80000) we interpret `'1KB'` to mean 1000 bytes, + and `'1KiB'` to mean 1024 bytes. In general, including a middle `'i'` will cause the unit to be interpreted as a power of 2, + rather than a power of 10 (so, for example, `'1 MB'` is treated as `1_000_000` bytes, whereas `'1 MiB'` is treated as `1_048_576` bytes). + + !!! info + Note that `1b` will be parsed as "1 byte" and not "1 bit". + + ```python + from pydantic import BaseModel, ByteSize + + class MyModel(BaseModel): + size: ByteSize + + print(MyModel(size=52000).size) + #> 52000 + print(MyModel(size='3000 KiB').size) + #> 3072000 + + m = MyModel(size='50 PB') + print(m.size.human_readable()) + #> 44.4PiB + print(m.size.human_readable(decimal=True)) + #> 50.0PB + print(m.size.human_readable(separator=' ')) + #> 44.4 PiB + + print(m.size.to('TiB')) + #> 45474.73508864641 + ``` + """ + + byte_sizes = { + 'b': 1, + 'kb': 10**3, + 'mb': 10**6, + 'gb': 10**9, + 'tb': 10**12, + 'pb': 10**15, + 'eb': 10**18, + 'kib': 2**10, + 'mib': 2**20, + 'gib': 2**30, + 'tib': 2**40, + 'pib': 2**50, + 'eib': 2**60, + 'bit': 1 / 8, + 'kbit': 10**3 / 8, + 'mbit': 10**6 / 8, + 'gbit': 10**9 / 8, + 'tbit': 10**12 / 8, + 'pbit': 10**15 / 8, + 'ebit': 10**18 / 8, + 'kibit': 2**10 / 8, + 'mibit': 2**20 / 8, + 'gibit': 2**30 / 8, + 'tibit': 2**40 / 8, + 'pibit': 2**50 / 8, + 'eibit': 2**60 / 8, + } + byte_sizes.update({k.lower()[0]: v for k, v in byte_sizes.items() if 'i' not in k}) + + byte_string_pattern = r'^\s*(\d*\.?\d+)\s*(\w+)?' + byte_string_re = re.compile(byte_string_pattern, re.IGNORECASE) + + @classmethod + def __get_pydantic_core_schema__(cls, source: type[Any], handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + return core_schema.with_info_after_validator_function( + function=cls._validate, + schema=core_schema.union_schema( + [ + core_schema.str_schema(pattern=cls.byte_string_pattern), + core_schema.int_schema(ge=0), + ], + custom_error_type='byte_size', + custom_error_message='could not parse value and unit from byte string', + ), + serialization=core_schema.plain_serializer_function_ser_schema( + int, return_schema=core_schema.int_schema(ge=0) + ), + ) + + @classmethod + def _validate(cls, input_value: Any, /, _: core_schema.ValidationInfo) -> ByteSize: + try: + return cls(int(input_value)) + except ValueError: + pass + + str_match = cls.byte_string_re.match(str(input_value)) + if str_match is None: + raise PydanticCustomError('byte_size', 'could not parse value and unit from byte string') + + scalar, unit = str_match.groups() + if unit is None: + unit = 'b' + + try: + unit_mult = cls.byte_sizes[unit.lower()] + except KeyError: + raise PydanticCustomError('byte_size_unit', 'could not interpret byte unit: {unit}', {'unit': unit}) + + return cls(int(float(scalar) * unit_mult)) + + def human_readable(self, decimal: bool = False, separator: str = '') -> str: + """Converts a byte size to a human readable string. + + Args: + decimal: If True, use decimal units (e.g. 1000 bytes per KB). If False, use binary units + (e.g. 1024 bytes per KiB). + separator: A string used to split the value and unit. Defaults to an empty string (''). + + Returns: + A human readable string representation of the byte size. + """ + if decimal: + divisor = 1000 + units = 'B', 'KB', 'MB', 'GB', 'TB', 'PB' + final_unit = 'EB' + else: + divisor = 1024 + units = 'B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB' + final_unit = 'EiB' + + num = float(self) + for unit in units: + if abs(num) < divisor: + if unit == 'B': + return f'{num:0.0f}{separator}{unit}' + else: + return f'{num:0.1f}{separator}{unit}' + num /= divisor + + return f'{num:0.1f}{separator}{final_unit}' + + def to(self, unit: str) -> float: + """Converts a byte size to another unit, including both byte and bit units. + + Args: + unit: The unit to convert to. Must be one of the following: B, KB, MB, GB, TB, PB, EB, + KiB, MiB, GiB, TiB, PiB, EiB (byte units) and + bit, kbit, mbit, gbit, tbit, pbit, ebit, + kibit, mibit, gibit, tibit, pibit, eibit (bit units). + + Returns: + The byte size in the new unit. + """ + try: + unit_div = self.byte_sizes[unit.lower()] + except KeyError: + raise PydanticCustomError('byte_size_unit', 'Could not interpret byte unit: {unit}', {'unit': unit}) + + return self / unit_div + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ DATE TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +def _check_annotated_type(annotated_type: str, expected_type: str, annotation: str) -> None: + if annotated_type != expected_type: + raise PydanticUserError(f"'{annotation}' cannot annotate '{annotated_type}'.", code='invalid-annotated-type') + + +if TYPE_CHECKING: + PastDate = Annotated[date, ...] + FutureDate = Annotated[date, ...] +else: + + class PastDate: + """A date in the past.""" + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[Any], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + if cls is source: + # used directly as a type + return core_schema.date_schema(now_op='past') + else: + schema = handler(source) + _check_annotated_type(schema['type'], 'date', cls.__name__) + schema['now_op'] = 'past' + return schema + + def __repr__(self) -> str: + return 'PastDate' + + class FutureDate: + """A date in the future.""" + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[Any], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + if cls is source: + # used directly as a type + return core_schema.date_schema(now_op='future') + else: + schema = handler(source) + _check_annotated_type(schema['type'], 'date', cls.__name__) + schema['now_op'] = 'future' + return schema + + def __repr__(self) -> str: + return 'FutureDate' + + +def condate( + *, + strict: bool | None = None, + gt: date | None = None, + ge: date | None = None, + lt: date | None = None, + le: date | None = None, +) -> type[date]: + """A wrapper for date that adds constraints. + + Args: + strict: Whether to validate the date value in strict mode. Defaults to `None`. + gt: The value must be greater than this. Defaults to `None`. + ge: The value must be greater than or equal to this. Defaults to `None`. + lt: The value must be less than this. Defaults to `None`. + le: The value must be less than or equal to this. Defaults to `None`. + + Returns: + A date type with the specified constraints. + """ + return Annotated[ # pyright: ignore[reportReturnType] + date, + Strict(strict) if strict is not None else None, + annotated_types.Interval(gt=gt, ge=ge, lt=lt, le=le), + ] + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ DATETIME TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +if TYPE_CHECKING: + AwareDatetime = Annotated[datetime, ...] + NaiveDatetime = Annotated[datetime, ...] + PastDatetime = Annotated[datetime, ...] + FutureDatetime = Annotated[datetime, ...] + +else: + + class AwareDatetime: + """A datetime that requires timezone info.""" + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[Any], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + if cls is source: + # used directly as a type + return core_schema.datetime_schema(tz_constraint='aware') + else: + schema = handler(source) + _check_annotated_type(schema['type'], 'datetime', cls.__name__) + schema['tz_constraint'] = 'aware' + return schema + + def __repr__(self) -> str: + return 'AwareDatetime' + + class NaiveDatetime: + """A datetime that doesn't require timezone info.""" + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[Any], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + if cls is source: + # used directly as a type + return core_schema.datetime_schema(tz_constraint='naive') + else: + schema = handler(source) + _check_annotated_type(schema['type'], 'datetime', cls.__name__) + schema['tz_constraint'] = 'naive' + return schema + + def __repr__(self) -> str: + return 'NaiveDatetime' + + class PastDatetime: + """A datetime that must be in the past.""" + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[Any], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + if cls is source: + # used directly as a type + return core_schema.datetime_schema(now_op='past') + else: + schema = handler(source) + _check_annotated_type(schema['type'], 'datetime', cls.__name__) + schema['now_op'] = 'past' + return schema + + def __repr__(self) -> str: + return 'PastDatetime' + + class FutureDatetime: + """A datetime that must be in the future.""" + + @classmethod + def __get_pydantic_core_schema__( + cls, source: type[Any], handler: GetCoreSchemaHandler + ) -> core_schema.CoreSchema: + if cls is source: + # used directly as a type + return core_schema.datetime_schema(now_op='future') + else: + schema = handler(source) + _check_annotated_type(schema['type'], 'datetime', cls.__name__) + schema['now_op'] = 'future' + return schema + + def __repr__(self) -> str: + return 'FutureDatetime' + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Encoded TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class EncoderProtocol(Protocol): + """Protocol for encoding and decoding data to and from bytes.""" + + @classmethod + def decode(cls, data: bytes) -> bytes: + """Decode the data using the encoder. + + Args: + data: The data to decode. + + Returns: + The decoded data. + """ + ... + + @classmethod + def encode(cls, value: bytes) -> bytes: + """Encode the data using the encoder. + + Args: + value: The data to encode. + + Returns: + The encoded data. + """ + ... + + @classmethod + def get_json_format(cls) -> str: + """Get the JSON format for the encoded data. + + Returns: + The JSON format for the encoded data. + """ + ... + + +class Base64Encoder(EncoderProtocol): + """Standard (non-URL-safe) Base64 encoder.""" + + @classmethod + def decode(cls, data: bytes) -> bytes: + """Decode the data from base64 encoded bytes to original bytes data. + + Args: + data: The data to decode. + + Returns: + The decoded data. + """ + try: + return base64.b64decode(data) + except ValueError as e: + raise PydanticCustomError('base64_decode', "Base64 decoding error: '{error}'", {'error': str(e)}) + + @classmethod + def encode(cls, value: bytes) -> bytes: + """Encode the data from bytes to a base64 encoded bytes. + + Args: + value: The data to encode. + + Returns: + The encoded data. + """ + return base64.b64encode(value) + + @classmethod + def get_json_format(cls) -> Literal['base64']: + """Get the JSON format for the encoded data. + + Returns: + The JSON format for the encoded data. + """ + return 'base64' + + +class Base64UrlEncoder(EncoderProtocol): + """URL-safe Base64 encoder.""" + + @classmethod + def decode(cls, data: bytes) -> bytes: + """Decode the data from base64 encoded bytes to original bytes data. + + Args: + data: The data to decode. + + Returns: + The decoded data. + """ + try: + return base64.urlsafe_b64decode(data) + except ValueError as e: + raise PydanticCustomError('base64_decode', "Base64 decoding error: '{error}'", {'error': str(e)}) + + @classmethod + def encode(cls, value: bytes) -> bytes: + """Encode the data from bytes to a base64 encoded bytes. + + Args: + value: The data to encode. + + Returns: + The encoded data. + """ + return base64.urlsafe_b64encode(value) + + @classmethod + def get_json_format(cls) -> Literal['base64url']: + """Get the JSON format for the encoded data. + + Returns: + The JSON format for the encoded data. + """ + return 'base64url' + + +@_dataclasses.dataclass(**_internal_dataclass.slots_true) +class EncodedBytes: + """A bytes type that is encoded and decoded using the specified encoder. + + `EncodedBytes` needs an encoder that implements `EncoderProtocol` to operate. + + ```python + from typing import Annotated + + from pydantic import BaseModel, EncodedBytes, EncoderProtocol, ValidationError + + class MyEncoder(EncoderProtocol): + @classmethod + def decode(cls, data: bytes) -> bytes: + if data == b'**undecodable**': + raise ValueError('Cannot decode data') + return data[13:] + + @classmethod + def encode(cls, value: bytes) -> bytes: + return b'**encoded**: ' + value + + @classmethod + def get_json_format(cls) -> str: + return 'my-encoder' + + MyEncodedBytes = Annotated[bytes, EncodedBytes(encoder=MyEncoder)] + + class Model(BaseModel): + my_encoded_bytes: MyEncodedBytes + + # Initialize the model with encoded data + m = Model(my_encoded_bytes=b'**encoded**: some bytes') + + # Access decoded value + print(m.my_encoded_bytes) + #> b'some bytes' + + # Serialize into the encoded form + print(m.model_dump()) + #> {'my_encoded_bytes': b'**encoded**: some bytes'} + + # Validate encoded data + try: + Model(my_encoded_bytes=b'**undecodable**') + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + my_encoded_bytes + Value error, Cannot decode data [type=value_error, input_value=b'**undecodable**', input_type=bytes] + ''' + ``` + """ + + encoder: type[EncoderProtocol] + + def __get_pydantic_json_schema__( + self, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = handler(core_schema) + field_schema.update(type='string', format=self.encoder.get_json_format()) + return field_schema + + def __get_pydantic_core_schema__(self, source: type[Any], handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + schema = handler(source) + _check_annotated_type(schema['type'], 'bytes', self.__class__.__name__) + return core_schema.with_info_after_validator_function( + function=self.decode, + schema=schema, + serialization=core_schema.plain_serializer_function_ser_schema(function=self.encode), + ) + + def decode(self, data: bytes, _: core_schema.ValidationInfo) -> bytes: + """Decode the data using the specified encoder. + + Args: + data: The data to decode. + + Returns: + The decoded data. + """ + return self.encoder.decode(data) + + def encode(self, value: bytes) -> bytes: + """Encode the data using the specified encoder. + + Args: + value: The data to encode. + + Returns: + The encoded data. + """ + return self.encoder.encode(value) + + def __hash__(self) -> int: + return hash(self.encoder) + + +@_dataclasses.dataclass(**_internal_dataclass.slots_true) +class EncodedStr: + """A str type that is encoded and decoded using the specified encoder. + + `EncodedStr` needs an encoder that implements `EncoderProtocol` to operate. + + ```python + from typing import Annotated + + from pydantic import BaseModel, EncodedStr, EncoderProtocol, ValidationError + + class MyEncoder(EncoderProtocol): + @classmethod + def decode(cls, data: bytes) -> bytes: + if data == b'**undecodable**': + raise ValueError('Cannot decode data') + return data[13:] + + @classmethod + def encode(cls, value: bytes) -> bytes: + return b'**encoded**: ' + value + + @classmethod + def get_json_format(cls) -> str: + return 'my-encoder' + + MyEncodedStr = Annotated[str, EncodedStr(encoder=MyEncoder)] + + class Model(BaseModel): + my_encoded_str: MyEncodedStr + + # Initialize the model with encoded data + m = Model(my_encoded_str='**encoded**: some str') + + # Access decoded value + print(m.my_encoded_str) + #> some str + + # Serialize into the encoded form + print(m.model_dump()) + #> {'my_encoded_str': '**encoded**: some str'} + + # Validate encoded data + try: + Model(my_encoded_str='**undecodable**') + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + my_encoded_str + Value error, Cannot decode data [type=value_error, input_value='**undecodable**', input_type=str] + ''' + ``` + """ + + encoder: type[EncoderProtocol] + + def __get_pydantic_json_schema__( + self, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler + ) -> JsonSchemaValue: + field_schema = handler(core_schema) + field_schema.update(type='string', format=self.encoder.get_json_format()) + return field_schema + + def __get_pydantic_core_schema__(self, source: type[Any], handler: GetCoreSchemaHandler) -> core_schema.CoreSchema: + schema = handler(source) + _check_annotated_type(schema['type'], 'str', self.__class__.__name__) + return core_schema.with_info_after_validator_function( + function=self.decode_str, + schema=schema, + serialization=core_schema.plain_serializer_function_ser_schema(function=self.encode_str), + ) + + def decode_str(self, data: str, _: core_schema.ValidationInfo) -> str: + """Decode the data using the specified encoder. + + Args: + data: The data to decode. + + Returns: + The decoded data. + """ + return self.encoder.decode(data.encode()).decode() + + def encode_str(self, value: str) -> str: + """Encode the data using the specified encoder. + + Args: + value: The data to encode. + + Returns: + The encoded data. + """ + return self.encoder.encode(value.encode()).decode() # noqa: UP008 + + def __hash__(self) -> int: + return hash(self.encoder) + + +Base64Bytes = Annotated[bytes, EncodedBytes(encoder=Base64Encoder)] +"""A bytes type that is encoded and decoded using the standard (non-URL-safe) base64 encoder. + +Note: + Under the hood, `Base64Bytes` uses the standard library `base64.b64encode` and `base64.b64decode` functions. + + As a result, attempting to decode url-safe base64 data using the `Base64Bytes` type may fail or produce an incorrect + decoding. + +Warning: + In versions of Pydantic prior to v2.10, `Base64Bytes` used [`base64.encodebytes`][base64.encodebytes] + and [`base64.decodebytes`][base64.decodebytes] functions. According to the [base64 documentation](https://docs.python.org/3/library/base64.html), + these methods are considered legacy implementation, and thus, Pydantic v2.10+ now uses the modern + [`base64.b64encode`][base64.b64encode] and [`base64.b64decode`][base64.b64decode] functions. + + If you'd still like to use these legacy encoders / decoders, you can achieve this by creating a custom annotated type, + like follows: + + ```python + import base64 + from typing import Annotated, Literal + + from pydantic_core import PydanticCustomError + + from pydantic import EncodedBytes, EncoderProtocol + + class LegacyBase64Encoder(EncoderProtocol): + @classmethod + def decode(cls, data: bytes) -> bytes: + try: + return base64.decodebytes(data) + except ValueError as e: + raise PydanticCustomError( + 'base64_decode', + "Base64 decoding error: '{error}'", + {'error': str(e)}, + ) + + @classmethod + def encode(cls, value: bytes) -> bytes: + return base64.encodebytes(value) + + @classmethod + def get_json_format(cls) -> Literal['base64']: + return 'base64' + + LegacyBase64Bytes = Annotated[bytes, EncodedBytes(encoder=LegacyBase64Encoder)] + ``` + +```python +from pydantic import Base64Bytes, BaseModel, ValidationError + +class Model(BaseModel): + base64_bytes: Base64Bytes + +# Initialize the model with base64 data +m = Model(base64_bytes=b'VGhpcyBpcyB0aGUgd2F5') + +# Access decoded value +print(m.base64_bytes) +#> b'This is the way' + +# Serialize into the base64 form +print(m.model_dump()) +#> {'base64_bytes': b'VGhpcyBpcyB0aGUgd2F5'} + +# Validate base64 data +try: + print(Model(base64_bytes=b'undecodable').base64_bytes) +except ValidationError as e: + print(e) + ''' + 1 validation error for Model + base64_bytes + Base64 decoding error: 'Incorrect padding' [type=base64_decode, input_value=b'undecodable', input_type=bytes] + ''' +``` +""" +Base64Str = Annotated[str, EncodedStr(encoder=Base64Encoder)] +"""A str type that is encoded and decoded using the standard (non-URL-safe) base64 encoder. + +Note: + Under the hood, `Base64Str` uses the standard library `base64.b64encode` and `base64.b64decode` functions. + + As a result, attempting to decode url-safe base64 data using the `Base64Str` type may fail or produce an incorrect + decoding. + +Warning: + In versions of Pydantic prior to v2.10, `Base64Str` used [`base64.encodebytes`][base64.encodebytes] + and [`base64.decodebytes`][base64.decodebytes] functions. According to the [base64 documentation](https://docs.python.org/3/library/base64.html), + these methods are considered legacy implementation, and thus, Pydantic v2.10+ now uses the modern + [`base64.b64encode`][base64.b64encode] and [`base64.b64decode`][base64.b64decode] functions. + + See the [`Base64Bytes`][pydantic.types.Base64Bytes] type for more information on how to + replicate the old behavior with the legacy encoders / decoders. + +```python +from pydantic import Base64Str, BaseModel, ValidationError + +class Model(BaseModel): + base64_str: Base64Str + +# Initialize the model with base64 data +m = Model(base64_str='VGhlc2UgYXJlbid0IHRoZSBkcm9pZHMgeW91J3JlIGxvb2tpbmcgZm9y') + +# Access decoded value +print(m.base64_str) +#> These aren't the droids you're looking for + +# Serialize into the base64 form +print(m.model_dump()) +#> {'base64_str': 'VGhlc2UgYXJlbid0IHRoZSBkcm9pZHMgeW91J3JlIGxvb2tpbmcgZm9y'} + +# Validate base64 data +try: + print(Model(base64_str='undecodable').base64_str) +except ValidationError as e: + print(e) + ''' + 1 validation error for Model + base64_str + Base64 decoding error: 'Incorrect padding' [type=base64_decode, input_value='undecodable', input_type=str] + ''' +``` +""" +Base64UrlBytes = Annotated[bytes, EncodedBytes(encoder=Base64UrlEncoder)] +"""A bytes type that is encoded and decoded using the URL-safe base64 encoder. + +Note: + Under the hood, `Base64UrlBytes` use standard library `base64.urlsafe_b64encode` and `base64.urlsafe_b64decode` + functions. + + As a result, the `Base64UrlBytes` type can be used to faithfully decode "vanilla" base64 data + (using `'+'` and `'/'`). + +```python +from pydantic import Base64UrlBytes, BaseModel + +class Model(BaseModel): + base64url_bytes: Base64UrlBytes + +# Initialize the model with base64 data +m = Model(base64url_bytes=b'SHc_dHc-TXc==') +print(m) +#> base64url_bytes=b'Hw?tw>Mw' +``` +""" +Base64UrlStr = Annotated[str, EncodedStr(encoder=Base64UrlEncoder)] +"""A str type that is encoded and decoded using the URL-safe base64 encoder. + +Note: + Under the hood, `Base64UrlStr` use standard library `base64.urlsafe_b64encode` and `base64.urlsafe_b64decode` + functions. + + As a result, the `Base64UrlStr` type can be used to faithfully decode "vanilla" base64 data (using `'+'` and `'/'`). + +```python +from pydantic import Base64UrlStr, BaseModel + +class Model(BaseModel): + base64url_str: Base64UrlStr + +# Initialize the model with base64 data +m = Model(base64url_str='SHc_dHc-TXc==') +print(m) +#> base64url_str='Hw?tw>Mw' +``` +""" + + +__getattr__ = getattr_migration(__name__) + + +@_dataclasses.dataclass(**_internal_dataclass.slots_true) +class GetPydanticSchema: + """!!! abstract "Usage Documentation" + [Using `GetPydanticSchema` to Reduce Boilerplate](../concepts/types.md#using-getpydanticschema-to-reduce-boilerplate) + + A convenience class for creating an annotation that provides pydantic custom type hooks. + + This class is intended to eliminate the need to create a custom "marker" which defines the + `__get_pydantic_core_schema__` and `__get_pydantic_json_schema__` custom hook methods. + + For example, to have a field treated by type checkers as `int`, but by pydantic as `Any`, you can do: + ```python + from typing import Annotated, Any + + from pydantic import BaseModel, GetPydanticSchema + + HandleAsAny = GetPydanticSchema(lambda _s, h: h(Any)) + + class Model(BaseModel): + x: Annotated[int, HandleAsAny] # pydantic sees `x: Any` + + print(repr(Model(x='abc').x)) + #> 'abc' + ``` + """ + + get_pydantic_core_schema: Callable[[Any, GetCoreSchemaHandler], CoreSchema] | None = None + get_pydantic_json_schema: Callable[[Any, GetJsonSchemaHandler], JsonSchemaValue] | None = None + + # Note: we may want to consider adding a convenience staticmethod `def for_type(type_: Any) -> GetPydanticSchema:` + # which returns `GetPydanticSchema(lambda _s, h: h(type_))` + + if not TYPE_CHECKING: + # We put `__getattr__` in a non-TYPE_CHECKING block because otherwise, mypy allows arbitrary attribute access + + def __getattr__(self, item: str) -> Any: + """Use this rather than defining `__get_pydantic_core_schema__` etc. to reduce the number of nested calls.""" + if item == '__get_pydantic_core_schema__' and self.get_pydantic_core_schema: + return self.get_pydantic_core_schema + elif item == '__get_pydantic_json_schema__' and self.get_pydantic_json_schema: + return self.get_pydantic_json_schema + else: + return object.__getattribute__(self, item) + + __hash__ = object.__hash__ + + +@_dataclasses.dataclass(**_internal_dataclass.slots_true, frozen=True) +class Tag: + """Provides a way to specify the expected tag to use for a case of a (callable) discriminated union. + + Also provides a way to label a union case in error messages. + + When using a callable `Discriminator`, attach a `Tag` to each case in the `Union` to specify the tag that + should be used to identify that case. For example, in the below example, the `Tag` is used to specify that + if `get_discriminator_value` returns `'apple'`, the input should be validated as an `ApplePie`, and if it + returns `'pumpkin'`, the input should be validated as a `PumpkinPie`. + + The primary role of the `Tag` here is to map the return value from the callable `Discriminator` function to + the appropriate member of the `Union` in question. + + ```python + from typing import Annotated, Any, Literal, Union + + from pydantic import BaseModel, Discriminator, Tag + + class Pie(BaseModel): + time_to_cook: int + num_ingredients: int + + class ApplePie(Pie): + fruit: Literal['apple'] = 'apple' + + class PumpkinPie(Pie): + filling: Literal['pumpkin'] = 'pumpkin' + + def get_discriminator_value(v: Any) -> str: + if isinstance(v, dict): + return v.get('fruit', v.get('filling')) + return getattr(v, 'fruit', getattr(v, 'filling', None)) + + class ThanksgivingDinner(BaseModel): + dessert: Annotated[ + Union[ + Annotated[ApplePie, Tag('apple')], + Annotated[PumpkinPie, Tag('pumpkin')], + ], + Discriminator(get_discriminator_value), + ] + + apple_variation = ThanksgivingDinner.model_validate( + {'dessert': {'fruit': 'apple', 'time_to_cook': 60, 'num_ingredients': 8}} + ) + print(repr(apple_variation)) + ''' + ThanksgivingDinner(dessert=ApplePie(time_to_cook=60, num_ingredients=8, fruit='apple')) + ''' + + pumpkin_variation = ThanksgivingDinner.model_validate( + { + 'dessert': { + 'filling': 'pumpkin', + 'time_to_cook': 40, + 'num_ingredients': 6, + } + } + ) + print(repr(pumpkin_variation)) + ''' + ThanksgivingDinner(dessert=PumpkinPie(time_to_cook=40, num_ingredients=6, filling='pumpkin')) + ''' + ``` + + !!! note + You must specify a `Tag` for every case in a `Tag` that is associated with a + callable `Discriminator`. Failing to do so will result in a `PydanticUserError` with code + [`callable-discriminator-no-tag`](../errors/usage_errors.md#callable-discriminator-no-tag). + + See the [Discriminated Unions] concepts docs for more details on how to use `Tag`s. + + [Discriminated Unions]: ../concepts/unions.md#discriminated-unions + """ + + tag: str + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema: + schema = handler(source_type) + metadata = cast('CoreMetadata', schema.setdefault('metadata', {})) + metadata['pydantic_internal_union_tag_key'] = self.tag + return schema + + +@_dataclasses.dataclass(**_internal_dataclass.slots_true, frozen=True) +class Discriminator: + """!!! abstract "Usage Documentation" + [Discriminated Unions with `Callable` `Discriminator`](../concepts/unions.md#discriminated-unions-with-callable-discriminator) + + Provides a way to use a custom callable as the way to extract the value of a union discriminator. + + This allows you to get validation behavior like you'd get from `Field(discriminator=)`, + but without needing to have a single shared field across all the union choices. This also makes it + possible to handle unions of models and primitive types with discriminated-union-style validation errors. + Finally, this allows you to use a custom callable as the way to identify which member of a union a value + belongs to, while still seeing all the performance benefits of a discriminated union. + + Consider this example, which is much more performant with the use of `Discriminator` and thus a `TaggedUnion` + than it would be as a normal `Union`. + + ```python + from typing import Annotated, Any, Literal, Union + + from pydantic import BaseModel, Discriminator, Tag + + class Pie(BaseModel): + time_to_cook: int + num_ingredients: int + + class ApplePie(Pie): + fruit: Literal['apple'] = 'apple' + + class PumpkinPie(Pie): + filling: Literal['pumpkin'] = 'pumpkin' + + def get_discriminator_value(v: Any) -> str: + if isinstance(v, dict): + return v.get('fruit', v.get('filling')) + return getattr(v, 'fruit', getattr(v, 'filling', None)) + + class ThanksgivingDinner(BaseModel): + dessert: Annotated[ + Union[ + Annotated[ApplePie, Tag('apple')], + Annotated[PumpkinPie, Tag('pumpkin')], + ], + Discriminator(get_discriminator_value), + ] + + apple_variation = ThanksgivingDinner.model_validate( + {'dessert': {'fruit': 'apple', 'time_to_cook': 60, 'num_ingredients': 8}} + ) + print(repr(apple_variation)) + ''' + ThanksgivingDinner(dessert=ApplePie(time_to_cook=60, num_ingredients=8, fruit='apple')) + ''' + + pumpkin_variation = ThanksgivingDinner.model_validate( + { + 'dessert': { + 'filling': 'pumpkin', + 'time_to_cook': 40, + 'num_ingredients': 6, + } + } + ) + print(repr(pumpkin_variation)) + ''' + ThanksgivingDinner(dessert=PumpkinPie(time_to_cook=40, num_ingredients=6, filling='pumpkin')) + ''' + ``` + + See the [Discriminated Unions] concepts docs for more details on how to use `Discriminator`s. + + [Discriminated Unions]: ../concepts/unions.md#discriminated-unions + """ + + discriminator: str | Callable[[Any], Hashable] + """The callable or field name for discriminating the type in a tagged union. + + A `Callable` discriminator must extract the value of the discriminator from the input. + A `str` discriminator must be the name of a field to discriminate against. + """ + custom_error_type: str | None = None + """Type to use in [custom errors](../errors/errors.md) replacing the standard discriminated union + validation errors. + """ + custom_error_message: str | None = None + """Message to use in custom errors.""" + custom_error_context: dict[str, int | str | float] | None = None + """Context to use in custom errors.""" + + def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema: + if not is_union_origin(get_origin(source_type)): + raise TypeError(f'{type(self).__name__} must be used with a Union type, not {source_type}') + + if isinstance(self.discriminator, str): + from pydantic import Field + + return handler(Annotated[source_type, Field(discriminator=self.discriminator)]) + else: + original_schema = handler(source_type) + return self._convert_schema(original_schema) + + def _convert_schema(self, original_schema: core_schema.CoreSchema) -> core_schema.TaggedUnionSchema: + if original_schema['type'] != 'union': + # This likely indicates that the schema was a single-item union that was simplified. + # In this case, we do the same thing we do in + # `pydantic._internal._discriminated_union._ApplyInferredDiscriminator._apply_to_root`, namely, + # package the generated schema back into a single-item union. + original_schema = core_schema.union_schema([original_schema]) + + tagged_union_choices = {} + for choice in original_schema['choices']: + tag = None + if isinstance(choice, tuple): + choice, tag = choice + metadata = cast('CoreMetadata | None', choice.get('metadata')) + if metadata is not None: + tag = metadata.get('pydantic_internal_union_tag_key') or tag + if tag is None: + raise PydanticUserError( + f'`Tag` not provided for choice {choice} used with `Discriminator`', + code='callable-discriminator-no-tag', + ) + tagged_union_choices[tag] = choice + + # Have to do these verbose checks to ensure falsy values ('' and {}) don't get ignored + custom_error_type = self.custom_error_type + if custom_error_type is None: + custom_error_type = original_schema.get('custom_error_type') + + custom_error_message = self.custom_error_message + if custom_error_message is None: + custom_error_message = original_schema.get('custom_error_message') + + custom_error_context = self.custom_error_context + if custom_error_context is None: + custom_error_context = original_schema.get('custom_error_context') + + custom_error_type = original_schema.get('custom_error_type') if custom_error_type is None else custom_error_type + return core_schema.tagged_union_schema( + tagged_union_choices, + self.discriminator, + custom_error_type=custom_error_type, + custom_error_message=custom_error_message, + custom_error_context=custom_error_context, + strict=original_schema.get('strict'), + ref=original_schema.get('ref'), + metadata=original_schema.get('metadata'), + serialization=original_schema.get('serialization'), + ) + + +_JSON_TYPES = {int, float, str, bool, list, dict, type(None)} + + +def _get_type_name(x: Any) -> str: + type_ = type(x) + if type_ in _JSON_TYPES: + return type_.__name__ + + # Handle proper subclasses; note we don't need to handle None or bool here + if isinstance(x, int): + return 'int' + if isinstance(x, float): + return 'float' + if isinstance(x, str): + return 'str' + if isinstance(x, list): + return 'list' + if isinstance(x, dict): + return 'dict' + + # Fail by returning the type's actual name + return getattr(type_, '__name__', '') + + +class _AllowAnyJson: + @classmethod + def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema: + python_schema = handler(source_type) + return core_schema.json_or_python_schema(json_schema=core_schema.any_schema(), python_schema=python_schema) + + +if TYPE_CHECKING: + # This seems to only be necessary for mypy + JsonValue: TypeAlias = Union[ + list['JsonValue'], + dict[str, 'JsonValue'], + str, + bool, + int, + float, + None, + ] + """A `JsonValue` is used to represent a value that can be serialized to JSON. + + It may be one of: + + * `list['JsonValue']` + * `dict[str, 'JsonValue']` + * `str` + * `bool` + * `int` + * `float` + * `None` + + The following example demonstrates how to use `JsonValue` to validate JSON data, + and what kind of errors to expect when input data is not json serializable. + + ```python + import json + + from pydantic import BaseModel, JsonValue, ValidationError + + class Model(BaseModel): + j: JsonValue + + valid_json_data = {'j': {'a': {'b': {'c': 1, 'd': [2, None]}}}} + invalid_json_data = {'j': {'a': {'b': ...}}} + + print(repr(Model.model_validate(valid_json_data))) + #> Model(j={'a': {'b': {'c': 1, 'd': [2, None]}}}) + print(repr(Model.model_validate_json(json.dumps(valid_json_data)))) + #> Model(j={'a': {'b': {'c': 1, 'd': [2, None]}}}) + + try: + Model.model_validate(invalid_json_data) + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + j.dict.a.dict.b + input was not a valid JSON value [type=invalid-json-value, input_value=Ellipsis, input_type=ellipsis] + ''' + ``` + """ + +else: + JsonValue = TypeAliasType( + 'JsonValue', + Annotated[ + Union[ + Annotated[list['JsonValue'], Tag('list')], + Annotated[dict[str, 'JsonValue'], Tag('dict')], + Annotated[str, Tag('str')], + Annotated[bool, Tag('bool')], + Annotated[int, Tag('int')], + Annotated[float, Tag('float')], + Annotated[None, Tag('NoneType')], + ], + Discriminator( + _get_type_name, + custom_error_type='invalid-json-value', + custom_error_message='input was not a valid JSON value', + ), + _AllowAnyJson, + ], + ) + + +class _OnErrorOmit: + @classmethod + def __get_pydantic_core_schema__(cls, source_type: Any, handler: GetCoreSchemaHandler) -> CoreSchema: + # there is no actual default value here but we use with_default_schema since it already has the on_error + # behavior implemented and it would be no more efficient to implement it on every other validator + # or as a standalone validator + return core_schema.with_default_schema(schema=handler(source_type), on_error='omit') + + +OnErrorOmit = Annotated[T, _OnErrorOmit] +""" +When used as an item in a list, the key type in a dict, optional values of a TypedDict, etc. +this annotation omits the item from the iteration if there is any error validating it. +That is, instead of a [`ValidationError`][pydantic_core.ValidationError] being propagated up and the entire iterable being discarded +any invalid items are discarded and the valid ones are returned. +""" + + +@_dataclasses.dataclass +class FailFast(_fields.PydanticMetadata, BaseMetadata): + """A `FailFast` annotation can be used to specify that validation should stop at the first error. + + This can be useful when you want to validate a large amount of data and you only need to know if it's valid or not. + + You might want to enable this setting if you want to validate your data faster (basically, if you use this, + validation will be more performant with the caveat that you get less information). + + ```python + from typing import Annotated + + from pydantic import BaseModel, FailFast, ValidationError + + class Model(BaseModel): + x: Annotated[list[int], FailFast()] + + # This will raise a single error for the first invalid value and stop validation + try: + obj = Model(x=[1, 2, 'a', 4, 5, 'b', 7, 8, 9, 'c']) + except ValidationError as e: + print(e) + ''' + 1 validation error for Model + x.2 + Input should be a valid integer, unable to parse string as an integer [type=int_parsing, input_value='a', input_type=str] + ''' + ``` + """ + + fail_fast: bool = True diff --git a/lib/python3.12/site-packages/pydantic/typing.py b/lib/python3.12/site-packages/pydantic/typing.py new file mode 100644 index 0000000000000000000000000000000000000000..0bda22d02ffc59f306ded475805343aef2855f49 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/typing.py @@ -0,0 +1,5 @@ +"""`typing` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/utils.py b/lib/python3.12/site-packages/pydantic/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8d1e2a81c5f0d1bd858bf12ccf15b1461008cf99 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/utils.py @@ -0,0 +1,5 @@ +"""The `utils` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/v1/__init__.py b/lib/python3.12/site-packages/pydantic/v1/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6ad3f466684df5c377480a8e181971a1b2f0016a --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/__init__.py @@ -0,0 +1,131 @@ +# flake8: noqa +from pydantic.v1 import dataclasses +from pydantic.v1.annotated_types import create_model_from_namedtuple, create_model_from_typeddict +from pydantic.v1.class_validators import root_validator, validator +from pydantic.v1.config import BaseConfig, ConfigDict, Extra +from pydantic.v1.decorator import validate_arguments +from pydantic.v1.env_settings import BaseSettings +from pydantic.v1.error_wrappers import ValidationError +from pydantic.v1.errors import * +from pydantic.v1.fields import Field, PrivateAttr, Required +from pydantic.v1.main import * +from pydantic.v1.networks import * +from pydantic.v1.parse import Protocol +from pydantic.v1.tools import * +from pydantic.v1.types import * +from pydantic.v1.version import VERSION, compiled + +__version__ = VERSION + +# WARNING __all__ from pydantic.errors is not included here, it will be removed as an export here in v2 +# please use "from pydantic.v1.errors import ..." instead +__all__ = [ + # annotated types utils + 'create_model_from_namedtuple', + 'create_model_from_typeddict', + # dataclasses + 'dataclasses', + # class_validators + 'root_validator', + 'validator', + # config + 'BaseConfig', + 'ConfigDict', + 'Extra', + # decorator + 'validate_arguments', + # env_settings + 'BaseSettings', + # error_wrappers + 'ValidationError', + # fields + 'Field', + 'Required', + # main + 'BaseModel', + 'create_model', + 'validate_model', + # network + 'AnyUrl', + 'AnyHttpUrl', + 'FileUrl', + 'HttpUrl', + 'stricturl', + 'EmailStr', + 'NameEmail', + 'IPvAnyAddress', + 'IPvAnyInterface', + 'IPvAnyNetwork', + 'PostgresDsn', + 'CockroachDsn', + 'AmqpDsn', + 'RedisDsn', + 'MongoDsn', + 'KafkaDsn', + 'validate_email', + # parse + 'Protocol', + # tools + 'parse_file_as', + 'parse_obj_as', + 'parse_raw_as', + 'schema_of', + 'schema_json_of', + # types + 'NoneStr', + 'NoneBytes', + 'StrBytes', + 'NoneStrBytes', + 'StrictStr', + 'ConstrainedBytes', + 'conbytes', + 'ConstrainedList', + 'conlist', + 'ConstrainedSet', + 'conset', + 'ConstrainedFrozenSet', + 'confrozenset', + 'ConstrainedStr', + 'constr', + 'PyObject', + 'ConstrainedInt', + 'conint', + 'PositiveInt', + 'NegativeInt', + 'NonNegativeInt', + 'NonPositiveInt', + 'ConstrainedFloat', + 'confloat', + 'PositiveFloat', + 'NegativeFloat', + 'NonNegativeFloat', + 'NonPositiveFloat', + 'FiniteFloat', + 'ConstrainedDecimal', + 'condecimal', + 'ConstrainedDate', + 'condate', + 'UUID1', + 'UUID3', + 'UUID4', + 'UUID5', + 'FilePath', + 'DirectoryPath', + 'Json', + 'JsonWrapper', + 'SecretField', + 'SecretStr', + 'SecretBytes', + 'StrictBool', + 'StrictBytes', + 'StrictInt', + 'StrictFloat', + 'PaymentCardNumber', + 'PrivateAttr', + 'ByteSize', + 'PastDate', + 'FutureDate', + # version + 'compiled', + 'VERSION', +] diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e56a9f142723216bafad6fb1cfeb5cf66c76740 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/_hypothesis_plugin.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/_hypothesis_plugin.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fea3d8d9bfb8a04af52f3606edf3b8e9bd7da4bd Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/_hypothesis_plugin.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/annotated_types.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/annotated_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..016e61fa29c8c6f45efa5b838e8c90262e500e20 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/annotated_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/class_validators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/class_validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ebe6467eb43318a0c63fafd0efd54094120a476 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/class_validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/color.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/color.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..190005fa6015b503ad3d1691406ee58cc48762b3 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/color.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b447524dc24541b1eb87e2b2f7e4763d38342fcb Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/dataclasses.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/dataclasses.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..178b03bbf461b8550318dc8c748718b41b6e7b6e Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/dataclasses.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/datetime_parse.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/datetime_parse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..974f8a26e92b4f3adacf1d2240a3669e06d12b1b Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/datetime_parse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/decorator.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/decorator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb1b29afc894ed8d3e251f3c1390521dd0d9c64a Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/decorator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/env_settings.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/env_settings.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb738e2ff4269279b6267f5cfc2fa5a136300727 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/env_settings.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/error_wrappers.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/error_wrappers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f6e01fb8702b2431f8c045f58c0272bbfed3973 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/error_wrappers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/errors.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/errors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..17310078d142ae0de3988da1d885672a6fa0a382 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/errors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/fields.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/fields.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee241c2af877ed3aa4c04a53f35efd2b7c92863e Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/fields.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/generics.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/generics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d830542ac1105a41542a4cf66d276905d03c1035 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/generics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/json.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/json.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..669b5b4ce285b6ee2d7532689f4aa304a423be09 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/json.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/main.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/main.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..969408add0ee35db92630ad7c66544df7dce5ac8 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/main.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/mypy.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/mypy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..afa53a9578cba857d98c0f512237fa1c9d7638ee Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/mypy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/networks.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/networks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b23e785664df9b39800549503f5c2eaeb686bf0 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/networks.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/parse.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/parse.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2662517a72b6b2c562e0a1dae3e956c4261230d1 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/parse.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/schema.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/schema.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7659de50d3ce0d5c1a4b14b518744d8d13a34c72 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/schema.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/tools.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/tools.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2285cf81b53d8c57039f5ea3e353ac49b597c43a Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/tools.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a6a8259efbc5cd96eaa674ac1d36db841fb3ec7 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/typing.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/typing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b2594058f5a434bd5db424d1e9f4361c5c32377d Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/typing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d70aed470bd34ebf6de6ac16511ec7369cab4c1d Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/validators.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1280034ce7bb0bbf608d664dc441ed43f4d42fc0 Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/__pycache__/version.cpython-312.pyc b/lib/python3.12/site-packages/pydantic/v1/__pycache__/version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9b8b70d82a31dd4d61632586e7ddaf7c9fe075e Binary files /dev/null and b/lib/python3.12/site-packages/pydantic/v1/__pycache__/version.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pydantic/v1/_hypothesis_plugin.py b/lib/python3.12/site-packages/pydantic/v1/_hypothesis_plugin.py new file mode 100644 index 0000000000000000000000000000000000000000..b62234d50750cbcf6d3ebba02d2ee54afb8c7347 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/_hypothesis_plugin.py @@ -0,0 +1,391 @@ +""" +Register Hypothesis strategies for Pydantic custom types. + +This enables fully-automatic generation of test data for most Pydantic classes. + +Note that this module has *no* runtime impact on Pydantic itself; instead it +is registered as a setuptools entry point and Hypothesis will import it if +Pydantic is installed. See also: + +https://hypothesis.readthedocs.io/en/latest/strategies.html#registering-strategies-via-setuptools-entry-points +https://hypothesis.readthedocs.io/en/latest/data.html#hypothesis.strategies.register_type_strategy +https://hypothesis.readthedocs.io/en/latest/strategies.html#interaction-with-pytest-cov +https://docs.pydantic.dev/usage/types/#pydantic-types + +Note that because our motivation is to *improve user experience*, the strategies +are always sound (never generate invalid data) but sacrifice completeness for +maintainability (ie may be unable to generate some tricky but valid data). + +Finally, this module makes liberal use of `# type: ignore[]` pragmas. +This is because Hypothesis annotates `register_type_strategy()` with +`(T, SearchStrategy[T])`, but in most cases we register e.g. `ConstrainedInt` +to generate instances of the builtin `int` type which match the constraints. +""" + +import contextlib +import datetime +import ipaddress +import json +import math +from fractions import Fraction +from typing import Callable, Dict, Type, Union, cast, overload + +import hypothesis.strategies as st + +import pydantic +import pydantic.color +import pydantic.types +from pydantic.v1.utils import lenient_issubclass + +# FilePath and DirectoryPath are explicitly unsupported, as we'd have to create +# them on-disk, and that's unsafe in general without being told *where* to do so. +# +# URLs are unsupported because it's easy for users to define their own strategy for +# "normal" URLs, and hard for us to define a general strategy which includes "weird" +# URLs but doesn't also have unpredictable performance problems. +# +# conlist() and conset() are unsupported for now, because the workarounds for +# Cython and Hypothesis to handle parametrized generic types are incompatible. +# We are rethinking Hypothesis compatibility in Pydantic v2. + +# Emails +try: + import email_validator +except ImportError: # pragma: no cover + pass +else: + + def is_valid_email(s: str) -> bool: + # Hypothesis' st.emails() occasionally generates emails like 0@A0--0.ac + # that are invalid according to email-validator, so we filter those out. + try: + email_validator.validate_email(s, check_deliverability=False) + return True + except email_validator.EmailNotValidError: # pragma: no cover + return False + + # Note that these strategies deliberately stay away from any tricky Unicode + # or other encoding issues; we're just trying to generate *something* valid. + st.register_type_strategy(pydantic.EmailStr, st.emails().filter(is_valid_email)) # type: ignore[arg-type] + st.register_type_strategy( + pydantic.NameEmail, + st.builds( + '{} <{}>'.format, # type: ignore[arg-type] + st.from_regex('[A-Za-z0-9_]+( [A-Za-z0-9_]+){0,5}', fullmatch=True), + st.emails().filter(is_valid_email), + ), + ) + +# PyObject - dotted names, in this case taken from the math module. +st.register_type_strategy( + pydantic.PyObject, # type: ignore[arg-type] + st.sampled_from( + [cast(pydantic.PyObject, f'math.{name}') for name in sorted(vars(math)) if not name.startswith('_')] + ), +) + +# CSS3 Colors; as name, hex, rgb(a) tuples or strings, or hsl strings +_color_regexes = ( + '|'.join( + ( + pydantic.color.r_hex_short, + pydantic.color.r_hex_long, + pydantic.color.r_rgb, + pydantic.color.r_rgba, + pydantic.color.r_hsl, + pydantic.color.r_hsla, + ) + ) + # Use more precise regex patterns to avoid value-out-of-range errors + .replace(pydantic.color._r_sl, r'(?:(\d\d?(?:\.\d+)?|100(?:\.0+)?)%)') + .replace(pydantic.color._r_alpha, r'(?:(0(?:\.\d+)?|1(?:\.0+)?|\.\d+|\d{1,2}%))') + .replace(pydantic.color._r_255, r'(?:((?:\d|\d\d|[01]\d\d|2[0-4]\d|25[0-4])(?:\.\d+)?|255(?:\.0+)?))') +) +st.register_type_strategy( + pydantic.color.Color, + st.one_of( + st.sampled_from(sorted(pydantic.color.COLORS_BY_NAME)), + st.tuples( + st.integers(0, 255), + st.integers(0, 255), + st.integers(0, 255), + st.none() | st.floats(0, 1) | st.floats(0, 100).map('{}%'.format), + ), + st.from_regex(_color_regexes, fullmatch=True), + ), +) + + +# Card numbers, valid according to the Luhn algorithm + + +def add_luhn_digit(card_number: str) -> str: + # See https://en.wikipedia.org/wiki/Luhn_algorithm + for digit in '0123456789': + with contextlib.suppress(Exception): + pydantic.PaymentCardNumber.validate_luhn_check_digit(card_number + digit) + return card_number + digit + raise AssertionError('Unreachable') # pragma: no cover + + +card_patterns = ( + # Note that these patterns omit the Luhn check digit; that's added by the function above + '4[0-9]{14}', # Visa + '5[12345][0-9]{13}', # Mastercard + '3[47][0-9]{12}', # American Express + '[0-26-9][0-9]{10,17}', # other (incomplete to avoid overlap) +) +st.register_type_strategy( + pydantic.PaymentCardNumber, + st.from_regex('|'.join(card_patterns), fullmatch=True).map(add_luhn_digit), # type: ignore[arg-type] +) + +# UUIDs +st.register_type_strategy(pydantic.UUID1, st.uuids(version=1)) +st.register_type_strategy(pydantic.UUID3, st.uuids(version=3)) +st.register_type_strategy(pydantic.UUID4, st.uuids(version=4)) +st.register_type_strategy(pydantic.UUID5, st.uuids(version=5)) + +# Secrets +st.register_type_strategy(pydantic.SecretBytes, st.binary().map(pydantic.SecretBytes)) +st.register_type_strategy(pydantic.SecretStr, st.text().map(pydantic.SecretStr)) + +# IP addresses, networks, and interfaces +st.register_type_strategy(pydantic.IPvAnyAddress, st.ip_addresses()) # type: ignore[arg-type] +st.register_type_strategy( + pydantic.IPvAnyInterface, + st.from_type(ipaddress.IPv4Interface) | st.from_type(ipaddress.IPv6Interface), # type: ignore[arg-type] +) +st.register_type_strategy( + pydantic.IPvAnyNetwork, + st.from_type(ipaddress.IPv4Network) | st.from_type(ipaddress.IPv6Network), # type: ignore[arg-type] +) + +# We hook into the con***() functions and the ConstrainedNumberMeta metaclass, +# so here we only have to register subclasses for other constrained types which +# don't go via those mechanisms. Then there are the registration hooks below. +st.register_type_strategy(pydantic.StrictBool, st.booleans()) +st.register_type_strategy(pydantic.StrictStr, st.text()) + + +# FutureDate, PastDate +st.register_type_strategy(pydantic.FutureDate, st.dates(min_value=datetime.date.today() + datetime.timedelta(days=1))) +st.register_type_strategy(pydantic.PastDate, st.dates(max_value=datetime.date.today() - datetime.timedelta(days=1))) + + +# Constrained-type resolver functions +# +# For these ones, we actually want to inspect the type in order to work out a +# satisfying strategy. First up, the machinery for tracking resolver functions: + +RESOLVERS: Dict[type, Callable[[type], st.SearchStrategy]] = {} # type: ignore[type-arg] + + +@overload +def _registered(typ: Type[pydantic.types.T]) -> Type[pydantic.types.T]: + pass + + +@overload +def _registered(typ: pydantic.types.ConstrainedNumberMeta) -> pydantic.types.ConstrainedNumberMeta: + pass + + +def _registered( + typ: Union[Type[pydantic.types.T], pydantic.types.ConstrainedNumberMeta] +) -> Union[Type[pydantic.types.T], pydantic.types.ConstrainedNumberMeta]: + # This function replaces the version in `pydantic.types`, in order to + # effect the registration of new constrained types so that Hypothesis + # can generate valid examples. + pydantic.types._DEFINED_TYPES.add(typ) + for supertype, resolver in RESOLVERS.items(): + if issubclass(typ, supertype): + st.register_type_strategy(typ, resolver(typ)) # type: ignore + return typ + raise NotImplementedError(f'Unknown type {typ!r} has no resolver to register') # pragma: no cover + + +def resolves( + typ: Union[type, pydantic.types.ConstrainedNumberMeta] +) -> Callable[[Callable[..., st.SearchStrategy]], Callable[..., st.SearchStrategy]]: # type: ignore[type-arg] + def inner(f): # type: ignore + assert f not in RESOLVERS + RESOLVERS[typ] = f + return f + + return inner + + +# Type-to-strategy resolver functions + + +@resolves(pydantic.JsonWrapper) +def resolve_json(cls): # type: ignore[no-untyped-def] + try: + inner = st.none() if cls.inner_type is None else st.from_type(cls.inner_type) + except Exception: # pragma: no cover + finite = st.floats(allow_infinity=False, allow_nan=False) + inner = st.recursive( + base=st.one_of(st.none(), st.booleans(), st.integers(), finite, st.text()), + extend=lambda x: st.lists(x) | st.dictionaries(st.text(), x), # type: ignore + ) + inner_type = getattr(cls, 'inner_type', None) + return st.builds( + cls.inner_type.json if lenient_issubclass(inner_type, pydantic.BaseModel) else json.dumps, + inner, + ensure_ascii=st.booleans(), + indent=st.none() | st.integers(0, 16), + sort_keys=st.booleans(), + ) + + +@resolves(pydantic.ConstrainedBytes) +def resolve_conbytes(cls): # type: ignore[no-untyped-def] # pragma: no cover + min_size = cls.min_length or 0 + max_size = cls.max_length + if not cls.strip_whitespace: + return st.binary(min_size=min_size, max_size=max_size) + # Fun with regex to ensure we neither start nor end with whitespace + repeats = '{{{},{}}}'.format( + min_size - 2 if min_size > 2 else 0, + max_size - 2 if (max_size or 0) > 2 else '', + ) + if min_size >= 2: + pattern = rf'\W.{repeats}\W' + elif min_size == 1: + pattern = rf'\W(.{repeats}\W)?' + else: + assert min_size == 0 + pattern = rf'(\W(.{repeats}\W)?)?' + return st.from_regex(pattern.encode(), fullmatch=True) + + +@resolves(pydantic.ConstrainedDecimal) +def resolve_condecimal(cls): # type: ignore[no-untyped-def] + min_value = cls.ge + max_value = cls.le + if cls.gt is not None: + assert min_value is None, 'Set `gt` or `ge`, but not both' + min_value = cls.gt + if cls.lt is not None: + assert max_value is None, 'Set `lt` or `le`, but not both' + max_value = cls.lt + s = st.decimals(min_value, max_value, allow_nan=False, places=cls.decimal_places) + if cls.lt is not None: + s = s.filter(lambda d: d < cls.lt) + if cls.gt is not None: + s = s.filter(lambda d: cls.gt < d) + return s + + +@resolves(pydantic.ConstrainedFloat) +def resolve_confloat(cls): # type: ignore[no-untyped-def] + min_value = cls.ge + max_value = cls.le + exclude_min = False + exclude_max = False + + if cls.gt is not None: + assert min_value is None, 'Set `gt` or `ge`, but not both' + min_value = cls.gt + exclude_min = True + if cls.lt is not None: + assert max_value is None, 'Set `lt` or `le`, but not both' + max_value = cls.lt + exclude_max = True + + if cls.multiple_of is None: + return st.floats(min_value, max_value, exclude_min=exclude_min, exclude_max=exclude_max, allow_nan=False) + + if min_value is not None: + min_value = math.ceil(min_value / cls.multiple_of) + if exclude_min: + min_value = min_value + 1 + if max_value is not None: + assert max_value >= cls.multiple_of, 'Cannot build model with max value smaller than multiple of' + max_value = math.floor(max_value / cls.multiple_of) + if exclude_max: + max_value = max_value - 1 + + return st.integers(min_value, max_value).map(lambda x: x * cls.multiple_of) + + +@resolves(pydantic.ConstrainedInt) +def resolve_conint(cls): # type: ignore[no-untyped-def] + min_value = cls.ge + max_value = cls.le + if cls.gt is not None: + assert min_value is None, 'Set `gt` or `ge`, but not both' + min_value = cls.gt + 1 + if cls.lt is not None: + assert max_value is None, 'Set `lt` or `le`, but not both' + max_value = cls.lt - 1 + + if cls.multiple_of is None or cls.multiple_of == 1: + return st.integers(min_value, max_value) + + # These adjustments and the .map handle integer-valued multiples, while the + # .filter handles trickier cases as for confloat. + if min_value is not None: + min_value = math.ceil(Fraction(min_value) / Fraction(cls.multiple_of)) + if max_value is not None: + max_value = math.floor(Fraction(max_value) / Fraction(cls.multiple_of)) + return st.integers(min_value, max_value).map(lambda x: x * cls.multiple_of) + + +@resolves(pydantic.ConstrainedDate) +def resolve_condate(cls): # type: ignore[no-untyped-def] + if cls.ge is not None: + assert cls.gt is None, 'Set `gt` or `ge`, but not both' + min_value = cls.ge + elif cls.gt is not None: + min_value = cls.gt + datetime.timedelta(days=1) + else: + min_value = datetime.date.min + if cls.le is not None: + assert cls.lt is None, 'Set `lt` or `le`, but not both' + max_value = cls.le + elif cls.lt is not None: + max_value = cls.lt - datetime.timedelta(days=1) + else: + max_value = datetime.date.max + return st.dates(min_value, max_value) + + +@resolves(pydantic.ConstrainedStr) +def resolve_constr(cls): # type: ignore[no-untyped-def] # pragma: no cover + min_size = cls.min_length or 0 + max_size = cls.max_length + + if cls.regex is None and not cls.strip_whitespace: + return st.text(min_size=min_size, max_size=max_size) + + if cls.regex is not None: + strategy = st.from_regex(cls.regex) + if cls.strip_whitespace: + strategy = strategy.filter(lambda s: s == s.strip()) + elif cls.strip_whitespace: + repeats = '{{{},{}}}'.format( + min_size - 2 if min_size > 2 else 0, + max_size - 2 if (max_size or 0) > 2 else '', + ) + if min_size >= 2: + strategy = st.from_regex(rf'\W.{repeats}\W') + elif min_size == 1: + strategy = st.from_regex(rf'\W(.{repeats}\W)?') + else: + assert min_size == 0 + strategy = st.from_regex(rf'(\W(.{repeats}\W)?)?') + + if min_size == 0 and max_size is None: + return strategy + elif max_size is None: + return strategy.filter(lambda s: min_size <= len(s)) + return strategy.filter(lambda s: min_size <= len(s) <= max_size) + + +# Finally, register all previously-defined types, and patch in our new function +for typ in list(pydantic.types._DEFINED_TYPES): + _registered(typ) +pydantic.types._registered = _registered +st.register_type_strategy(pydantic.Json, resolve_json) diff --git a/lib/python3.12/site-packages/pydantic/v1/annotated_types.py b/lib/python3.12/site-packages/pydantic/v1/annotated_types.py new file mode 100644 index 0000000000000000000000000000000000000000..d9eaaafd5b6b2617baf3b3c889659c5bac807ff2 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/annotated_types.py @@ -0,0 +1,72 @@ +import sys +from typing import TYPE_CHECKING, Any, Dict, FrozenSet, NamedTuple, Type + +from pydantic.v1.fields import Required +from pydantic.v1.main import BaseModel, create_model +from pydantic.v1.typing import is_typeddict, is_typeddict_special + +if TYPE_CHECKING: + from typing_extensions import TypedDict + +if sys.version_info < (3, 11): + + def is_legacy_typeddict(typeddict_cls: Type['TypedDict']) -> bool: # type: ignore[valid-type] + return is_typeddict(typeddict_cls) and type(typeddict_cls).__module__ == 'typing' + +else: + + def is_legacy_typeddict(_: Any) -> Any: + return False + + +def create_model_from_typeddict( + # Mypy bug: `Type[TypedDict]` is resolved as `Any` https://github.com/python/mypy/issues/11030 + typeddict_cls: Type['TypedDict'], # type: ignore[valid-type] + **kwargs: Any, +) -> Type['BaseModel']: + """ + Create a `BaseModel` based on the fields of a `TypedDict`. + Since `typing.TypedDict` in Python 3.8 does not store runtime information about optional keys, + we raise an error if this happens (see https://bugs.python.org/issue38834). + """ + field_definitions: Dict[str, Any] + + # Best case scenario: with python 3.9+ or when `TypedDict` is imported from `typing_extensions` + if not hasattr(typeddict_cls, '__required_keys__'): + raise TypeError( + 'You should use `typing_extensions.TypedDict` instead of `typing.TypedDict` with Python < 3.9.2. ' + 'Without it, there is no way to differentiate required and optional fields when subclassed.' + ) + + if is_legacy_typeddict(typeddict_cls) and any( + is_typeddict_special(t) for t in typeddict_cls.__annotations__.values() + ): + raise TypeError( + 'You should use `typing_extensions.TypedDict` instead of `typing.TypedDict` with Python < 3.11. ' + 'Without it, there is no way to reflect Required/NotRequired keys.' + ) + + required_keys: FrozenSet[str] = typeddict_cls.__required_keys__ # type: ignore[attr-defined] + field_definitions = { + field_name: (field_type, Required if field_name in required_keys else None) + for field_name, field_type in typeddict_cls.__annotations__.items() + } + + return create_model(typeddict_cls.__name__, **kwargs, **field_definitions) + + +def create_model_from_namedtuple(namedtuple_cls: Type['NamedTuple'], **kwargs: Any) -> Type['BaseModel']: + """ + Create a `BaseModel` based on the fields of a named tuple. + A named tuple can be created with `typing.NamedTuple` and declared annotations + but also with `collections.namedtuple`, in this case we consider all fields + to have type `Any`. + """ + # With python 3.10+, `__annotations__` always exists but can be empty hence the `getattr... or...` logic + namedtuple_annotations: Dict[str, Type[Any]] = getattr(namedtuple_cls, '__annotations__', None) or { + k: Any for k in namedtuple_cls._fields + } + field_definitions: Dict[str, Any] = { + field_name: (field_type, Required) for field_name, field_type in namedtuple_annotations.items() + } + return create_model(namedtuple_cls.__name__, **kwargs, **field_definitions) diff --git a/lib/python3.12/site-packages/pydantic/v1/class_validators.py b/lib/python3.12/site-packages/pydantic/v1/class_validators.py new file mode 100644 index 0000000000000000000000000000000000000000..2f68fc8600038d8de10017b0d02a3fde77a06ba6 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/class_validators.py @@ -0,0 +1,361 @@ +import warnings +from collections import ChainMap +from functools import partial, partialmethod, wraps +from itertools import chain +from types import FunctionType +from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Type, Union, overload + +from pydantic.v1.errors import ConfigError +from pydantic.v1.typing import AnyCallable +from pydantic.v1.utils import ROOT_KEY, in_ipython + +if TYPE_CHECKING: + from pydantic.v1.typing import AnyClassMethod + + +class Validator: + __slots__ = 'func', 'pre', 'each_item', 'always', 'check_fields', 'skip_on_failure' + + def __init__( + self, + func: AnyCallable, + pre: bool = False, + each_item: bool = False, + always: bool = False, + check_fields: bool = False, + skip_on_failure: bool = False, + ): + self.func = func + self.pre = pre + self.each_item = each_item + self.always = always + self.check_fields = check_fields + self.skip_on_failure = skip_on_failure + + +if TYPE_CHECKING: + from inspect import Signature + + from pydantic.v1.config import BaseConfig + from pydantic.v1.fields import ModelField + from pydantic.v1.types import ModelOrDc + + ValidatorCallable = Callable[[Optional[ModelOrDc], Any, Dict[str, Any], ModelField, Type[BaseConfig]], Any] + ValidatorsList = List[ValidatorCallable] + ValidatorListDict = Dict[str, List[Validator]] + +_FUNCS: Set[str] = set() +VALIDATOR_CONFIG_KEY = '__validator_config__' +ROOT_VALIDATOR_CONFIG_KEY = '__root_validator_config__' + + +def validator( + *fields: str, + pre: bool = False, + each_item: bool = False, + always: bool = False, + check_fields: bool = True, + whole: Optional[bool] = None, + allow_reuse: bool = False, +) -> Callable[[AnyCallable], 'AnyClassMethod']: + """ + Decorate methods on the class indicating that they should be used to validate fields + :param fields: which field(s) the method should be called on + :param pre: whether or not this validator should be called before the standard validators (else after) + :param each_item: for complex objects (sets, lists etc.) whether to validate individual elements rather than the + whole object + :param always: whether this method and other validators should be called even if the value is missing + :param check_fields: whether to check that the fields actually exist on the model + :param allow_reuse: whether to track and raise an error if another validator refers to the decorated function + """ + if not fields: + raise ConfigError('validator with no fields specified') + elif isinstance(fields[0], FunctionType): + raise ConfigError( + "validators should be used with fields and keyword arguments, not bare. " # noqa: Q000 + "E.g. usage should be `@validator('', ...)`" + ) + elif not all(isinstance(field, str) for field in fields): + raise ConfigError( + "validator fields should be passed as separate string args. " # noqa: Q000 + "E.g. usage should be `@validator('', '', ...)`" + ) + + if whole is not None: + warnings.warn( + 'The "whole" keyword argument is deprecated, use "each_item" (inverse meaning, default False) instead', + DeprecationWarning, + ) + assert each_item is False, '"each_item" and "whole" conflict, remove "whole"' + each_item = not whole + + def dec(f: AnyCallable) -> 'AnyClassMethod': + f_cls = _prepare_validator(f, allow_reuse) + setattr( + f_cls, + VALIDATOR_CONFIG_KEY, + ( + fields, + Validator(func=f_cls.__func__, pre=pre, each_item=each_item, always=always, check_fields=check_fields), + ), + ) + return f_cls + + return dec + + +@overload +def root_validator(_func: AnyCallable) -> 'AnyClassMethod': + ... + + +@overload +def root_validator( + *, pre: bool = False, allow_reuse: bool = False, skip_on_failure: bool = False +) -> Callable[[AnyCallable], 'AnyClassMethod']: + ... + + +def root_validator( + _func: Optional[AnyCallable] = None, *, pre: bool = False, allow_reuse: bool = False, skip_on_failure: bool = False +) -> Union['AnyClassMethod', Callable[[AnyCallable], 'AnyClassMethod']]: + """ + Decorate methods on a model indicating that they should be used to validate (and perhaps modify) data either + before or after standard model parsing/validation is performed. + """ + if _func: + f_cls = _prepare_validator(_func, allow_reuse) + setattr( + f_cls, ROOT_VALIDATOR_CONFIG_KEY, Validator(func=f_cls.__func__, pre=pre, skip_on_failure=skip_on_failure) + ) + return f_cls + + def dec(f: AnyCallable) -> 'AnyClassMethod': + f_cls = _prepare_validator(f, allow_reuse) + setattr( + f_cls, ROOT_VALIDATOR_CONFIG_KEY, Validator(func=f_cls.__func__, pre=pre, skip_on_failure=skip_on_failure) + ) + return f_cls + + return dec + + +def _prepare_validator(function: AnyCallable, allow_reuse: bool) -> 'AnyClassMethod': + """ + Avoid validators with duplicated names since without this, validators can be overwritten silently + which generally isn't the intended behaviour, don't run in ipython (see #312) or if allow_reuse is False. + """ + f_cls = function if isinstance(function, classmethod) else classmethod(function) + if not in_ipython() and not allow_reuse: + ref = ( + getattr(f_cls.__func__, '__module__', '') + + '.' + + getattr(f_cls.__func__, '__qualname__', f'') + ) + if ref in _FUNCS: + raise ConfigError(f'duplicate validator function "{ref}"; if this is intended, set `allow_reuse=True`') + _FUNCS.add(ref) + return f_cls + + +class ValidatorGroup: + def __init__(self, validators: 'ValidatorListDict') -> None: + self.validators = validators + self.used_validators = {'*'} + + def get_validators(self, name: str) -> Optional[Dict[str, Validator]]: + self.used_validators.add(name) + validators = self.validators.get(name, []) + if name != ROOT_KEY: + validators += self.validators.get('*', []) + if validators: + return {getattr(v.func, '__name__', f''): v for v in validators} + else: + return None + + def check_for_unused(self) -> None: + unused_validators = set( + chain.from_iterable( + ( + getattr(v.func, '__name__', f'') + for v in self.validators[f] + if v.check_fields + ) + for f in (self.validators.keys() - self.used_validators) + ) + ) + if unused_validators: + fn = ', '.join(unused_validators) + raise ConfigError( + f"Validators defined with incorrect fields: {fn} " # noqa: Q000 + f"(use check_fields=False if you're inheriting from the model and intended this)" + ) + + +def extract_validators(namespace: Dict[str, Any]) -> Dict[str, List[Validator]]: + validators: Dict[str, List[Validator]] = {} + for var_name, value in namespace.items(): + validator_config = getattr(value, VALIDATOR_CONFIG_KEY, None) + if validator_config: + fields, v = validator_config + for field in fields: + if field in validators: + validators[field].append(v) + else: + validators[field] = [v] + return validators + + +def extract_root_validators(namespace: Dict[str, Any]) -> Tuple[List[AnyCallable], List[Tuple[bool, AnyCallable]]]: + from inspect import signature + + pre_validators: List[AnyCallable] = [] + post_validators: List[Tuple[bool, AnyCallable]] = [] + for name, value in namespace.items(): + validator_config: Optional[Validator] = getattr(value, ROOT_VALIDATOR_CONFIG_KEY, None) + if validator_config: + sig = signature(validator_config.func) + args = list(sig.parameters.keys()) + if args[0] == 'self': + raise ConfigError( + f'Invalid signature for root validator {name}: {sig}, "self" not permitted as first argument, ' + f'should be: (cls, values).' + ) + if len(args) != 2: + raise ConfigError(f'Invalid signature for root validator {name}: {sig}, should be: (cls, values).') + # check function signature + if validator_config.pre: + pre_validators.append(validator_config.func) + else: + post_validators.append((validator_config.skip_on_failure, validator_config.func)) + return pre_validators, post_validators + + +def inherit_validators(base_validators: 'ValidatorListDict', validators: 'ValidatorListDict') -> 'ValidatorListDict': + for field, field_validators in base_validators.items(): + if field not in validators: + validators[field] = [] + validators[field] += field_validators + return validators + + +def make_generic_validator(validator: AnyCallable) -> 'ValidatorCallable': + """ + Make a generic function which calls a validator with the right arguments. + + Unfortunately other approaches (eg. return a partial of a function that builds the arguments) is slow, + hence this laborious way of doing things. + + It's done like this so validators don't all need **kwargs in their signature, eg. any combination of + the arguments "values", "fields" and/or "config" are permitted. + """ + from inspect import signature + + if not isinstance(validator, (partial, partialmethod)): + # This should be the default case, so overhead is reduced + sig = signature(validator) + args = list(sig.parameters.keys()) + else: + # Fix the generated argument lists of partial methods + sig = signature(validator.func) + args = [ + k + for k in signature(validator.func).parameters.keys() + if k not in validator.args | validator.keywords.keys() + ] + + first_arg = args.pop(0) + if first_arg == 'self': + raise ConfigError( + f'Invalid signature for validator {validator}: {sig}, "self" not permitted as first argument, ' + f'should be: (cls, value, values, config, field), "values", "config" and "field" are all optional.' + ) + elif first_arg == 'cls': + # assume the second argument is value + return wraps(validator)(_generic_validator_cls(validator, sig, set(args[1:]))) + else: + # assume the first argument was value which has already been removed + return wraps(validator)(_generic_validator_basic(validator, sig, set(args))) + + +def prep_validators(v_funcs: Iterable[AnyCallable]) -> 'ValidatorsList': + return [make_generic_validator(f) for f in v_funcs if f] + + +all_kwargs = {'values', 'field', 'config'} + + +def _generic_validator_cls(validator: AnyCallable, sig: 'Signature', args: Set[str]) -> 'ValidatorCallable': + # assume the first argument is value + has_kwargs = False + if 'kwargs' in args: + has_kwargs = True + args -= {'kwargs'} + + if not args.issubset(all_kwargs): + raise ConfigError( + f'Invalid signature for validator {validator}: {sig}, should be: ' + f'(cls, value, values, config, field), "values", "config" and "field" are all optional.' + ) + + if has_kwargs: + return lambda cls, v, values, field, config: validator(cls, v, values=values, field=field, config=config) + elif args == set(): + return lambda cls, v, values, field, config: validator(cls, v) + elif args == {'values'}: + return lambda cls, v, values, field, config: validator(cls, v, values=values) + elif args == {'field'}: + return lambda cls, v, values, field, config: validator(cls, v, field=field) + elif args == {'config'}: + return lambda cls, v, values, field, config: validator(cls, v, config=config) + elif args == {'values', 'field'}: + return lambda cls, v, values, field, config: validator(cls, v, values=values, field=field) + elif args == {'values', 'config'}: + return lambda cls, v, values, field, config: validator(cls, v, values=values, config=config) + elif args == {'field', 'config'}: + return lambda cls, v, values, field, config: validator(cls, v, field=field, config=config) + else: + # args == {'values', 'field', 'config'} + return lambda cls, v, values, field, config: validator(cls, v, values=values, field=field, config=config) + + +def _generic_validator_basic(validator: AnyCallable, sig: 'Signature', args: Set[str]) -> 'ValidatorCallable': + has_kwargs = False + if 'kwargs' in args: + has_kwargs = True + args -= {'kwargs'} + + if not args.issubset(all_kwargs): + raise ConfigError( + f'Invalid signature for validator {validator}: {sig}, should be: ' + f'(value, values, config, field), "values", "config" and "field" are all optional.' + ) + + if has_kwargs: + return lambda cls, v, values, field, config: validator(v, values=values, field=field, config=config) + elif args == set(): + return lambda cls, v, values, field, config: validator(v) + elif args == {'values'}: + return lambda cls, v, values, field, config: validator(v, values=values) + elif args == {'field'}: + return lambda cls, v, values, field, config: validator(v, field=field) + elif args == {'config'}: + return lambda cls, v, values, field, config: validator(v, config=config) + elif args == {'values', 'field'}: + return lambda cls, v, values, field, config: validator(v, values=values, field=field) + elif args == {'values', 'config'}: + return lambda cls, v, values, field, config: validator(v, values=values, config=config) + elif args == {'field', 'config'}: + return lambda cls, v, values, field, config: validator(v, field=field, config=config) + else: + # args == {'values', 'field', 'config'} + return lambda cls, v, values, field, config: validator(v, values=values, field=field, config=config) + + +def gather_all_validators(type_: 'ModelOrDc') -> Dict[str, 'AnyClassMethod']: + all_attributes = ChainMap(*[cls.__dict__ for cls in type_.__mro__]) # type: ignore[arg-type,var-annotated] + return { + k: v + for k, v in all_attributes.items() + if hasattr(v, VALIDATOR_CONFIG_KEY) or hasattr(v, ROOT_VALIDATOR_CONFIG_KEY) + } diff --git a/lib/python3.12/site-packages/pydantic/v1/color.py b/lib/python3.12/site-packages/pydantic/v1/color.py new file mode 100644 index 0000000000000000000000000000000000000000..b0bbf78f4367c3e13a5aac0df09b5824d8f7e6ea --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/color.py @@ -0,0 +1,494 @@ +""" +Color definitions are used as per CSS3 specification: +http://www.w3.org/TR/css3-color/#svg-color + +A few colors have multiple names referring to the sames colors, eg. `grey` and `gray` or `aqua` and `cyan`. + +In these cases the LAST color when sorted alphabetically takes preferences, +eg. Color((0, 255, 255)).as_named() == 'cyan' because "cyan" comes after "aqua". +""" +import math +import re +from colorsys import hls_to_rgb, rgb_to_hls +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union, cast + +from pydantic.v1.errors import ColorError +from pydantic.v1.utils import Representation, almost_equal_floats + +if TYPE_CHECKING: + from pydantic.v1.typing import CallableGenerator, ReprArgs + +ColorTuple = Union[Tuple[int, int, int], Tuple[int, int, int, float]] +ColorType = Union[ColorTuple, str] +HslColorTuple = Union[Tuple[float, float, float], Tuple[float, float, float, float]] + + +class RGBA: + """ + Internal use only as a representation of a color. + """ + + __slots__ = 'r', 'g', 'b', 'alpha', '_tuple' + + def __init__(self, r: float, g: float, b: float, alpha: Optional[float]): + self.r = r + self.g = g + self.b = b + self.alpha = alpha + + self._tuple: Tuple[float, float, float, Optional[float]] = (r, g, b, alpha) + + def __getitem__(self, item: Any) -> Any: + return self._tuple[item] + + +# these are not compiled here to avoid import slowdown, they'll be compiled the first time they're used, then cached +r_hex_short = r'\s*(?:#|0x)?([0-9a-f])([0-9a-f])([0-9a-f])([0-9a-f])?\s*' +r_hex_long = r'\s*(?:#|0x)?([0-9a-f]{2})([0-9a-f]{2})([0-9a-f]{2})([0-9a-f]{2})?\s*' +_r_255 = r'(\d{1,3}(?:\.\d+)?)' +_r_comma = r'\s*,\s*' +r_rgb = fr'\s*rgb\(\s*{_r_255}{_r_comma}{_r_255}{_r_comma}{_r_255}\)\s*' +_r_alpha = r'(\d(?:\.\d+)?|\.\d+|\d{1,2}%)' +r_rgba = fr'\s*rgba\(\s*{_r_255}{_r_comma}{_r_255}{_r_comma}{_r_255}{_r_comma}{_r_alpha}\s*\)\s*' +_r_h = r'(-?\d+(?:\.\d+)?|-?\.\d+)(deg|rad|turn)?' +_r_sl = r'(\d{1,3}(?:\.\d+)?)%' +r_hsl = fr'\s*hsl\(\s*{_r_h}{_r_comma}{_r_sl}{_r_comma}{_r_sl}\s*\)\s*' +r_hsla = fr'\s*hsl\(\s*{_r_h}{_r_comma}{_r_sl}{_r_comma}{_r_sl}{_r_comma}{_r_alpha}\s*\)\s*' + +# colors where the two hex characters are the same, if all colors match this the short version of hex colors can be used +repeat_colors = {int(c * 2, 16) for c in '0123456789abcdef'} +rads = 2 * math.pi + + +class Color(Representation): + __slots__ = '_original', '_rgba' + + def __init__(self, value: ColorType) -> None: + self._rgba: RGBA + self._original: ColorType + if isinstance(value, (tuple, list)): + self._rgba = parse_tuple(value) + elif isinstance(value, str): + self._rgba = parse_str(value) + elif isinstance(value, Color): + self._rgba = value._rgba + value = value._original + else: + raise ColorError(reason='value must be a tuple, list or string') + + # if we've got here value must be a valid color + self._original = value + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='string', format='color') + + def original(self) -> ColorType: + """ + Original value passed to Color + """ + return self._original + + def as_named(self, *, fallback: bool = False) -> str: + if self._rgba.alpha is None: + rgb = cast(Tuple[int, int, int], self.as_rgb_tuple()) + try: + return COLORS_BY_VALUE[rgb] + except KeyError as e: + if fallback: + return self.as_hex() + else: + raise ValueError('no named color found, use fallback=True, as_hex() or as_rgb()') from e + else: + return self.as_hex() + + def as_hex(self) -> str: + """ + Hex string representing the color can be 3, 4, 6 or 8 characters depending on whether the string + a "short" representation of the color is possible and whether there's an alpha channel. + """ + values = [float_to_255(c) for c in self._rgba[:3]] + if self._rgba.alpha is not None: + values.append(float_to_255(self._rgba.alpha)) + + as_hex = ''.join(f'{v:02x}' for v in values) + if all(c in repeat_colors for c in values): + as_hex = ''.join(as_hex[c] for c in range(0, len(as_hex), 2)) + return '#' + as_hex + + def as_rgb(self) -> str: + """ + Color as an rgb(, , ) or rgba(, , , ) string. + """ + if self._rgba.alpha is None: + return f'rgb({float_to_255(self._rgba.r)}, {float_to_255(self._rgba.g)}, {float_to_255(self._rgba.b)})' + else: + return ( + f'rgba({float_to_255(self._rgba.r)}, {float_to_255(self._rgba.g)}, {float_to_255(self._rgba.b)}, ' + f'{round(self._alpha_float(), 2)})' + ) + + def as_rgb_tuple(self, *, alpha: Optional[bool] = None) -> ColorTuple: + """ + Color as an RGB or RGBA tuple; red, green and blue are in the range 0 to 255, alpha if included is + in the range 0 to 1. + + :param alpha: whether to include the alpha channel, options are + None - (default) include alpha only if it's set (e.g. not None) + True - always include alpha, + False - always omit alpha, + """ + r, g, b = (float_to_255(c) for c in self._rgba[:3]) + if alpha is None: + if self._rgba.alpha is None: + return r, g, b + else: + return r, g, b, self._alpha_float() + elif alpha: + return r, g, b, self._alpha_float() + else: + # alpha is False + return r, g, b + + def as_hsl(self) -> str: + """ + Color as an hsl(, , ) or hsl(, , , ) string. + """ + if self._rgba.alpha is None: + h, s, li = self.as_hsl_tuple(alpha=False) # type: ignore + return f'hsl({h * 360:0.0f}, {s:0.0%}, {li:0.0%})' + else: + h, s, li, a = self.as_hsl_tuple(alpha=True) # type: ignore + return f'hsl({h * 360:0.0f}, {s:0.0%}, {li:0.0%}, {round(a, 2)})' + + def as_hsl_tuple(self, *, alpha: Optional[bool] = None) -> HslColorTuple: + """ + Color as an HSL or HSLA tuple, e.g. hue, saturation, lightness and optionally alpha; all elements are in + the range 0 to 1. + + NOTE: this is HSL as used in HTML and most other places, not HLS as used in python's colorsys. + + :param alpha: whether to include the alpha channel, options are + None - (default) include alpha only if it's set (e.g. not None) + True - always include alpha, + False - always omit alpha, + """ + h, l, s = rgb_to_hls(self._rgba.r, self._rgba.g, self._rgba.b) + if alpha is None: + if self._rgba.alpha is None: + return h, s, l + else: + return h, s, l, self._alpha_float() + if alpha: + return h, s, l, self._alpha_float() + else: + # alpha is False + return h, s, l + + def _alpha_float(self) -> float: + return 1 if self._rgba.alpha is None else self._rgba.alpha + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls + + def __str__(self) -> str: + return self.as_named(fallback=True) + + def __repr_args__(self) -> 'ReprArgs': + return [(None, self.as_named(fallback=True))] + [('rgb', self.as_rgb_tuple())] # type: ignore + + def __eq__(self, other: Any) -> bool: + return isinstance(other, Color) and self.as_rgb_tuple() == other.as_rgb_tuple() + + def __hash__(self) -> int: + return hash(self.as_rgb_tuple()) + + +def parse_tuple(value: Tuple[Any, ...]) -> RGBA: + """ + Parse a tuple or list as a color. + """ + if len(value) == 3: + r, g, b = (parse_color_value(v) for v in value) + return RGBA(r, g, b, None) + elif len(value) == 4: + r, g, b = (parse_color_value(v) for v in value[:3]) + return RGBA(r, g, b, parse_float_alpha(value[3])) + else: + raise ColorError(reason='tuples must have length 3 or 4') + + +def parse_str(value: str) -> RGBA: + """ + Parse a string to an RGBA tuple, trying the following formats (in this order): + * named color, see COLORS_BY_NAME below + * hex short eg. `fff` (prefix can be `#`, `0x` or nothing) + * hex long eg. `ffffff` (prefix can be `#`, `0x` or nothing) + * `rgb(, , ) ` + * `rgba(, , , )` + """ + value_lower = value.lower() + try: + r, g, b = COLORS_BY_NAME[value_lower] + except KeyError: + pass + else: + return ints_to_rgba(r, g, b, None) + + m = re.fullmatch(r_hex_short, value_lower) + if m: + *rgb, a = m.groups() + r, g, b = (int(v * 2, 16) for v in rgb) + if a: + alpha: Optional[float] = int(a * 2, 16) / 255 + else: + alpha = None + return ints_to_rgba(r, g, b, alpha) + + m = re.fullmatch(r_hex_long, value_lower) + if m: + *rgb, a = m.groups() + r, g, b = (int(v, 16) for v in rgb) + if a: + alpha = int(a, 16) / 255 + else: + alpha = None + return ints_to_rgba(r, g, b, alpha) + + m = re.fullmatch(r_rgb, value_lower) + if m: + return ints_to_rgba(*m.groups(), None) # type: ignore + + m = re.fullmatch(r_rgba, value_lower) + if m: + return ints_to_rgba(*m.groups()) # type: ignore + + m = re.fullmatch(r_hsl, value_lower) + if m: + h, h_units, s, l_ = m.groups() + return parse_hsl(h, h_units, s, l_) + + m = re.fullmatch(r_hsla, value_lower) + if m: + h, h_units, s, l_, a = m.groups() + return parse_hsl(h, h_units, s, l_, parse_float_alpha(a)) + + raise ColorError(reason='string not recognised as a valid color') + + +def ints_to_rgba(r: Union[int, str], g: Union[int, str], b: Union[int, str], alpha: Optional[float]) -> RGBA: + return RGBA(parse_color_value(r), parse_color_value(g), parse_color_value(b), parse_float_alpha(alpha)) + + +def parse_color_value(value: Union[int, str], max_val: int = 255) -> float: + """ + Parse a value checking it's a valid int in the range 0 to max_val and divide by max_val to give a number + in the range 0 to 1 + """ + try: + color = float(value) + except ValueError: + raise ColorError(reason='color values must be a valid number') + if 0 <= color <= max_val: + return color / max_val + else: + raise ColorError(reason=f'color values must be in the range 0 to {max_val}') + + +def parse_float_alpha(value: Union[None, str, float, int]) -> Optional[float]: + """ + Parse a value checking it's a valid float in the range 0 to 1 + """ + if value is None: + return None + try: + if isinstance(value, str) and value.endswith('%'): + alpha = float(value[:-1]) / 100 + else: + alpha = float(value) + except ValueError: + raise ColorError(reason='alpha values must be a valid float') + + if almost_equal_floats(alpha, 1): + return None + elif 0 <= alpha <= 1: + return alpha + else: + raise ColorError(reason='alpha values must be in the range 0 to 1') + + +def parse_hsl(h: str, h_units: str, sat: str, light: str, alpha: Optional[float] = None) -> RGBA: + """ + Parse raw hue, saturation, lightness and alpha values and convert to RGBA. + """ + s_value, l_value = parse_color_value(sat, 100), parse_color_value(light, 100) + + h_value = float(h) + if h_units in {None, 'deg'}: + h_value = h_value % 360 / 360 + elif h_units == 'rad': + h_value = h_value % rads / rads + else: + # turns + h_value = h_value % 1 + + r, g, b = hls_to_rgb(h_value, l_value, s_value) + return RGBA(r, g, b, alpha) + + +def float_to_255(c: float) -> int: + return int(round(c * 255)) + + +COLORS_BY_NAME = { + 'aliceblue': (240, 248, 255), + 'antiquewhite': (250, 235, 215), + 'aqua': (0, 255, 255), + 'aquamarine': (127, 255, 212), + 'azure': (240, 255, 255), + 'beige': (245, 245, 220), + 'bisque': (255, 228, 196), + 'black': (0, 0, 0), + 'blanchedalmond': (255, 235, 205), + 'blue': (0, 0, 255), + 'blueviolet': (138, 43, 226), + 'brown': (165, 42, 42), + 'burlywood': (222, 184, 135), + 'cadetblue': (95, 158, 160), + 'chartreuse': (127, 255, 0), + 'chocolate': (210, 105, 30), + 'coral': (255, 127, 80), + 'cornflowerblue': (100, 149, 237), + 'cornsilk': (255, 248, 220), + 'crimson': (220, 20, 60), + 'cyan': (0, 255, 255), + 'darkblue': (0, 0, 139), + 'darkcyan': (0, 139, 139), + 'darkgoldenrod': (184, 134, 11), + 'darkgray': (169, 169, 169), + 'darkgreen': (0, 100, 0), + 'darkgrey': (169, 169, 169), + 'darkkhaki': (189, 183, 107), + 'darkmagenta': (139, 0, 139), + 'darkolivegreen': (85, 107, 47), + 'darkorange': (255, 140, 0), + 'darkorchid': (153, 50, 204), + 'darkred': (139, 0, 0), + 'darksalmon': (233, 150, 122), + 'darkseagreen': (143, 188, 143), + 'darkslateblue': (72, 61, 139), + 'darkslategray': (47, 79, 79), + 'darkslategrey': (47, 79, 79), + 'darkturquoise': (0, 206, 209), + 'darkviolet': (148, 0, 211), + 'deeppink': (255, 20, 147), + 'deepskyblue': (0, 191, 255), + 'dimgray': (105, 105, 105), + 'dimgrey': (105, 105, 105), + 'dodgerblue': (30, 144, 255), + 'firebrick': (178, 34, 34), + 'floralwhite': (255, 250, 240), + 'forestgreen': (34, 139, 34), + 'fuchsia': (255, 0, 255), + 'gainsboro': (220, 220, 220), + 'ghostwhite': (248, 248, 255), + 'gold': (255, 215, 0), + 'goldenrod': (218, 165, 32), + 'gray': (128, 128, 128), + 'green': (0, 128, 0), + 'greenyellow': (173, 255, 47), + 'grey': (128, 128, 128), + 'honeydew': (240, 255, 240), + 'hotpink': (255, 105, 180), + 'indianred': (205, 92, 92), + 'indigo': (75, 0, 130), + 'ivory': (255, 255, 240), + 'khaki': (240, 230, 140), + 'lavender': (230, 230, 250), + 'lavenderblush': (255, 240, 245), + 'lawngreen': (124, 252, 0), + 'lemonchiffon': (255, 250, 205), + 'lightblue': (173, 216, 230), + 'lightcoral': (240, 128, 128), + 'lightcyan': (224, 255, 255), + 'lightgoldenrodyellow': (250, 250, 210), + 'lightgray': (211, 211, 211), + 'lightgreen': (144, 238, 144), + 'lightgrey': (211, 211, 211), + 'lightpink': (255, 182, 193), + 'lightsalmon': (255, 160, 122), + 'lightseagreen': (32, 178, 170), + 'lightskyblue': (135, 206, 250), + 'lightslategray': (119, 136, 153), + 'lightslategrey': (119, 136, 153), + 'lightsteelblue': (176, 196, 222), + 'lightyellow': (255, 255, 224), + 'lime': (0, 255, 0), + 'limegreen': (50, 205, 50), + 'linen': (250, 240, 230), + 'magenta': (255, 0, 255), + 'maroon': (128, 0, 0), + 'mediumaquamarine': (102, 205, 170), + 'mediumblue': (0, 0, 205), + 'mediumorchid': (186, 85, 211), + 'mediumpurple': (147, 112, 219), + 'mediumseagreen': (60, 179, 113), + 'mediumslateblue': (123, 104, 238), + 'mediumspringgreen': (0, 250, 154), + 'mediumturquoise': (72, 209, 204), + 'mediumvioletred': (199, 21, 133), + 'midnightblue': (25, 25, 112), + 'mintcream': (245, 255, 250), + 'mistyrose': (255, 228, 225), + 'moccasin': (255, 228, 181), + 'navajowhite': (255, 222, 173), + 'navy': (0, 0, 128), + 'oldlace': (253, 245, 230), + 'olive': (128, 128, 0), + 'olivedrab': (107, 142, 35), + 'orange': (255, 165, 0), + 'orangered': (255, 69, 0), + 'orchid': (218, 112, 214), + 'palegoldenrod': (238, 232, 170), + 'palegreen': (152, 251, 152), + 'paleturquoise': (175, 238, 238), + 'palevioletred': (219, 112, 147), + 'papayawhip': (255, 239, 213), + 'peachpuff': (255, 218, 185), + 'peru': (205, 133, 63), + 'pink': (255, 192, 203), + 'plum': (221, 160, 221), + 'powderblue': (176, 224, 230), + 'purple': (128, 0, 128), + 'red': (255, 0, 0), + 'rosybrown': (188, 143, 143), + 'royalblue': (65, 105, 225), + 'saddlebrown': (139, 69, 19), + 'salmon': (250, 128, 114), + 'sandybrown': (244, 164, 96), + 'seagreen': (46, 139, 87), + 'seashell': (255, 245, 238), + 'sienna': (160, 82, 45), + 'silver': (192, 192, 192), + 'skyblue': (135, 206, 235), + 'slateblue': (106, 90, 205), + 'slategray': (112, 128, 144), + 'slategrey': (112, 128, 144), + 'snow': (255, 250, 250), + 'springgreen': (0, 255, 127), + 'steelblue': (70, 130, 180), + 'tan': (210, 180, 140), + 'teal': (0, 128, 128), + 'thistle': (216, 191, 216), + 'tomato': (255, 99, 71), + 'turquoise': (64, 224, 208), + 'violet': (238, 130, 238), + 'wheat': (245, 222, 179), + 'white': (255, 255, 255), + 'whitesmoke': (245, 245, 245), + 'yellow': (255, 255, 0), + 'yellowgreen': (154, 205, 50), +} + +COLORS_BY_VALUE = {v: k for k, v in COLORS_BY_NAME.items()} diff --git a/lib/python3.12/site-packages/pydantic/v1/config.py b/lib/python3.12/site-packages/pydantic/v1/config.py new file mode 100644 index 0000000000000000000000000000000000000000..18f7c999bee0fab95293b2434047fd20532a6446 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/config.py @@ -0,0 +1,191 @@ +import json +from enum import Enum +from typing import TYPE_CHECKING, Any, Callable, Dict, ForwardRef, Optional, Tuple, Type, Union + +from typing_extensions import Literal, Protocol + +from pydantic.v1.typing import AnyArgTCallable, AnyCallable +from pydantic.v1.utils import GetterDict +from pydantic.v1.version import compiled + +if TYPE_CHECKING: + from typing import overload + + from pydantic.v1.fields import ModelField + from pydantic.v1.main import BaseModel + + ConfigType = Type['BaseConfig'] + + class SchemaExtraCallable(Protocol): + @overload + def __call__(self, schema: Dict[str, Any]) -> None: + pass + + @overload + def __call__(self, schema: Dict[str, Any], model_class: Type[BaseModel]) -> None: + pass + +else: + SchemaExtraCallable = Callable[..., None] + +__all__ = 'BaseConfig', 'ConfigDict', 'get_config', 'Extra', 'inherit_config', 'prepare_config' + + +class Extra(str, Enum): + allow = 'allow' + ignore = 'ignore' + forbid = 'forbid' + + +# https://github.com/cython/cython/issues/4003 +# Fixed in Cython 3 and Pydantic v1 won't support Cython 3. +# Pydantic v2 doesn't depend on Cython at all. +if not compiled: + from typing_extensions import TypedDict + + class ConfigDict(TypedDict, total=False): + title: Optional[str] + anystr_lower: bool + anystr_strip_whitespace: bool + min_anystr_length: int + max_anystr_length: Optional[int] + validate_all: bool + extra: Extra + allow_mutation: bool + frozen: bool + allow_population_by_field_name: bool + use_enum_values: bool + fields: Dict[str, Union[str, Dict[str, str]]] + validate_assignment: bool + error_msg_templates: Dict[str, str] + arbitrary_types_allowed: bool + orm_mode: bool + getter_dict: Type[GetterDict] + alias_generator: Optional[Callable[[str], str]] + keep_untouched: Tuple[type, ...] + schema_extra: Union[Dict[str, object], 'SchemaExtraCallable'] + json_loads: Callable[[str], object] + json_dumps: AnyArgTCallable[str] + json_encoders: Dict[Type[object], AnyCallable] + underscore_attrs_are_private: bool + allow_inf_nan: bool + copy_on_model_validation: Literal['none', 'deep', 'shallow'] + # whether dataclass `__post_init__` should be run after validation + post_init_call: Literal['before_validation', 'after_validation'] + +else: + ConfigDict = dict # type: ignore + + +class BaseConfig: + title: Optional[str] = None + anystr_lower: bool = False + anystr_upper: bool = False + anystr_strip_whitespace: bool = False + min_anystr_length: int = 0 + max_anystr_length: Optional[int] = None + validate_all: bool = False + extra: Extra = Extra.ignore + allow_mutation: bool = True + frozen: bool = False + allow_population_by_field_name: bool = False + use_enum_values: bool = False + fields: Dict[str, Union[str, Dict[str, str]]] = {} + validate_assignment: bool = False + error_msg_templates: Dict[str, str] = {} + arbitrary_types_allowed: bool = False + orm_mode: bool = False + getter_dict: Type[GetterDict] = GetterDict + alias_generator: Optional[Callable[[str], str]] = None + keep_untouched: Tuple[type, ...] = () + schema_extra: Union[Dict[str, Any], 'SchemaExtraCallable'] = {} + json_loads: Callable[[str], Any] = json.loads + json_dumps: Callable[..., str] = json.dumps + json_encoders: Dict[Union[Type[Any], str, ForwardRef], AnyCallable] = {} + underscore_attrs_are_private: bool = False + allow_inf_nan: bool = True + + # whether inherited models as fields should be reconstructed as base model, + # and whether such a copy should be shallow or deep + copy_on_model_validation: Literal['none', 'deep', 'shallow'] = 'shallow' + + # whether `Union` should check all allowed types before even trying to coerce + smart_union: bool = False + # whether dataclass `__post_init__` should be run before or after validation + post_init_call: Literal['before_validation', 'after_validation'] = 'before_validation' + + @classmethod + def get_field_info(cls, name: str) -> Dict[str, Any]: + """ + Get properties of FieldInfo from the `fields` property of the config class. + """ + + fields_value = cls.fields.get(name) + + if isinstance(fields_value, str): + field_info: Dict[str, Any] = {'alias': fields_value} + elif isinstance(fields_value, dict): + field_info = fields_value + else: + field_info = {} + + if 'alias' in field_info: + field_info.setdefault('alias_priority', 2) + + if field_info.get('alias_priority', 0) <= 1 and cls.alias_generator: + alias = cls.alias_generator(name) + if not isinstance(alias, str): + raise TypeError(f'Config.alias_generator must return str, not {alias.__class__}') + field_info.update(alias=alias, alias_priority=1) + return field_info + + @classmethod + def prepare_field(cls, field: 'ModelField') -> None: + """ + Optional hook to check or modify fields during model creation. + """ + pass + + +def get_config(config: Union[ConfigDict, Type[object], None]) -> Type[BaseConfig]: + if config is None: + return BaseConfig + + else: + config_dict = ( + config + if isinstance(config, dict) + else {k: getattr(config, k) for k in dir(config) if not k.startswith('__')} + ) + + class Config(BaseConfig): + ... + + for k, v in config_dict.items(): + setattr(Config, k, v) + return Config + + +def inherit_config(self_config: 'ConfigType', parent_config: 'ConfigType', **namespace: Any) -> 'ConfigType': + if not self_config: + base_classes: Tuple['ConfigType', ...] = (parent_config,) + elif self_config == parent_config: + base_classes = (self_config,) + else: + base_classes = self_config, parent_config + + namespace['json_encoders'] = { + **getattr(parent_config, 'json_encoders', {}), + **getattr(self_config, 'json_encoders', {}), + **namespace.get('json_encoders', {}), + } + + return type('Config', base_classes, namespace) + + +def prepare_config(config: Type[BaseConfig], cls_name: str) -> None: + if not isinstance(config.extra, Extra): + try: + config.extra = Extra(config.extra) + except ValueError: + raise ValueError(f'"{cls_name}": {config.extra} is not a valid value for "extra"') diff --git a/lib/python3.12/site-packages/pydantic/v1/dataclasses.py b/lib/python3.12/site-packages/pydantic/v1/dataclasses.py new file mode 100644 index 0000000000000000000000000000000000000000..bd1670291d6622a28ecb5b8e9ad68f4bac9305d7 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/dataclasses.py @@ -0,0 +1,500 @@ +""" +The main purpose is to enhance stdlib dataclasses by adding validation +A pydantic dataclass can be generated from scratch or from a stdlib one. + +Behind the scene, a pydantic dataclass is just like a regular one on which we attach +a `BaseModel` and magic methods to trigger the validation of the data. +`__init__` and `__post_init__` are hence overridden and have extra logic to be +able to validate input data. + +When a pydantic dataclass is generated from scratch, it's just a plain dataclass +with validation triggered at initialization + +The tricky part if for stdlib dataclasses that are converted after into pydantic ones e.g. + +```py +@dataclasses.dataclass +class M: + x: int + +ValidatedM = pydantic.dataclasses.dataclass(M) +``` + +We indeed still want to support equality, hashing, repr, ... as if it was the stdlib one! + +```py +assert isinstance(ValidatedM(x=1), M) +assert ValidatedM(x=1) == M(x=1) +``` + +This means we **don't want to create a new dataclass that inherits from it** +The trick is to create a wrapper around `M` that will act as a proxy to trigger +validation without altering default `M` behaviour. +""" +import copy +import dataclasses +import sys +from contextlib import contextmanager +from functools import wraps + +try: + from functools import cached_property +except ImportError: + # cached_property available only for python3.8+ + pass + +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generator, Optional, Type, TypeVar, Union, overload + +from typing_extensions import dataclass_transform + +from pydantic.v1.class_validators import gather_all_validators +from pydantic.v1.config import BaseConfig, ConfigDict, Extra, get_config +from pydantic.v1.error_wrappers import ValidationError +from pydantic.v1.errors import DataclassTypeError +from pydantic.v1.fields import Field, FieldInfo, Required, Undefined +from pydantic.v1.main import create_model, validate_model +from pydantic.v1.utils import ClassAttribute + +if TYPE_CHECKING: + from pydantic.v1.main import BaseModel + from pydantic.v1.typing import CallableGenerator, NoArgAnyCallable + + DataclassT = TypeVar('DataclassT', bound='Dataclass') + + DataclassClassOrWrapper = Union[Type['Dataclass'], 'DataclassProxy'] + + class Dataclass: + # stdlib attributes + __dataclass_fields__: ClassVar[Dict[str, Any]] + __dataclass_params__: ClassVar[Any] # in reality `dataclasses._DataclassParams` + __post_init__: ClassVar[Callable[..., None]] + + # Added by pydantic + __pydantic_run_validation__: ClassVar[bool] + __post_init_post_parse__: ClassVar[Callable[..., None]] + __pydantic_initialised__: ClassVar[bool] + __pydantic_model__: ClassVar[Type[BaseModel]] + __pydantic_validate_values__: ClassVar[Callable[['Dataclass'], None]] + __pydantic_has_field_info_default__: ClassVar[bool] # whether a `pydantic.Field` is used as default value + + def __init__(self, *args: object, **kwargs: object) -> None: + pass + + @classmethod + def __get_validators__(cls: Type['Dataclass']) -> 'CallableGenerator': + pass + + @classmethod + def __validate__(cls: Type['DataclassT'], v: Any) -> 'DataclassT': + pass + + +__all__ = [ + 'dataclass', + 'set_validation', + 'create_pydantic_model_from_dataclass', + 'is_builtin_dataclass', + 'make_dataclass_validator', +] + +_T = TypeVar('_T') + +if sys.version_info >= (3, 10): + + @dataclass_transform(field_specifiers=(dataclasses.field, Field)) + @overload + def dataclass( + *, + init: bool = True, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool = False, + config: Union[ConfigDict, Type[object], None] = None, + validate_on_init: Optional[bool] = None, + use_proxy: Optional[bool] = None, + kw_only: bool = ..., + ) -> Callable[[Type[_T]], 'DataclassClassOrWrapper']: + ... + + @dataclass_transform(field_specifiers=(dataclasses.field, Field)) + @overload + def dataclass( + _cls: Type[_T], + *, + init: bool = True, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool = False, + config: Union[ConfigDict, Type[object], None] = None, + validate_on_init: Optional[bool] = None, + use_proxy: Optional[bool] = None, + kw_only: bool = ..., + ) -> 'DataclassClassOrWrapper': + ... + +else: + + @dataclass_transform(field_specifiers=(dataclasses.field, Field)) + @overload + def dataclass( + *, + init: bool = True, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool = False, + config: Union[ConfigDict, Type[object], None] = None, + validate_on_init: Optional[bool] = None, + use_proxy: Optional[bool] = None, + ) -> Callable[[Type[_T]], 'DataclassClassOrWrapper']: + ... + + @dataclass_transform(field_specifiers=(dataclasses.field, Field)) + @overload + def dataclass( + _cls: Type[_T], + *, + init: bool = True, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool = False, + config: Union[ConfigDict, Type[object], None] = None, + validate_on_init: Optional[bool] = None, + use_proxy: Optional[bool] = None, + ) -> 'DataclassClassOrWrapper': + ... + + +@dataclass_transform(field_specifiers=(dataclasses.field, Field)) +def dataclass( + _cls: Optional[Type[_T]] = None, + *, + init: bool = True, + repr: bool = True, + eq: bool = True, + order: bool = False, + unsafe_hash: bool = False, + frozen: bool = False, + config: Union[ConfigDict, Type[object], None] = None, + validate_on_init: Optional[bool] = None, + use_proxy: Optional[bool] = None, + kw_only: bool = False, +) -> Union[Callable[[Type[_T]], 'DataclassClassOrWrapper'], 'DataclassClassOrWrapper']: + """ + Like the python standard lib dataclasses but with type validation. + The result is either a pydantic dataclass that will validate input data + or a wrapper that will trigger validation around a stdlib dataclass + to avoid modifying it directly + """ + the_config = get_config(config) + + def wrap(cls: Type[Any]) -> 'DataclassClassOrWrapper': + should_use_proxy = ( + use_proxy + if use_proxy is not None + else ( + is_builtin_dataclass(cls) + and (cls.__bases__[0] is object or set(dir(cls)) == set(dir(cls.__bases__[0]))) + ) + ) + if should_use_proxy: + dc_cls_doc = '' + dc_cls = DataclassProxy(cls) + default_validate_on_init = False + else: + dc_cls_doc = cls.__doc__ or '' # needs to be done before generating dataclass + if sys.version_info >= (3, 10): + dc_cls = dataclasses.dataclass( + cls, + init=init, + repr=repr, + eq=eq, + order=order, + unsafe_hash=unsafe_hash, + frozen=frozen, + kw_only=kw_only, + ) + else: + dc_cls = dataclasses.dataclass( # type: ignore + cls, init=init, repr=repr, eq=eq, order=order, unsafe_hash=unsafe_hash, frozen=frozen + ) + default_validate_on_init = True + + should_validate_on_init = default_validate_on_init if validate_on_init is None else validate_on_init + _add_pydantic_validation_attributes(cls, the_config, should_validate_on_init, dc_cls_doc) + dc_cls.__pydantic_model__.__try_update_forward_refs__(**{cls.__name__: cls}) + return dc_cls + + if _cls is None: + return wrap + + return wrap(_cls) + + +@contextmanager +def set_validation(cls: Type['DataclassT'], value: bool) -> Generator[Type['DataclassT'], None, None]: + original_run_validation = cls.__pydantic_run_validation__ + try: + cls.__pydantic_run_validation__ = value + yield cls + finally: + cls.__pydantic_run_validation__ = original_run_validation + + +class DataclassProxy: + __slots__ = '__dataclass__' + + def __init__(self, dc_cls: Type['Dataclass']) -> None: + object.__setattr__(self, '__dataclass__', dc_cls) + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + with set_validation(self.__dataclass__, True): + return self.__dataclass__(*args, **kwargs) + + def __getattr__(self, name: str) -> Any: + return getattr(self.__dataclass__, name) + + def __setattr__(self, __name: str, __value: Any) -> None: + return setattr(self.__dataclass__, __name, __value) + + def __instancecheck__(self, instance: Any) -> bool: + return isinstance(instance, self.__dataclass__) + + def __copy__(self) -> 'DataclassProxy': + return DataclassProxy(copy.copy(self.__dataclass__)) + + def __deepcopy__(self, memo: Any) -> 'DataclassProxy': + return DataclassProxy(copy.deepcopy(self.__dataclass__, memo)) + + +def _add_pydantic_validation_attributes( # noqa: C901 (ignore complexity) + dc_cls: Type['Dataclass'], + config: Type[BaseConfig], + validate_on_init: bool, + dc_cls_doc: str, +) -> None: + """ + We need to replace the right method. If no `__post_init__` has been set in the stdlib dataclass + it won't even exist (code is generated on the fly by `dataclasses`) + By default, we run validation after `__init__` or `__post_init__` if defined + """ + init = dc_cls.__init__ + + @wraps(init) + def handle_extra_init(self: 'Dataclass', *args: Any, **kwargs: Any) -> None: + if config.extra == Extra.ignore: + init(self, *args, **{k: v for k, v in kwargs.items() if k in self.__dataclass_fields__}) + + elif config.extra == Extra.allow: + for k, v in kwargs.items(): + self.__dict__.setdefault(k, v) + init(self, *args, **{k: v for k, v in kwargs.items() if k in self.__dataclass_fields__}) + + else: + init(self, *args, **kwargs) + + if hasattr(dc_cls, '__post_init__'): + try: + post_init = dc_cls.__post_init__.__wrapped__ # type: ignore[attr-defined] + except AttributeError: + post_init = dc_cls.__post_init__ + + @wraps(post_init) + def new_post_init(self: 'Dataclass', *args: Any, **kwargs: Any) -> None: + if config.post_init_call == 'before_validation': + post_init(self, *args, **kwargs) + + if self.__class__.__pydantic_run_validation__: + self.__pydantic_validate_values__() + if hasattr(self, '__post_init_post_parse__'): + self.__post_init_post_parse__(*args, **kwargs) + + if config.post_init_call == 'after_validation': + post_init(self, *args, **kwargs) + + setattr(dc_cls, '__init__', handle_extra_init) + setattr(dc_cls, '__post_init__', new_post_init) + + else: + + @wraps(init) + def new_init(self: 'Dataclass', *args: Any, **kwargs: Any) -> None: + handle_extra_init(self, *args, **kwargs) + + if self.__class__.__pydantic_run_validation__: + self.__pydantic_validate_values__() + + if hasattr(self, '__post_init_post_parse__'): + # We need to find again the initvars. To do that we use `__dataclass_fields__` instead of + # public method `dataclasses.fields` + + # get all initvars and their default values + initvars_and_values: Dict[str, Any] = {} + for i, f in enumerate(self.__class__.__dataclass_fields__.values()): + if f._field_type is dataclasses._FIELD_INITVAR: # type: ignore[attr-defined] + try: + # set arg value by default + initvars_and_values[f.name] = args[i] + except IndexError: + initvars_and_values[f.name] = kwargs.get(f.name, f.default) + + self.__post_init_post_parse__(**initvars_and_values) + + setattr(dc_cls, '__init__', new_init) + + setattr(dc_cls, '__pydantic_run_validation__', ClassAttribute('__pydantic_run_validation__', validate_on_init)) + setattr(dc_cls, '__pydantic_initialised__', False) + setattr(dc_cls, '__pydantic_model__', create_pydantic_model_from_dataclass(dc_cls, config, dc_cls_doc)) + setattr(dc_cls, '__pydantic_validate_values__', _dataclass_validate_values) + setattr(dc_cls, '__validate__', classmethod(_validate_dataclass)) + setattr(dc_cls, '__get_validators__', classmethod(_get_validators)) + + if dc_cls.__pydantic_model__.__config__.validate_assignment and not dc_cls.__dataclass_params__.frozen: + setattr(dc_cls, '__setattr__', _dataclass_validate_assignment_setattr) + + +def _get_validators(cls: 'DataclassClassOrWrapper') -> 'CallableGenerator': + yield cls.__validate__ + + +def _validate_dataclass(cls: Type['DataclassT'], v: Any) -> 'DataclassT': + with set_validation(cls, True): + if isinstance(v, cls): + v.__pydantic_validate_values__() + return v + elif isinstance(v, (list, tuple)): + return cls(*v) + elif isinstance(v, dict): + return cls(**v) + else: + raise DataclassTypeError(class_name=cls.__name__) + + +def create_pydantic_model_from_dataclass( + dc_cls: Type['Dataclass'], + config: Type[Any] = BaseConfig, + dc_cls_doc: Optional[str] = None, +) -> Type['BaseModel']: + field_definitions: Dict[str, Any] = {} + for field in dataclasses.fields(dc_cls): + default: Any = Undefined + default_factory: Optional['NoArgAnyCallable'] = None + field_info: FieldInfo + + if field.default is not dataclasses.MISSING: + default = field.default + elif field.default_factory is not dataclasses.MISSING: + default_factory = field.default_factory + else: + default = Required + + if isinstance(default, FieldInfo): + field_info = default + dc_cls.__pydantic_has_field_info_default__ = True + else: + field_info = Field(default=default, default_factory=default_factory, **field.metadata) + + field_definitions[field.name] = (field.type, field_info) + + validators = gather_all_validators(dc_cls) + model: Type['BaseModel'] = create_model( + dc_cls.__name__, + __config__=config, + __module__=dc_cls.__module__, + __validators__=validators, + __cls_kwargs__={'__resolve_forward_refs__': False}, + **field_definitions, + ) + model.__doc__ = dc_cls_doc if dc_cls_doc is not None else dc_cls.__doc__ or '' + return model + + +if sys.version_info >= (3, 8): + + def _is_field_cached_property(obj: 'Dataclass', k: str) -> bool: + return isinstance(getattr(type(obj), k, None), cached_property) + +else: + + def _is_field_cached_property(obj: 'Dataclass', k: str) -> bool: + return False + + +def _dataclass_validate_values(self: 'Dataclass') -> None: + # validation errors can occur if this function is called twice on an already initialised dataclass. + # for example if Extra.forbid is enabled, it would consider __pydantic_initialised__ an invalid extra property + if getattr(self, '__pydantic_initialised__'): + return + if getattr(self, '__pydantic_has_field_info_default__', False): + # We need to remove `FieldInfo` values since they are not valid as input + # It's ok to do that because they are obviously the default values! + input_data = { + k: v + for k, v in self.__dict__.items() + if not (isinstance(v, FieldInfo) or _is_field_cached_property(self, k)) + } + else: + input_data = {k: v for k, v in self.__dict__.items() if not _is_field_cached_property(self, k)} + d, _, validation_error = validate_model(self.__pydantic_model__, input_data, cls=self.__class__) + if validation_error: + raise validation_error + self.__dict__.update(d) + object.__setattr__(self, '__pydantic_initialised__', True) + + +def _dataclass_validate_assignment_setattr(self: 'Dataclass', name: str, value: Any) -> None: + if self.__pydantic_initialised__: + d = dict(self.__dict__) + d.pop(name, None) + known_field = self.__pydantic_model__.__fields__.get(name, None) + if known_field: + value, error_ = known_field.validate(value, d, loc=name, cls=self.__class__) + if error_: + raise ValidationError([error_], self.__class__) + + object.__setattr__(self, name, value) + + +def is_builtin_dataclass(_cls: Type[Any]) -> bool: + """ + Whether a class is a stdlib dataclass + (useful to discriminated a pydantic dataclass that is actually a wrapper around a stdlib dataclass) + + we check that + - `_cls` is a dataclass + - `_cls` is not a processed pydantic dataclass (with a basemodel attached) + - `_cls` is not a pydantic dataclass inheriting directly from a stdlib dataclass + e.g. + ``` + @dataclasses.dataclass + class A: + x: int + + @pydantic.dataclasses.dataclass + class B(A): + y: int + ``` + In this case, when we first check `B`, we make an extra check and look at the annotations ('y'), + which won't be a superset of all the dataclass fields (only the stdlib fields i.e. 'x') + """ + return ( + dataclasses.is_dataclass(_cls) + and not hasattr(_cls, '__pydantic_model__') + and set(_cls.__dataclass_fields__).issuperset(set(getattr(_cls, '__annotations__', {}))) + ) + + +def make_dataclass_validator(dc_cls: Type['Dataclass'], config: Type[BaseConfig]) -> 'CallableGenerator': + """ + Create a pydantic.dataclass from a builtin dataclass to add type validation + and yield the validators + It retrieves the parameters of the dataclass and forwards them to the newly created dataclass + """ + yield from _get_validators(dataclass(dc_cls, config=config, use_proxy=True)) diff --git a/lib/python3.12/site-packages/pydantic/v1/datetime_parse.py b/lib/python3.12/site-packages/pydantic/v1/datetime_parse.py new file mode 100644 index 0000000000000000000000000000000000000000..a7598fc6c56688c4dd6526fc6e8b8e03682f62e8 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/datetime_parse.py @@ -0,0 +1,248 @@ +""" +Functions to parse datetime objects. + +We're using regular expressions rather than time.strptime because: +- They provide both validation and parsing. +- They're more flexible for datetimes. +- The date/datetime/time constructors produce friendlier error messages. + +Stolen from https://raw.githubusercontent.com/django/django/main/django/utils/dateparse.py at +9718fa2e8abe430c3526a9278dd976443d4ae3c6 + +Changed to: +* use standard python datetime types not django.utils.timezone +* raise ValueError when regex doesn't match rather than returning None +* support parsing unix timestamps for dates and datetimes +""" +import re +from datetime import date, datetime, time, timedelta, timezone +from typing import Dict, Optional, Type, Union + +from pydantic.v1 import errors + +date_expr = r'(?P\d{4})-(?P\d{1,2})-(?P\d{1,2})' +time_expr = ( + r'(?P\d{1,2}):(?P\d{1,2})' + r'(?::(?P\d{1,2})(?:\.(?P\d{1,6})\d{0,6})?)?' + r'(?PZ|[+-]\d{2}(?::?\d{2})?)?$' +) + +date_re = re.compile(f'{date_expr}$') +time_re = re.compile(time_expr) +datetime_re = re.compile(f'{date_expr}[T ]{time_expr}') + +standard_duration_re = re.compile( + r'^' + r'(?:(?P-?\d+) (days?, )?)?' + r'((?:(?P-?\d+):)(?=\d+:\d+))?' + r'(?:(?P-?\d+):)?' + r'(?P-?\d+)' + r'(?:\.(?P\d{1,6})\d{0,6})?' + r'$' +) + +# Support the sections of ISO 8601 date representation that are accepted by timedelta +iso8601_duration_re = re.compile( + r'^(?P[-+]?)' + r'P' + r'(?:(?P\d+(.\d+)?)D)?' + r'(?:T' + r'(?:(?P\d+(.\d+)?)H)?' + r'(?:(?P\d+(.\d+)?)M)?' + r'(?:(?P\d+(.\d+)?)S)?' + r')?' + r'$' +) + +EPOCH = datetime(1970, 1, 1) +# if greater than this, the number is in ms, if less than or equal it's in seconds +# (in seconds this is 11th October 2603, in ms it's 20th August 1970) +MS_WATERSHED = int(2e10) +# slightly more than datetime.max in ns - (datetime.max - EPOCH).total_seconds() * 1e9 +MAX_NUMBER = int(3e20) +StrBytesIntFloat = Union[str, bytes, int, float] + + +def get_numeric(value: StrBytesIntFloat, native_expected_type: str) -> Union[None, int, float]: + if isinstance(value, (int, float)): + return value + try: + return float(value) + except ValueError: + return None + except TypeError: + raise TypeError(f'invalid type; expected {native_expected_type}, string, bytes, int or float') + + +def from_unix_seconds(seconds: Union[int, float]) -> datetime: + if seconds > MAX_NUMBER: + return datetime.max + elif seconds < -MAX_NUMBER: + return datetime.min + + while abs(seconds) > MS_WATERSHED: + seconds /= 1000 + dt = EPOCH + timedelta(seconds=seconds) + return dt.replace(tzinfo=timezone.utc) + + +def _parse_timezone(value: Optional[str], error: Type[Exception]) -> Union[None, int, timezone]: + if value == 'Z': + return timezone.utc + elif value is not None: + offset_mins = int(value[-2:]) if len(value) > 3 else 0 + offset = 60 * int(value[1:3]) + offset_mins + if value[0] == '-': + offset = -offset + try: + return timezone(timedelta(minutes=offset)) + except ValueError: + raise error() + else: + return None + + +def parse_date(value: Union[date, StrBytesIntFloat]) -> date: + """ + Parse a date/int/float/string and return a datetime.date. + + Raise ValueError if the input is well formatted but not a valid date. + Raise ValueError if the input isn't well formatted. + """ + if isinstance(value, date): + if isinstance(value, datetime): + return value.date() + else: + return value + + number = get_numeric(value, 'date') + if number is not None: + return from_unix_seconds(number).date() + + if isinstance(value, bytes): + value = value.decode() + + match = date_re.match(value) # type: ignore + if match is None: + raise errors.DateError() + + kw = {k: int(v) for k, v in match.groupdict().items()} + + try: + return date(**kw) + except ValueError: + raise errors.DateError() + + +def parse_time(value: Union[time, StrBytesIntFloat]) -> time: + """ + Parse a time/string and return a datetime.time. + + Raise ValueError if the input is well formatted but not a valid time. + Raise ValueError if the input isn't well formatted, in particular if it contains an offset. + """ + if isinstance(value, time): + return value + + number = get_numeric(value, 'time') + if number is not None: + if number >= 86400: + # doesn't make sense since the time time loop back around to 0 + raise errors.TimeError() + return (datetime.min + timedelta(seconds=number)).time() + + if isinstance(value, bytes): + value = value.decode() + + match = time_re.match(value) # type: ignore + if match is None: + raise errors.TimeError() + + kw = match.groupdict() + if kw['microsecond']: + kw['microsecond'] = kw['microsecond'].ljust(6, '0') + + tzinfo = _parse_timezone(kw.pop('tzinfo'), errors.TimeError) + kw_: Dict[str, Union[None, int, timezone]] = {k: int(v) for k, v in kw.items() if v is not None} + kw_['tzinfo'] = tzinfo + + try: + return time(**kw_) # type: ignore + except ValueError: + raise errors.TimeError() + + +def parse_datetime(value: Union[datetime, StrBytesIntFloat]) -> datetime: + """ + Parse a datetime/int/float/string and return a datetime.datetime. + + This function supports time zone offsets. When the input contains one, + the output uses a timezone with a fixed offset from UTC. + + Raise ValueError if the input is well formatted but not a valid datetime. + Raise ValueError if the input isn't well formatted. + """ + if isinstance(value, datetime): + return value + + number = get_numeric(value, 'datetime') + if number is not None: + return from_unix_seconds(number) + + if isinstance(value, bytes): + value = value.decode() + + match = datetime_re.match(value) # type: ignore + if match is None: + raise errors.DateTimeError() + + kw = match.groupdict() + if kw['microsecond']: + kw['microsecond'] = kw['microsecond'].ljust(6, '0') + + tzinfo = _parse_timezone(kw.pop('tzinfo'), errors.DateTimeError) + kw_: Dict[str, Union[None, int, timezone]] = {k: int(v) for k, v in kw.items() if v is not None} + kw_['tzinfo'] = tzinfo + + try: + return datetime(**kw_) # type: ignore + except ValueError: + raise errors.DateTimeError() + + +def parse_duration(value: StrBytesIntFloat) -> timedelta: + """ + Parse a duration int/float/string and return a datetime.timedelta. + + The preferred format for durations in Django is '%d %H:%M:%S.%f'. + + Also supports ISO 8601 representation. + """ + if isinstance(value, timedelta): + return value + + if isinstance(value, (int, float)): + # below code requires a string + value = f'{value:f}' + elif isinstance(value, bytes): + value = value.decode() + + try: + match = standard_duration_re.match(value) or iso8601_duration_re.match(value) + except TypeError: + raise TypeError('invalid type; expected timedelta, string, bytes, int or float') + + if not match: + raise errors.DurationError() + + kw = match.groupdict() + sign = -1 if kw.pop('sign', '+') == '-' else 1 + if kw.get('microseconds'): + kw['microseconds'] = kw['microseconds'].ljust(6, '0') + + if kw.get('seconds') and kw.get('microseconds') and kw['seconds'].startswith('-'): + kw['microseconds'] = '-' + kw['microseconds'] + + kw_ = {k: float(v) for k, v in kw.items() if v is not None} + + return sign * timedelta(**kw_) diff --git a/lib/python3.12/site-packages/pydantic/v1/decorator.py b/lib/python3.12/site-packages/pydantic/v1/decorator.py new file mode 100644 index 0000000000000000000000000000000000000000..2c7c2c2ffdb45ca50cd2b7e57bde1fc711adb851 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/decorator.py @@ -0,0 +1,264 @@ +from functools import wraps +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, Tuple, Type, TypeVar, Union, overload + +from pydantic.v1 import validator +from pydantic.v1.config import Extra +from pydantic.v1.errors import ConfigError +from pydantic.v1.main import BaseModel, create_model +from pydantic.v1.typing import get_all_type_hints +from pydantic.v1.utils import to_camel + +__all__ = ('validate_arguments',) + +if TYPE_CHECKING: + from pydantic.v1.typing import AnyCallable + + AnyCallableT = TypeVar('AnyCallableT', bound=AnyCallable) + ConfigType = Union[None, Type[Any], Dict[str, Any]] + + +@overload +def validate_arguments(func: None = None, *, config: 'ConfigType' = None) -> Callable[['AnyCallableT'], 'AnyCallableT']: + ... + + +@overload +def validate_arguments(func: 'AnyCallableT') -> 'AnyCallableT': + ... + + +def validate_arguments(func: Optional['AnyCallableT'] = None, *, config: 'ConfigType' = None) -> Any: + """ + Decorator to validate the arguments passed to a function. + """ + + def validate(_func: 'AnyCallable') -> 'AnyCallable': + vd = ValidatedFunction(_func, config) + + @wraps(_func) + def wrapper_function(*args: Any, **kwargs: Any) -> Any: + return vd.call(*args, **kwargs) + + wrapper_function.vd = vd # type: ignore + wrapper_function.validate = vd.init_model_instance # type: ignore + wrapper_function.raw_function = vd.raw_function # type: ignore + wrapper_function.model = vd.model # type: ignore + return wrapper_function + + if func: + return validate(func) + else: + return validate + + +ALT_V_ARGS = 'v__args' +ALT_V_KWARGS = 'v__kwargs' +V_POSITIONAL_ONLY_NAME = 'v__positional_only' +V_DUPLICATE_KWARGS = 'v__duplicate_kwargs' + + +class ValidatedFunction: + def __init__(self, function: 'AnyCallableT', config: 'ConfigType'): # noqa C901 + from inspect import Parameter, signature + + parameters: Mapping[str, Parameter] = signature(function).parameters + + if parameters.keys() & {ALT_V_ARGS, ALT_V_KWARGS, V_POSITIONAL_ONLY_NAME, V_DUPLICATE_KWARGS}: + raise ConfigError( + f'"{ALT_V_ARGS}", "{ALT_V_KWARGS}", "{V_POSITIONAL_ONLY_NAME}" and "{V_DUPLICATE_KWARGS}" ' + f'are not permitted as argument names when using the "{validate_arguments.__name__}" decorator' + ) + + self.raw_function = function + self.arg_mapping: Dict[int, str] = {} + self.positional_only_args = set() + self.v_args_name = 'args' + self.v_kwargs_name = 'kwargs' + + type_hints = get_all_type_hints(function) + takes_args = False + takes_kwargs = False + fields: Dict[str, Tuple[Any, Any]] = {} + for i, (name, p) in enumerate(parameters.items()): + if p.annotation is p.empty: + annotation = Any + else: + annotation = type_hints[name] + + default = ... if p.default is p.empty else p.default + if p.kind == Parameter.POSITIONAL_ONLY: + self.arg_mapping[i] = name + fields[name] = annotation, default + fields[V_POSITIONAL_ONLY_NAME] = List[str], None + self.positional_only_args.add(name) + elif p.kind == Parameter.POSITIONAL_OR_KEYWORD: + self.arg_mapping[i] = name + fields[name] = annotation, default + fields[V_DUPLICATE_KWARGS] = List[str], None + elif p.kind == Parameter.KEYWORD_ONLY: + fields[name] = annotation, default + elif p.kind == Parameter.VAR_POSITIONAL: + self.v_args_name = name + fields[name] = Tuple[annotation, ...], None + takes_args = True + else: + assert p.kind == Parameter.VAR_KEYWORD, p.kind + self.v_kwargs_name = name + fields[name] = Dict[str, annotation], None # type: ignore + takes_kwargs = True + + # these checks avoid a clash between "args" and a field with that name + if not takes_args and self.v_args_name in fields: + self.v_args_name = ALT_V_ARGS + + # same with "kwargs" + if not takes_kwargs and self.v_kwargs_name in fields: + self.v_kwargs_name = ALT_V_KWARGS + + if not takes_args: + # we add the field so validation below can raise the correct exception + fields[self.v_args_name] = List[Any], None + + if not takes_kwargs: + # same with kwargs + fields[self.v_kwargs_name] = Dict[Any, Any], None + + self.create_model(fields, takes_args, takes_kwargs, config) + + def init_model_instance(self, *args: Any, **kwargs: Any) -> BaseModel: + values = self.build_values(args, kwargs) + return self.model(**values) + + def call(self, *args: Any, **kwargs: Any) -> Any: + m = self.init_model_instance(*args, **kwargs) + return self.execute(m) + + def build_values(self, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> Dict[str, Any]: + values: Dict[str, Any] = {} + if args: + arg_iter = enumerate(args) + while True: + try: + i, a = next(arg_iter) + except StopIteration: + break + arg_name = self.arg_mapping.get(i) + if arg_name is not None: + values[arg_name] = a + else: + values[self.v_args_name] = [a] + [a for _, a in arg_iter] + break + + var_kwargs: Dict[str, Any] = {} + wrong_positional_args = [] + duplicate_kwargs = [] + fields_alias = [ + field.alias + for name, field in self.model.__fields__.items() + if name not in (self.v_args_name, self.v_kwargs_name) + ] + non_var_fields = set(self.model.__fields__) - {self.v_args_name, self.v_kwargs_name} + for k, v in kwargs.items(): + if k in non_var_fields or k in fields_alias: + if k in self.positional_only_args: + wrong_positional_args.append(k) + if k in values: + duplicate_kwargs.append(k) + values[k] = v + else: + var_kwargs[k] = v + + if var_kwargs: + values[self.v_kwargs_name] = var_kwargs + if wrong_positional_args: + values[V_POSITIONAL_ONLY_NAME] = wrong_positional_args + if duplicate_kwargs: + values[V_DUPLICATE_KWARGS] = duplicate_kwargs + return values + + def execute(self, m: BaseModel) -> Any: + d = {k: v for k, v in m._iter() if k in m.__fields_set__ or m.__fields__[k].default_factory} + var_kwargs = d.pop(self.v_kwargs_name, {}) + + if self.v_args_name in d: + args_: List[Any] = [] + in_kwargs = False + kwargs = {} + for name, value in d.items(): + if in_kwargs: + kwargs[name] = value + elif name == self.v_args_name: + args_ += value + in_kwargs = True + else: + args_.append(value) + return self.raw_function(*args_, **kwargs, **var_kwargs) + elif self.positional_only_args: + args_ = [] + kwargs = {} + for name, value in d.items(): + if name in self.positional_only_args: + args_.append(value) + else: + kwargs[name] = value + return self.raw_function(*args_, **kwargs, **var_kwargs) + else: + return self.raw_function(**d, **var_kwargs) + + def create_model(self, fields: Dict[str, Any], takes_args: bool, takes_kwargs: bool, config: 'ConfigType') -> None: + pos_args = len(self.arg_mapping) + + class CustomConfig: + pass + + if not TYPE_CHECKING: # pragma: no branch + if isinstance(config, dict): + CustomConfig = type('Config', (), config) # noqa: F811 + elif config is not None: + CustomConfig = config # noqa: F811 + + if hasattr(CustomConfig, 'fields') or hasattr(CustomConfig, 'alias_generator'): + raise ConfigError( + 'Setting the "fields" and "alias_generator" property on custom Config for ' + '@validate_arguments is not yet supported, please remove.' + ) + + class DecoratorBaseModel(BaseModel): + @validator(self.v_args_name, check_fields=False, allow_reuse=True) + def check_args(cls, v: Optional[List[Any]]) -> Optional[List[Any]]: + if takes_args or v is None: + return v + + raise TypeError(f'{pos_args} positional arguments expected but {pos_args + len(v)} given') + + @validator(self.v_kwargs_name, check_fields=False, allow_reuse=True) + def check_kwargs(cls, v: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + if takes_kwargs or v is None: + return v + + plural = '' if len(v) == 1 else 's' + keys = ', '.join(map(repr, v.keys())) + raise TypeError(f'unexpected keyword argument{plural}: {keys}') + + @validator(V_POSITIONAL_ONLY_NAME, check_fields=False, allow_reuse=True) + def check_positional_only(cls, v: Optional[List[str]]) -> None: + if v is None: + return + + plural = '' if len(v) == 1 else 's' + keys = ', '.join(map(repr, v)) + raise TypeError(f'positional-only argument{plural} passed as keyword argument{plural}: {keys}') + + @validator(V_DUPLICATE_KWARGS, check_fields=False, allow_reuse=True) + def check_duplicate_kwargs(cls, v: Optional[List[str]]) -> None: + if v is None: + return + + plural = '' if len(v) == 1 else 's' + keys = ', '.join(map(repr, v)) + raise TypeError(f'multiple values for argument{plural}: {keys}') + + class Config(CustomConfig): + extra = getattr(CustomConfig, 'extra', Extra.forbid) + + self.model = create_model(to_camel(self.raw_function.__name__), __base__=DecoratorBaseModel, **fields) diff --git a/lib/python3.12/site-packages/pydantic/v1/env_settings.py b/lib/python3.12/site-packages/pydantic/v1/env_settings.py new file mode 100644 index 0000000000000000000000000000000000000000..5f6f217500d8ca1e7717f59ae372ee9a27660fe3 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/env_settings.py @@ -0,0 +1,350 @@ +import os +import warnings +from pathlib import Path +from typing import AbstractSet, Any, Callable, ClassVar, Dict, List, Mapping, Optional, Tuple, Type, Union + +from pydantic.v1.config import BaseConfig, Extra +from pydantic.v1.fields import ModelField +from pydantic.v1.main import BaseModel +from pydantic.v1.types import JsonWrapper +from pydantic.v1.typing import StrPath, display_as_type, get_origin, is_union +from pydantic.v1.utils import deep_update, lenient_issubclass, path_type, sequence_like + +env_file_sentinel = str(object()) + +SettingsSourceCallable = Callable[['BaseSettings'], Dict[str, Any]] +DotenvType = Union[StrPath, List[StrPath], Tuple[StrPath, ...]] + + +class SettingsError(ValueError): + pass + + +class BaseSettings(BaseModel): + """ + Base class for settings, allowing values to be overridden by environment variables. + + This is useful in production for secrets you do not wish to save in code, it plays nicely with docker(-compose), + Heroku and any 12 factor app design. + """ + + def __init__( + __pydantic_self__, + _env_file: Optional[DotenvType] = env_file_sentinel, + _env_file_encoding: Optional[str] = None, + _env_nested_delimiter: Optional[str] = None, + _secrets_dir: Optional[StrPath] = None, + **values: Any, + ) -> None: + # Uses something other than `self` the first arg to allow "self" as a settable attribute + super().__init__( + **__pydantic_self__._build_values( + values, + _env_file=_env_file, + _env_file_encoding=_env_file_encoding, + _env_nested_delimiter=_env_nested_delimiter, + _secrets_dir=_secrets_dir, + ) + ) + + def _build_values( + self, + init_kwargs: Dict[str, Any], + _env_file: Optional[DotenvType] = None, + _env_file_encoding: Optional[str] = None, + _env_nested_delimiter: Optional[str] = None, + _secrets_dir: Optional[StrPath] = None, + ) -> Dict[str, Any]: + # Configure built-in sources + init_settings = InitSettingsSource(init_kwargs=init_kwargs) + env_settings = EnvSettingsSource( + env_file=(_env_file if _env_file != env_file_sentinel else self.__config__.env_file), + env_file_encoding=( + _env_file_encoding if _env_file_encoding is not None else self.__config__.env_file_encoding + ), + env_nested_delimiter=( + _env_nested_delimiter if _env_nested_delimiter is not None else self.__config__.env_nested_delimiter + ), + env_prefix_len=len(self.__config__.env_prefix), + ) + file_secret_settings = SecretsSettingsSource(secrets_dir=_secrets_dir or self.__config__.secrets_dir) + # Provide a hook to set built-in sources priority and add / remove sources + sources = self.__config__.customise_sources( + init_settings=init_settings, env_settings=env_settings, file_secret_settings=file_secret_settings + ) + if sources: + return deep_update(*reversed([source(self) for source in sources])) + else: + # no one should mean to do this, but I think returning an empty dict is marginally preferable + # to an informative error and much better than a confusing error + return {} + + class Config(BaseConfig): + env_prefix: str = '' + env_file: Optional[DotenvType] = None + env_file_encoding: Optional[str] = None + env_nested_delimiter: Optional[str] = None + secrets_dir: Optional[StrPath] = None + validate_all: bool = True + extra: Extra = Extra.forbid + arbitrary_types_allowed: bool = True + case_sensitive: bool = False + + @classmethod + def prepare_field(cls, field: ModelField) -> None: + env_names: Union[List[str], AbstractSet[str]] + field_info_from_config = cls.get_field_info(field.name) + + env = field_info_from_config.get('env') or field.field_info.extra.get('env') + if env is None: + if field.has_alias: + warnings.warn( + 'aliases are no longer used by BaseSettings to define which environment variables to read. ' + 'Instead use the "env" field setting. ' + 'See https://pydantic-docs.helpmanual.io/usage/settings/#environment-variable-names', + FutureWarning, + ) + env_names = {cls.env_prefix + field.name} + elif isinstance(env, str): + env_names = {env} + elif isinstance(env, (set, frozenset)): + env_names = env + elif sequence_like(env): + env_names = list(env) + else: + raise TypeError(f'invalid field env: {env!r} ({display_as_type(env)}); should be string, list or set') + + if not cls.case_sensitive: + env_names = env_names.__class__(n.lower() for n in env_names) + field.field_info.extra['env_names'] = env_names + + @classmethod + def customise_sources( + cls, + init_settings: SettingsSourceCallable, + env_settings: SettingsSourceCallable, + file_secret_settings: SettingsSourceCallable, + ) -> Tuple[SettingsSourceCallable, ...]: + return init_settings, env_settings, file_secret_settings + + @classmethod + def parse_env_var(cls, field_name: str, raw_val: str) -> Any: + return cls.json_loads(raw_val) + + # populated by the metaclass using the Config class defined above, annotated here to help IDEs only + __config__: ClassVar[Type[Config]] + + +class InitSettingsSource: + __slots__ = ('init_kwargs',) + + def __init__(self, init_kwargs: Dict[str, Any]): + self.init_kwargs = init_kwargs + + def __call__(self, settings: BaseSettings) -> Dict[str, Any]: + return self.init_kwargs + + def __repr__(self) -> str: + return f'InitSettingsSource(init_kwargs={self.init_kwargs!r})' + + +class EnvSettingsSource: + __slots__ = ('env_file', 'env_file_encoding', 'env_nested_delimiter', 'env_prefix_len') + + def __init__( + self, + env_file: Optional[DotenvType], + env_file_encoding: Optional[str], + env_nested_delimiter: Optional[str] = None, + env_prefix_len: int = 0, + ): + self.env_file: Optional[DotenvType] = env_file + self.env_file_encoding: Optional[str] = env_file_encoding + self.env_nested_delimiter: Optional[str] = env_nested_delimiter + self.env_prefix_len: int = env_prefix_len + + def __call__(self, settings: BaseSettings) -> Dict[str, Any]: # noqa C901 + """ + Build environment variables suitable for passing to the Model. + """ + d: Dict[str, Any] = {} + + if settings.__config__.case_sensitive: + env_vars: Mapping[str, Optional[str]] = os.environ + else: + env_vars = {k.lower(): v for k, v in os.environ.items()} + + dotenv_vars = self._read_env_files(settings.__config__.case_sensitive) + if dotenv_vars: + env_vars = {**dotenv_vars, **env_vars} + + for field in settings.__fields__.values(): + env_val: Optional[str] = None + for env_name in field.field_info.extra['env_names']: + env_val = env_vars.get(env_name) + if env_val is not None: + break + + is_complex, allow_parse_failure = self.field_is_complex(field) + if is_complex: + if env_val is None: + # field is complex but no value found so far, try explode_env_vars + env_val_built = self.explode_env_vars(field, env_vars) + if env_val_built: + d[field.alias] = env_val_built + else: + # field is complex and there's a value, decode that as JSON, then add explode_env_vars + try: + env_val = settings.__config__.parse_env_var(field.name, env_val) + except ValueError as e: + if not allow_parse_failure: + raise SettingsError(f'error parsing env var "{env_name}"') from e + + if isinstance(env_val, dict): + d[field.alias] = deep_update(env_val, self.explode_env_vars(field, env_vars)) + else: + d[field.alias] = env_val + elif env_val is not None: + # simplest case, field is not complex, we only need to add the value if it was found + d[field.alias] = env_val + + return d + + def _read_env_files(self, case_sensitive: bool) -> Dict[str, Optional[str]]: + env_files = self.env_file + if env_files is None: + return {} + + if isinstance(env_files, (str, os.PathLike)): + env_files = [env_files] + + dotenv_vars = {} + for env_file in env_files: + env_path = Path(env_file).expanduser() + if env_path.is_file(): + dotenv_vars.update( + read_env_file(env_path, encoding=self.env_file_encoding, case_sensitive=case_sensitive) + ) + + return dotenv_vars + + def field_is_complex(self, field: ModelField) -> Tuple[bool, bool]: + """ + Find out if a field is complex, and if so whether JSON errors should be ignored + """ + if lenient_issubclass(field.annotation, JsonWrapper): + return False, False + + if field.is_complex(): + allow_parse_failure = False + elif is_union(get_origin(field.type_)) and field.sub_fields and any(f.is_complex() for f in field.sub_fields): + allow_parse_failure = True + else: + return False, False + + return True, allow_parse_failure + + def explode_env_vars(self, field: ModelField, env_vars: Mapping[str, Optional[str]]) -> Dict[str, Any]: + """ + Process env_vars and extract the values of keys containing env_nested_delimiter into nested dictionaries. + + This is applied to a single field, hence filtering by env_var prefix. + """ + prefixes = [f'{env_name}{self.env_nested_delimiter}' for env_name in field.field_info.extra['env_names']] + result: Dict[str, Any] = {} + for env_name, env_val in env_vars.items(): + if not any(env_name.startswith(prefix) for prefix in prefixes): + continue + # we remove the prefix before splitting in case the prefix has characters in common with the delimiter + env_name_without_prefix = env_name[self.env_prefix_len :] + _, *keys, last_key = env_name_without_prefix.split(self.env_nested_delimiter) + env_var = result + for key in keys: + env_var = env_var.setdefault(key, {}) + env_var[last_key] = env_val + + return result + + def __repr__(self) -> str: + return ( + f'EnvSettingsSource(env_file={self.env_file!r}, env_file_encoding={self.env_file_encoding!r}, ' + f'env_nested_delimiter={self.env_nested_delimiter!r})' + ) + + +class SecretsSettingsSource: + __slots__ = ('secrets_dir',) + + def __init__(self, secrets_dir: Optional[StrPath]): + self.secrets_dir: Optional[StrPath] = secrets_dir + + def __call__(self, settings: BaseSettings) -> Dict[str, Any]: + """ + Build fields from "secrets" files. + """ + secrets: Dict[str, Optional[str]] = {} + + if self.secrets_dir is None: + return secrets + + secrets_path = Path(self.secrets_dir).expanduser() + + if not secrets_path.exists(): + warnings.warn(f'directory "{secrets_path}" does not exist') + return secrets + + if not secrets_path.is_dir(): + raise SettingsError(f'secrets_dir must reference a directory, not a {path_type(secrets_path)}') + + for field in settings.__fields__.values(): + for env_name in field.field_info.extra['env_names']: + path = find_case_path(secrets_path, env_name, settings.__config__.case_sensitive) + if not path: + # path does not exist, we currently don't return a warning for this + continue + + if path.is_file(): + secret_value = path.read_text().strip() + if field.is_complex(): + try: + secret_value = settings.__config__.parse_env_var(field.name, secret_value) + except ValueError as e: + raise SettingsError(f'error parsing env var "{env_name}"') from e + + secrets[field.alias] = secret_value + else: + warnings.warn( + f'attempted to load secret file "{path}" but found a {path_type(path)} instead.', + stacklevel=4, + ) + return secrets + + def __repr__(self) -> str: + return f'SecretsSettingsSource(secrets_dir={self.secrets_dir!r})' + + +def read_env_file( + file_path: StrPath, *, encoding: str = None, case_sensitive: bool = False +) -> Dict[str, Optional[str]]: + try: + from dotenv import dotenv_values + except ImportError as e: + raise ImportError('python-dotenv is not installed, run `pip install pydantic[dotenv]`') from e + + file_vars: Dict[str, Optional[str]] = dotenv_values(file_path, encoding=encoding or 'utf8') + if not case_sensitive: + return {k.lower(): v for k, v in file_vars.items()} + else: + return file_vars + + +def find_case_path(dir_path: Path, file_name: str, case_sensitive: bool) -> Optional[Path]: + """ + Find a file within path's directory matching filename, optionally ignoring case. + """ + for f in dir_path.iterdir(): + if f.name == file_name: + return f + elif not case_sensitive and f.name.lower() == file_name.lower(): + return f + return None diff --git a/lib/python3.12/site-packages/pydantic/v1/error_wrappers.py b/lib/python3.12/site-packages/pydantic/v1/error_wrappers.py new file mode 100644 index 0000000000000000000000000000000000000000..bc7f263146ee4f1701391a2436ab38e5e16f0e1a --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/error_wrappers.py @@ -0,0 +1,161 @@ +import json +from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Sequence, Tuple, Type, Union + +from pydantic.v1.json import pydantic_encoder +from pydantic.v1.utils import Representation + +if TYPE_CHECKING: + from typing_extensions import TypedDict + + from pydantic.v1.config import BaseConfig + from pydantic.v1.types import ModelOrDc + from pydantic.v1.typing import ReprArgs + + Loc = Tuple[Union[int, str], ...] + + class _ErrorDictRequired(TypedDict): + loc: Loc + msg: str + type: str + + class ErrorDict(_ErrorDictRequired, total=False): + ctx: Dict[str, Any] + + +__all__ = 'ErrorWrapper', 'ValidationError' + + +class ErrorWrapper(Representation): + __slots__ = 'exc', '_loc' + + def __init__(self, exc: Exception, loc: Union[str, 'Loc']) -> None: + self.exc = exc + self._loc = loc + + def loc_tuple(self) -> 'Loc': + if isinstance(self._loc, tuple): + return self._loc + else: + return (self._loc,) + + def __repr_args__(self) -> 'ReprArgs': + return [('exc', self.exc), ('loc', self.loc_tuple())] + + +# ErrorList is something like Union[List[Union[List[ErrorWrapper], ErrorWrapper]], ErrorWrapper] +# but recursive, therefore just use: +ErrorList = Union[Sequence[Any], ErrorWrapper] + + +class ValidationError(Representation, ValueError): + __slots__ = 'raw_errors', 'model', '_error_cache' + + def __init__(self, errors: Sequence[ErrorList], model: 'ModelOrDc') -> None: + self.raw_errors = errors + self.model = model + self._error_cache: Optional[List['ErrorDict']] = None + + def errors(self) -> List['ErrorDict']: + if self._error_cache is None: + try: + config = self.model.__config__ # type: ignore + except AttributeError: + config = self.model.__pydantic_model__.__config__ # type: ignore + self._error_cache = list(flatten_errors(self.raw_errors, config)) + return self._error_cache + + def json(self, *, indent: Union[None, int, str] = 2) -> str: + return json.dumps(self.errors(), indent=indent, default=pydantic_encoder) + + def __str__(self) -> str: + errors = self.errors() + no_errors = len(errors) + return ( + f'{no_errors} validation error{"" if no_errors == 1 else "s"} for {self.model.__name__}\n' + f'{display_errors(errors)}' + ) + + def __repr_args__(self) -> 'ReprArgs': + return [('model', self.model.__name__), ('errors', self.errors())] + + +def display_errors(errors: List['ErrorDict']) -> str: + return '\n'.join(f'{_display_error_loc(e)}\n {e["msg"]} ({_display_error_type_and_ctx(e)})' for e in errors) + + +def _display_error_loc(error: 'ErrorDict') -> str: + return ' -> '.join(str(e) for e in error['loc']) + + +def _display_error_type_and_ctx(error: 'ErrorDict') -> str: + t = 'type=' + error['type'] + ctx = error.get('ctx') + if ctx: + return t + ''.join(f'; {k}={v}' for k, v in ctx.items()) + else: + return t + + +def flatten_errors( + errors: Sequence[Any], config: Type['BaseConfig'], loc: Optional['Loc'] = None +) -> Generator['ErrorDict', None, None]: + for error in errors: + if isinstance(error, ErrorWrapper): + if loc: + error_loc = loc + error.loc_tuple() + else: + error_loc = error.loc_tuple() + + if isinstance(error.exc, ValidationError): + yield from flatten_errors(error.exc.raw_errors, config, error_loc) + else: + yield error_dict(error.exc, config, error_loc) + elif isinstance(error, list): + yield from flatten_errors(error, config, loc=loc) + else: + raise RuntimeError(f'Unknown error object: {error}') + + +def error_dict(exc: Exception, config: Type['BaseConfig'], loc: 'Loc') -> 'ErrorDict': + type_ = get_exc_type(exc.__class__) + msg_template = config.error_msg_templates.get(type_) or getattr(exc, 'msg_template', None) + ctx = exc.__dict__ + if msg_template: + msg = msg_template.format(**ctx) + else: + msg = str(exc) + + d: 'ErrorDict' = {'loc': loc, 'msg': msg, 'type': type_} + + if ctx: + d['ctx'] = ctx + + return d + + +_EXC_TYPE_CACHE: Dict[Type[Exception], str] = {} + + +def get_exc_type(cls: Type[Exception]) -> str: + # slightly more efficient than using lru_cache since we don't need to worry about the cache filling up + try: + return _EXC_TYPE_CACHE[cls] + except KeyError: + r = _get_exc_type(cls) + _EXC_TYPE_CACHE[cls] = r + return r + + +def _get_exc_type(cls: Type[Exception]) -> str: + if issubclass(cls, AssertionError): + return 'assertion_error' + + base_name = 'type_error' if issubclass(cls, TypeError) else 'value_error' + if cls in (TypeError, ValueError): + # just TypeError or ValueError, no extra code + return base_name + + # if it's not a TypeError or ValueError, we just take the lowercase of the exception name + # no chaining or snake case logic, use "code" for more complex error types. + code = getattr(cls, 'code', None) or cls.__name__.replace('Error', '').lower() + return base_name + '.' + code diff --git a/lib/python3.12/site-packages/pydantic/v1/errors.py b/lib/python3.12/site-packages/pydantic/v1/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..6e8644258a6ccf999a19f69c359452a109e54ffe --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/errors.py @@ -0,0 +1,646 @@ +from decimal import Decimal +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Sequence, Set, Tuple, Type, Union + +from pydantic.v1.typing import display_as_type + +if TYPE_CHECKING: + from pydantic.v1.typing import DictStrAny + +# explicitly state exports to avoid "from pydantic.v1.errors import *" also importing Decimal, Path etc. +__all__ = ( + 'PydanticTypeError', + 'PydanticValueError', + 'ConfigError', + 'MissingError', + 'ExtraError', + 'NoneIsNotAllowedError', + 'NoneIsAllowedError', + 'WrongConstantError', + 'NotNoneError', + 'BoolError', + 'BytesError', + 'DictError', + 'EmailError', + 'UrlError', + 'UrlSchemeError', + 'UrlSchemePermittedError', + 'UrlUserInfoError', + 'UrlHostError', + 'UrlHostTldError', + 'UrlPortError', + 'UrlExtraError', + 'EnumError', + 'IntEnumError', + 'EnumMemberError', + 'IntegerError', + 'FloatError', + 'PathError', + 'PathNotExistsError', + 'PathNotAFileError', + 'PathNotADirectoryError', + 'PyObjectError', + 'SequenceError', + 'ListError', + 'SetError', + 'FrozenSetError', + 'TupleError', + 'TupleLengthError', + 'ListMinLengthError', + 'ListMaxLengthError', + 'ListUniqueItemsError', + 'SetMinLengthError', + 'SetMaxLengthError', + 'FrozenSetMinLengthError', + 'FrozenSetMaxLengthError', + 'AnyStrMinLengthError', + 'AnyStrMaxLengthError', + 'StrError', + 'StrRegexError', + 'NumberNotGtError', + 'NumberNotGeError', + 'NumberNotLtError', + 'NumberNotLeError', + 'NumberNotMultipleError', + 'DecimalError', + 'DecimalIsNotFiniteError', + 'DecimalMaxDigitsError', + 'DecimalMaxPlacesError', + 'DecimalWholeDigitsError', + 'DateTimeError', + 'DateError', + 'DateNotInThePastError', + 'DateNotInTheFutureError', + 'TimeError', + 'DurationError', + 'HashableError', + 'UUIDError', + 'UUIDVersionError', + 'ArbitraryTypeError', + 'ClassError', + 'SubclassError', + 'JsonError', + 'JsonTypeError', + 'PatternError', + 'DataclassTypeError', + 'CallableError', + 'IPvAnyAddressError', + 'IPvAnyInterfaceError', + 'IPvAnyNetworkError', + 'IPv4AddressError', + 'IPv6AddressError', + 'IPv4NetworkError', + 'IPv6NetworkError', + 'IPv4InterfaceError', + 'IPv6InterfaceError', + 'ColorError', + 'StrictBoolError', + 'NotDigitError', + 'LuhnValidationError', + 'InvalidLengthForBrand', + 'InvalidByteSize', + 'InvalidByteSizeUnit', + 'MissingDiscriminator', + 'InvalidDiscriminator', +) + + +def cls_kwargs(cls: Type['PydanticErrorMixin'], ctx: 'DictStrAny') -> 'PydanticErrorMixin': + """ + For built-in exceptions like ValueError or TypeError, we need to implement + __reduce__ to override the default behaviour (instead of __getstate__/__setstate__) + By default pickle protocol 2 calls `cls.__new__(cls, *args)`. + Since we only use kwargs, we need a little constructor to change that. + Note: the callable can't be a lambda as pickle looks in the namespace to find it + """ + return cls(**ctx) + + +class PydanticErrorMixin: + code: str + msg_template: str + + def __init__(self, **ctx: Any) -> None: + self.__dict__ = ctx + + def __str__(self) -> str: + return self.msg_template.format(**self.__dict__) + + def __reduce__(self) -> Tuple[Callable[..., 'PydanticErrorMixin'], Tuple[Type['PydanticErrorMixin'], 'DictStrAny']]: + return cls_kwargs, (self.__class__, self.__dict__) + + +class PydanticTypeError(PydanticErrorMixin, TypeError): + pass + + +class PydanticValueError(PydanticErrorMixin, ValueError): + pass + + +class ConfigError(RuntimeError): + pass + + +class MissingError(PydanticValueError): + msg_template = 'field required' + + +class ExtraError(PydanticValueError): + msg_template = 'extra fields not permitted' + + +class NoneIsNotAllowedError(PydanticTypeError): + code = 'none.not_allowed' + msg_template = 'none is not an allowed value' + + +class NoneIsAllowedError(PydanticTypeError): + code = 'none.allowed' + msg_template = 'value is not none' + + +class WrongConstantError(PydanticValueError): + code = 'const' + + def __str__(self) -> str: + permitted = ', '.join(repr(v) for v in self.permitted) # type: ignore + return f'unexpected value; permitted: {permitted}' + + +class NotNoneError(PydanticTypeError): + code = 'not_none' + msg_template = 'value is not None' + + +class BoolError(PydanticTypeError): + msg_template = 'value could not be parsed to a boolean' + + +class BytesError(PydanticTypeError): + msg_template = 'byte type expected' + + +class DictError(PydanticTypeError): + msg_template = 'value is not a valid dict' + + +class EmailError(PydanticValueError): + msg_template = 'value is not a valid email address' + + +class UrlError(PydanticValueError): + code = 'url' + + +class UrlSchemeError(UrlError): + code = 'url.scheme' + msg_template = 'invalid or missing URL scheme' + + +class UrlSchemePermittedError(UrlError): + code = 'url.scheme' + msg_template = 'URL scheme not permitted' + + def __init__(self, allowed_schemes: Set[str]): + super().__init__(allowed_schemes=allowed_schemes) + + +class UrlUserInfoError(UrlError): + code = 'url.userinfo' + msg_template = 'userinfo required in URL but missing' + + +class UrlHostError(UrlError): + code = 'url.host' + msg_template = 'URL host invalid' + + +class UrlHostTldError(UrlError): + code = 'url.host' + msg_template = 'URL host invalid, top level domain required' + + +class UrlPortError(UrlError): + code = 'url.port' + msg_template = 'URL port invalid, port cannot exceed 65535' + + +class UrlExtraError(UrlError): + code = 'url.extra' + msg_template = 'URL invalid, extra characters found after valid URL: {extra!r}' + + +class EnumMemberError(PydanticTypeError): + code = 'enum' + + def __str__(self) -> str: + permitted = ', '.join(repr(v.value) for v in self.enum_values) # type: ignore + return f'value is not a valid enumeration member; permitted: {permitted}' + + +class IntegerError(PydanticTypeError): + msg_template = 'value is not a valid integer' + + +class FloatError(PydanticTypeError): + msg_template = 'value is not a valid float' + + +class PathError(PydanticTypeError): + msg_template = 'value is not a valid path' + + +class _PathValueError(PydanticValueError): + def __init__(self, *, path: Path) -> None: + super().__init__(path=str(path)) + + +class PathNotExistsError(_PathValueError): + code = 'path.not_exists' + msg_template = 'file or directory at path "{path}" does not exist' + + +class PathNotAFileError(_PathValueError): + code = 'path.not_a_file' + msg_template = 'path "{path}" does not point to a file' + + +class PathNotADirectoryError(_PathValueError): + code = 'path.not_a_directory' + msg_template = 'path "{path}" does not point to a directory' + + +class PyObjectError(PydanticTypeError): + msg_template = 'ensure this value contains valid import path or valid callable: {error_message}' + + +class SequenceError(PydanticTypeError): + msg_template = 'value is not a valid sequence' + + +class IterableError(PydanticTypeError): + msg_template = 'value is not a valid iterable' + + +class ListError(PydanticTypeError): + msg_template = 'value is not a valid list' + + +class SetError(PydanticTypeError): + msg_template = 'value is not a valid set' + + +class FrozenSetError(PydanticTypeError): + msg_template = 'value is not a valid frozenset' + + +class DequeError(PydanticTypeError): + msg_template = 'value is not a valid deque' + + +class TupleError(PydanticTypeError): + msg_template = 'value is not a valid tuple' + + +class TupleLengthError(PydanticValueError): + code = 'tuple.length' + msg_template = 'wrong tuple length {actual_length}, expected {expected_length}' + + def __init__(self, *, actual_length: int, expected_length: int) -> None: + super().__init__(actual_length=actual_length, expected_length=expected_length) + + +class ListMinLengthError(PydanticValueError): + code = 'list.min_items' + msg_template = 'ensure this value has at least {limit_value} items' + + def __init__(self, *, limit_value: int) -> None: + super().__init__(limit_value=limit_value) + + +class ListMaxLengthError(PydanticValueError): + code = 'list.max_items' + msg_template = 'ensure this value has at most {limit_value} items' + + def __init__(self, *, limit_value: int) -> None: + super().__init__(limit_value=limit_value) + + +class ListUniqueItemsError(PydanticValueError): + code = 'list.unique_items' + msg_template = 'the list has duplicated items' + + +class SetMinLengthError(PydanticValueError): + code = 'set.min_items' + msg_template = 'ensure this value has at least {limit_value} items' + + def __init__(self, *, limit_value: int) -> None: + super().__init__(limit_value=limit_value) + + +class SetMaxLengthError(PydanticValueError): + code = 'set.max_items' + msg_template = 'ensure this value has at most {limit_value} items' + + def __init__(self, *, limit_value: int) -> None: + super().__init__(limit_value=limit_value) + + +class FrozenSetMinLengthError(PydanticValueError): + code = 'frozenset.min_items' + msg_template = 'ensure this value has at least {limit_value} items' + + def __init__(self, *, limit_value: int) -> None: + super().__init__(limit_value=limit_value) + + +class FrozenSetMaxLengthError(PydanticValueError): + code = 'frozenset.max_items' + msg_template = 'ensure this value has at most {limit_value} items' + + def __init__(self, *, limit_value: int) -> None: + super().__init__(limit_value=limit_value) + + +class AnyStrMinLengthError(PydanticValueError): + code = 'any_str.min_length' + msg_template = 'ensure this value has at least {limit_value} characters' + + def __init__(self, *, limit_value: int) -> None: + super().__init__(limit_value=limit_value) + + +class AnyStrMaxLengthError(PydanticValueError): + code = 'any_str.max_length' + msg_template = 'ensure this value has at most {limit_value} characters' + + def __init__(self, *, limit_value: int) -> None: + super().__init__(limit_value=limit_value) + + +class StrError(PydanticTypeError): + msg_template = 'str type expected' + + +class StrRegexError(PydanticValueError): + code = 'str.regex' + msg_template = 'string does not match regex "{pattern}"' + + def __init__(self, *, pattern: str) -> None: + super().__init__(pattern=pattern) + + +class _NumberBoundError(PydanticValueError): + def __init__(self, *, limit_value: Union[int, float, Decimal]) -> None: + super().__init__(limit_value=limit_value) + + +class NumberNotGtError(_NumberBoundError): + code = 'number.not_gt' + msg_template = 'ensure this value is greater than {limit_value}' + + +class NumberNotGeError(_NumberBoundError): + code = 'number.not_ge' + msg_template = 'ensure this value is greater than or equal to {limit_value}' + + +class NumberNotLtError(_NumberBoundError): + code = 'number.not_lt' + msg_template = 'ensure this value is less than {limit_value}' + + +class NumberNotLeError(_NumberBoundError): + code = 'number.not_le' + msg_template = 'ensure this value is less than or equal to {limit_value}' + + +class NumberNotFiniteError(PydanticValueError): + code = 'number.not_finite_number' + msg_template = 'ensure this value is a finite number' + + +class NumberNotMultipleError(PydanticValueError): + code = 'number.not_multiple' + msg_template = 'ensure this value is a multiple of {multiple_of}' + + def __init__(self, *, multiple_of: Union[int, float, Decimal]) -> None: + super().__init__(multiple_of=multiple_of) + + +class DecimalError(PydanticTypeError): + msg_template = 'value is not a valid decimal' + + +class DecimalIsNotFiniteError(PydanticValueError): + code = 'decimal.not_finite' + msg_template = 'value is not a valid decimal' + + +class DecimalMaxDigitsError(PydanticValueError): + code = 'decimal.max_digits' + msg_template = 'ensure that there are no more than {max_digits} digits in total' + + def __init__(self, *, max_digits: int) -> None: + super().__init__(max_digits=max_digits) + + +class DecimalMaxPlacesError(PydanticValueError): + code = 'decimal.max_places' + msg_template = 'ensure that there are no more than {decimal_places} decimal places' + + def __init__(self, *, decimal_places: int) -> None: + super().__init__(decimal_places=decimal_places) + + +class DecimalWholeDigitsError(PydanticValueError): + code = 'decimal.whole_digits' + msg_template = 'ensure that there are no more than {whole_digits} digits before the decimal point' + + def __init__(self, *, whole_digits: int) -> None: + super().__init__(whole_digits=whole_digits) + + +class DateTimeError(PydanticValueError): + msg_template = 'invalid datetime format' + + +class DateError(PydanticValueError): + msg_template = 'invalid date format' + + +class DateNotInThePastError(PydanticValueError): + code = 'date.not_in_the_past' + msg_template = 'date is not in the past' + + +class DateNotInTheFutureError(PydanticValueError): + code = 'date.not_in_the_future' + msg_template = 'date is not in the future' + + +class TimeError(PydanticValueError): + msg_template = 'invalid time format' + + +class DurationError(PydanticValueError): + msg_template = 'invalid duration format' + + +class HashableError(PydanticTypeError): + msg_template = 'value is not a valid hashable' + + +class UUIDError(PydanticTypeError): + msg_template = 'value is not a valid uuid' + + +class UUIDVersionError(PydanticValueError): + code = 'uuid.version' + msg_template = 'uuid version {required_version} expected' + + def __init__(self, *, required_version: int) -> None: + super().__init__(required_version=required_version) + + +class ArbitraryTypeError(PydanticTypeError): + code = 'arbitrary_type' + msg_template = 'instance of {expected_arbitrary_type} expected' + + def __init__(self, *, expected_arbitrary_type: Type[Any]) -> None: + super().__init__(expected_arbitrary_type=display_as_type(expected_arbitrary_type)) + + +class ClassError(PydanticTypeError): + code = 'class' + msg_template = 'a class is expected' + + +class SubclassError(PydanticTypeError): + code = 'subclass' + msg_template = 'subclass of {expected_class} expected' + + def __init__(self, *, expected_class: Type[Any]) -> None: + super().__init__(expected_class=display_as_type(expected_class)) + + +class JsonError(PydanticValueError): + msg_template = 'Invalid JSON' + + +class JsonTypeError(PydanticTypeError): + code = 'json' + msg_template = 'JSON object must be str, bytes or bytearray' + + +class PatternError(PydanticValueError): + code = 'regex_pattern' + msg_template = 'Invalid regular expression' + + +class DataclassTypeError(PydanticTypeError): + code = 'dataclass' + msg_template = 'instance of {class_name}, tuple or dict expected' + + +class CallableError(PydanticTypeError): + msg_template = '{value} is not callable' + + +class EnumError(PydanticTypeError): + code = 'enum_instance' + msg_template = '{value} is not a valid Enum instance' + + +class IntEnumError(PydanticTypeError): + code = 'int_enum_instance' + msg_template = '{value} is not a valid IntEnum instance' + + +class IPvAnyAddressError(PydanticValueError): + msg_template = 'value is not a valid IPv4 or IPv6 address' + + +class IPvAnyInterfaceError(PydanticValueError): + msg_template = 'value is not a valid IPv4 or IPv6 interface' + + +class IPvAnyNetworkError(PydanticValueError): + msg_template = 'value is not a valid IPv4 or IPv6 network' + + +class IPv4AddressError(PydanticValueError): + msg_template = 'value is not a valid IPv4 address' + + +class IPv6AddressError(PydanticValueError): + msg_template = 'value is not a valid IPv6 address' + + +class IPv4NetworkError(PydanticValueError): + msg_template = 'value is not a valid IPv4 network' + + +class IPv6NetworkError(PydanticValueError): + msg_template = 'value is not a valid IPv6 network' + + +class IPv4InterfaceError(PydanticValueError): + msg_template = 'value is not a valid IPv4 interface' + + +class IPv6InterfaceError(PydanticValueError): + msg_template = 'value is not a valid IPv6 interface' + + +class ColorError(PydanticValueError): + msg_template = 'value is not a valid color: {reason}' + + +class StrictBoolError(PydanticValueError): + msg_template = 'value is not a valid boolean' + + +class NotDigitError(PydanticValueError): + code = 'payment_card_number.digits' + msg_template = 'card number is not all digits' + + +class LuhnValidationError(PydanticValueError): + code = 'payment_card_number.luhn_check' + msg_template = 'card number is not luhn valid' + + +class InvalidLengthForBrand(PydanticValueError): + code = 'payment_card_number.invalid_length_for_brand' + msg_template = 'Length for a {brand} card must be {required_length}' + + +class InvalidByteSize(PydanticValueError): + msg_template = 'could not parse value and unit from byte string' + + +class InvalidByteSizeUnit(PydanticValueError): + msg_template = 'could not interpret byte unit: {unit}' + + +class MissingDiscriminator(PydanticValueError): + code = 'discriminated_union.missing_discriminator' + msg_template = 'Discriminator {discriminator_key!r} is missing in value' + + +class InvalidDiscriminator(PydanticValueError): + code = 'discriminated_union.invalid_discriminator' + msg_template = ( + 'No match for discriminator {discriminator_key!r} and value {discriminator_value!r} ' + '(allowed values: {allowed_values})' + ) + + def __init__(self, *, discriminator_key: str, discriminator_value: Any, allowed_values: Sequence[Any]) -> None: + super().__init__( + discriminator_key=discriminator_key, + discriminator_value=discriminator_value, + allowed_values=', '.join(map(repr, allowed_values)), + ) diff --git a/lib/python3.12/site-packages/pydantic/v1/fields.py b/lib/python3.12/site-packages/pydantic/v1/fields.py new file mode 100644 index 0000000000000000000000000000000000000000..002b60cde1bf28bbbf20af824638503d598fbc60 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/fields.py @@ -0,0 +1,1253 @@ +import copy +import re +from collections import Counter as CollectionCounter, defaultdict, deque +from collections.abc import Callable, Hashable as CollectionsHashable, Iterable as CollectionsIterable +from typing import ( + TYPE_CHECKING, + Any, + Counter, + DefaultDict, + Deque, + Dict, + ForwardRef, + FrozenSet, + Generator, + Iterable, + Iterator, + List, + Mapping, + Optional, + Pattern, + Sequence, + Set, + Tuple, + Type, + TypeVar, + Union, +) + +from typing_extensions import Annotated, Final + +from pydantic.v1 import errors as errors_ +from pydantic.v1.class_validators import Validator, make_generic_validator, prep_validators +from pydantic.v1.error_wrappers import ErrorWrapper +from pydantic.v1.errors import ConfigError, InvalidDiscriminator, MissingDiscriminator, NoneIsNotAllowedError +from pydantic.v1.types import Json, JsonWrapper +from pydantic.v1.typing import ( + NoArgAnyCallable, + convert_generics, + display_as_type, + get_args, + get_origin, + is_finalvar, + is_literal_type, + is_new_type, + is_none_type, + is_typeddict, + is_typeddict_special, + is_union, + new_type_supertype, +) +from pydantic.v1.utils import ( + PyObjectStr, + Representation, + ValueItems, + get_discriminator_alias_and_values, + get_unique_discriminator_alias, + lenient_isinstance, + lenient_issubclass, + sequence_like, + smart_deepcopy, +) +from pydantic.v1.validators import constant_validator, dict_validator, find_validators, validate_json + +Required: Any = Ellipsis + +T = TypeVar('T') + + +class UndefinedType: + def __repr__(self) -> str: + return 'PydanticUndefined' + + def __copy__(self: T) -> T: + return self + + def __reduce__(self) -> str: + return 'Undefined' + + def __deepcopy__(self: T, _: Any) -> T: + return self + + +Undefined = UndefinedType() + +if TYPE_CHECKING: + from pydantic.v1.class_validators import ValidatorsList + from pydantic.v1.config import BaseConfig + from pydantic.v1.error_wrappers import ErrorList + from pydantic.v1.types import ModelOrDc + from pydantic.v1.typing import AbstractSetIntStr, MappingIntStrAny, ReprArgs + + ValidateReturn = Tuple[Optional[Any], Optional[ErrorList]] + LocStr = Union[Tuple[Union[int, str], ...], str] + BoolUndefined = Union[bool, UndefinedType] + + +class FieldInfo(Representation): + """ + Captures extra information about a field. + """ + + __slots__ = ( + 'default', + 'default_factory', + 'alias', + 'alias_priority', + 'title', + 'description', + 'exclude', + 'include', + 'const', + 'gt', + 'ge', + 'lt', + 'le', + 'multiple_of', + 'allow_inf_nan', + 'max_digits', + 'decimal_places', + 'min_items', + 'max_items', + 'unique_items', + 'min_length', + 'max_length', + 'allow_mutation', + 'repr', + 'regex', + 'discriminator', + 'extra', + ) + + # field constraints with the default value, it's also used in update_from_config below + __field_constraints__ = { + 'min_length': None, + 'max_length': None, + 'regex': None, + 'gt': None, + 'lt': None, + 'ge': None, + 'le': None, + 'multiple_of': None, + 'allow_inf_nan': None, + 'max_digits': None, + 'decimal_places': None, + 'min_items': None, + 'max_items': None, + 'unique_items': None, + 'allow_mutation': True, + } + + def __init__(self, default: Any = Undefined, **kwargs: Any) -> None: + self.default = default + self.default_factory = kwargs.pop('default_factory', None) + self.alias = kwargs.pop('alias', None) + self.alias_priority = kwargs.pop('alias_priority', 2 if self.alias is not None else None) + self.title = kwargs.pop('title', None) + self.description = kwargs.pop('description', None) + self.exclude = kwargs.pop('exclude', None) + self.include = kwargs.pop('include', None) + self.const = kwargs.pop('const', None) + self.gt = kwargs.pop('gt', None) + self.ge = kwargs.pop('ge', None) + self.lt = kwargs.pop('lt', None) + self.le = kwargs.pop('le', None) + self.multiple_of = kwargs.pop('multiple_of', None) + self.allow_inf_nan = kwargs.pop('allow_inf_nan', None) + self.max_digits = kwargs.pop('max_digits', None) + self.decimal_places = kwargs.pop('decimal_places', None) + self.min_items = kwargs.pop('min_items', None) + self.max_items = kwargs.pop('max_items', None) + self.unique_items = kwargs.pop('unique_items', None) + self.min_length = kwargs.pop('min_length', None) + self.max_length = kwargs.pop('max_length', None) + self.allow_mutation = kwargs.pop('allow_mutation', True) + self.regex = kwargs.pop('regex', None) + self.discriminator = kwargs.pop('discriminator', None) + self.repr = kwargs.pop('repr', True) + self.extra = kwargs + + def __repr_args__(self) -> 'ReprArgs': + field_defaults_to_hide: Dict[str, Any] = { + 'repr': True, + **self.__field_constraints__, + } + + attrs = ((s, getattr(self, s)) for s in self.__slots__) + return [(a, v) for a, v in attrs if v != field_defaults_to_hide.get(a, None)] + + def get_constraints(self) -> Set[str]: + """ + Gets the constraints set on the field by comparing the constraint value with its default value + + :return: the constraints set on field_info + """ + return {attr for attr, default in self.__field_constraints__.items() if getattr(self, attr) != default} + + def update_from_config(self, from_config: Dict[str, Any]) -> None: + """ + Update this FieldInfo based on a dict from get_field_info, only fields which have not been set are dated. + """ + for attr_name, value in from_config.items(): + try: + current_value = getattr(self, attr_name) + except AttributeError: + # attr_name is not an attribute of FieldInfo, it should therefore be added to extra + # (except if extra already has this value!) + self.extra.setdefault(attr_name, value) + else: + if current_value is self.__field_constraints__.get(attr_name, None): + setattr(self, attr_name, value) + elif attr_name == 'exclude': + self.exclude = ValueItems.merge(value, current_value) + elif attr_name == 'include': + self.include = ValueItems.merge(value, current_value, intersect=True) + + def _validate(self) -> None: + if self.default is not Undefined and self.default_factory is not None: + raise ValueError('cannot specify both default and default_factory') + + +def Field( + default: Any = Undefined, + *, + default_factory: Optional[NoArgAnyCallable] = None, + alias: Optional[str] = None, + title: Optional[str] = None, + description: Optional[str] = None, + exclude: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny', Any]] = None, + include: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny', Any]] = None, + const: Optional[bool] = None, + gt: Optional[float] = None, + ge: Optional[float] = None, + lt: Optional[float] = None, + le: Optional[float] = None, + multiple_of: Optional[float] = None, + allow_inf_nan: Optional[bool] = None, + max_digits: Optional[int] = None, + decimal_places: Optional[int] = None, + min_items: Optional[int] = None, + max_items: Optional[int] = None, + unique_items: Optional[bool] = None, + min_length: Optional[int] = None, + max_length: Optional[int] = None, + allow_mutation: bool = True, + regex: Optional[str] = None, + discriminator: Optional[str] = None, + repr: bool = True, + **extra: Any, +) -> Any: + """ + Used to provide extra information about a field, either for the model schema or complex validation. Some arguments + apply only to number fields (``int``, ``float``, ``Decimal``) and some apply only to ``str``. + + :param default: since this is replacing the field’s default, its first argument is used + to set the default, use ellipsis (``...``) to indicate the field is required + :param default_factory: callable that will be called when a default value is needed for this field + If both `default` and `default_factory` are set, an error is raised. + :param alias: the public name of the field + :param title: can be any string, used in the schema + :param description: can be any string, used in the schema + :param exclude: exclude this field while dumping. + Takes same values as the ``include`` and ``exclude`` arguments on the ``.dict`` method. + :param include: include this field while dumping. + Takes same values as the ``include`` and ``exclude`` arguments on the ``.dict`` method. + :param const: this field is required and *must* take it's default value + :param gt: only applies to numbers, requires the field to be "greater than". The schema + will have an ``exclusiveMinimum`` validation keyword + :param ge: only applies to numbers, requires the field to be "greater than or equal to". The + schema will have a ``minimum`` validation keyword + :param lt: only applies to numbers, requires the field to be "less than". The schema + will have an ``exclusiveMaximum`` validation keyword + :param le: only applies to numbers, requires the field to be "less than or equal to". The + schema will have a ``maximum`` validation keyword + :param multiple_of: only applies to numbers, requires the field to be "a multiple of". The + schema will have a ``multipleOf`` validation keyword + :param allow_inf_nan: only applies to numbers, allows the field to be NaN or infinity (+inf or -inf), + which is a valid Python float. Default True, set to False for compatibility with JSON. + :param max_digits: only applies to Decimals, requires the field to have a maximum number + of digits within the decimal. It does not include a zero before the decimal point or trailing decimal zeroes. + :param decimal_places: only applies to Decimals, requires the field to have at most a number of decimal places + allowed. It does not include trailing decimal zeroes. + :param min_items: only applies to lists, requires the field to have a minimum number of + elements. The schema will have a ``minItems`` validation keyword + :param max_items: only applies to lists, requires the field to have a maximum number of + elements. The schema will have a ``maxItems`` validation keyword + :param unique_items: only applies to lists, requires the field not to have duplicated + elements. The schema will have a ``uniqueItems`` validation keyword + :param min_length: only applies to strings, requires the field to have a minimum length. The + schema will have a ``minLength`` validation keyword + :param max_length: only applies to strings, requires the field to have a maximum length. The + schema will have a ``maxLength`` validation keyword + :param allow_mutation: a boolean which defaults to True. When False, the field raises a TypeError if the field is + assigned on an instance. The BaseModel Config must set validate_assignment to True + :param regex: only applies to strings, requires the field match against a regular expression + pattern string. The schema will have a ``pattern`` validation keyword + :param discriminator: only useful with a (discriminated a.k.a. tagged) `Union` of sub models with a common field. + The `discriminator` is the name of this common field to shorten validation and improve generated schema + :param repr: show this field in the representation + :param **extra: any additional keyword arguments will be added as is to the schema + """ + field_info = FieldInfo( + default, + default_factory=default_factory, + alias=alias, + title=title, + description=description, + exclude=exclude, + include=include, + const=const, + gt=gt, + ge=ge, + lt=lt, + le=le, + multiple_of=multiple_of, + allow_inf_nan=allow_inf_nan, + max_digits=max_digits, + decimal_places=decimal_places, + min_items=min_items, + max_items=max_items, + unique_items=unique_items, + min_length=min_length, + max_length=max_length, + allow_mutation=allow_mutation, + regex=regex, + discriminator=discriminator, + repr=repr, + **extra, + ) + field_info._validate() + return field_info + + +# used to be an enum but changed to int's for small performance improvement as less access overhead +SHAPE_SINGLETON = 1 +SHAPE_LIST = 2 +SHAPE_SET = 3 +SHAPE_MAPPING = 4 +SHAPE_TUPLE = 5 +SHAPE_TUPLE_ELLIPSIS = 6 +SHAPE_SEQUENCE = 7 +SHAPE_FROZENSET = 8 +SHAPE_ITERABLE = 9 +SHAPE_GENERIC = 10 +SHAPE_DEQUE = 11 +SHAPE_DICT = 12 +SHAPE_DEFAULTDICT = 13 +SHAPE_COUNTER = 14 +SHAPE_NAME_LOOKUP = { + SHAPE_LIST: 'List[{}]', + SHAPE_SET: 'Set[{}]', + SHAPE_TUPLE_ELLIPSIS: 'Tuple[{}, ...]', + SHAPE_SEQUENCE: 'Sequence[{}]', + SHAPE_FROZENSET: 'FrozenSet[{}]', + SHAPE_ITERABLE: 'Iterable[{}]', + SHAPE_DEQUE: 'Deque[{}]', + SHAPE_DICT: 'Dict[{}]', + SHAPE_DEFAULTDICT: 'DefaultDict[{}]', + SHAPE_COUNTER: 'Counter[{}]', +} + +MAPPING_LIKE_SHAPES: Set[int] = {SHAPE_DEFAULTDICT, SHAPE_DICT, SHAPE_MAPPING, SHAPE_COUNTER} + + +class ModelField(Representation): + __slots__ = ( + 'type_', + 'outer_type_', + 'annotation', + 'sub_fields', + 'sub_fields_mapping', + 'key_field', + 'validators', + 'pre_validators', + 'post_validators', + 'default', + 'default_factory', + 'required', + 'final', + 'model_config', + 'name', + 'alias', + 'has_alias', + 'field_info', + 'discriminator_key', + 'discriminator_alias', + 'validate_always', + 'allow_none', + 'shape', + 'class_validators', + 'parse_json', + ) + + def __init__( + self, + *, + name: str, + type_: Type[Any], + class_validators: Optional[Dict[str, Validator]], + model_config: Type['BaseConfig'], + default: Any = None, + default_factory: Optional[NoArgAnyCallable] = None, + required: 'BoolUndefined' = Undefined, + final: bool = False, + alias: Optional[str] = None, + field_info: Optional[FieldInfo] = None, + ) -> None: + self.name: str = name + self.has_alias: bool = alias is not None + self.alias: str = alias if alias is not None else name + self.annotation = type_ + self.type_: Any = convert_generics(type_) + self.outer_type_: Any = type_ + self.class_validators = class_validators or {} + self.default: Any = default + self.default_factory: Optional[NoArgAnyCallable] = default_factory + self.required: 'BoolUndefined' = required + self.final: bool = final + self.model_config = model_config + self.field_info: FieldInfo = field_info or FieldInfo(default) + self.discriminator_key: Optional[str] = self.field_info.discriminator + self.discriminator_alias: Optional[str] = self.discriminator_key + + self.allow_none: bool = False + self.validate_always: bool = False + self.sub_fields: Optional[List[ModelField]] = None + self.sub_fields_mapping: Optional[Dict[str, 'ModelField']] = None # used for discriminated union + self.key_field: Optional[ModelField] = None + self.validators: 'ValidatorsList' = [] + self.pre_validators: Optional['ValidatorsList'] = None + self.post_validators: Optional['ValidatorsList'] = None + self.parse_json: bool = False + self.shape: int = SHAPE_SINGLETON + self.model_config.prepare_field(self) + self.prepare() + + def get_default(self) -> Any: + return smart_deepcopy(self.default) if self.default_factory is None else self.default_factory() + + @staticmethod + def _get_field_info( + field_name: str, annotation: Any, value: Any, config: Type['BaseConfig'] + ) -> Tuple[FieldInfo, Any]: + """ + Get a FieldInfo from a root typing.Annotated annotation, value, or config default. + + The FieldInfo may be set in typing.Annotated or the value, but not both. If neither contain + a FieldInfo, a new one will be created using the config. + + :param field_name: name of the field for use in error messages + :param annotation: a type hint such as `str` or `Annotated[str, Field(..., min_length=5)]` + :param value: the field's assigned value + :param config: the model's config object + :return: the FieldInfo contained in the `annotation`, the value, or a new one from the config. + """ + field_info_from_config = config.get_field_info(field_name) + + field_info = None + if get_origin(annotation) is Annotated: + field_infos = [arg for arg in get_args(annotation)[1:] if isinstance(arg, FieldInfo)] + if len(field_infos) > 1: + raise ValueError(f'cannot specify multiple `Annotated` `Field`s for {field_name!r}') + field_info = next(iter(field_infos), None) + if field_info is not None: + field_info = copy.copy(field_info) + field_info.update_from_config(field_info_from_config) + if field_info.default not in (Undefined, Required): + raise ValueError(f'`Field` default cannot be set in `Annotated` for {field_name!r}') + if value is not Undefined and value is not Required: + # check also `Required` because of `validate_arguments` that sets `...` as default value + field_info.default = value + + if isinstance(value, FieldInfo): + if field_info is not None: + raise ValueError(f'cannot specify `Annotated` and value `Field`s together for {field_name!r}') + field_info = value + field_info.update_from_config(field_info_from_config) + elif field_info is None: + field_info = FieldInfo(value, **field_info_from_config) + value = None if field_info.default_factory is not None else field_info.default + field_info._validate() + return field_info, value + + @classmethod + def infer( + cls, + *, + name: str, + value: Any, + annotation: Any, + class_validators: Optional[Dict[str, Validator]], + config: Type['BaseConfig'], + ) -> 'ModelField': + from pydantic.v1.schema import get_annotation_from_field_info + + field_info, value = cls._get_field_info(name, annotation, value, config) + required: 'BoolUndefined' = Undefined + if value is Required: + required = True + value = None + elif value is not Undefined: + required = False + annotation = get_annotation_from_field_info(annotation, field_info, name, config.validate_assignment) + + return cls( + name=name, + type_=annotation, + alias=field_info.alias, + class_validators=class_validators, + default=value, + default_factory=field_info.default_factory, + required=required, + model_config=config, + field_info=field_info, + ) + + def set_config(self, config: Type['BaseConfig']) -> None: + self.model_config = config + info_from_config = config.get_field_info(self.name) + config.prepare_field(self) + new_alias = info_from_config.get('alias') + new_alias_priority = info_from_config.get('alias_priority') or 0 + if new_alias and new_alias_priority >= (self.field_info.alias_priority or 0): + self.field_info.alias = new_alias + self.field_info.alias_priority = new_alias_priority + self.alias = new_alias + new_exclude = info_from_config.get('exclude') + if new_exclude is not None: + self.field_info.exclude = ValueItems.merge(self.field_info.exclude, new_exclude) + new_include = info_from_config.get('include') + if new_include is not None: + self.field_info.include = ValueItems.merge(self.field_info.include, new_include, intersect=True) + + @property + def alt_alias(self) -> bool: + return self.name != self.alias + + def prepare(self) -> None: + """ + Prepare the field but inspecting self.default, self.type_ etc. + + Note: this method is **not** idempotent (because _type_analysis is not idempotent), + e.g. calling it it multiple times may modify the field and configure it incorrectly. + """ + self._set_default_and_type() + if self.type_.__class__ is ForwardRef or self.type_.__class__ is DeferredType: + # self.type_ is currently a ForwardRef and there's nothing we can do now, + # user will need to call model.update_forward_refs() + return + + self._type_analysis() + if self.required is Undefined: + self.required = True + if self.default is Undefined and self.default_factory is None: + self.default = None + self.populate_validators() + + def _set_default_and_type(self) -> None: + """ + Set the default value, infer the type if needed and check if `None` value is valid. + """ + if self.default_factory is not None: + if self.type_ is Undefined: + raise errors_.ConfigError( + f'you need to set the type of field {self.name!r} when using `default_factory`' + ) + return + + default_value = self.get_default() + + if default_value is not None and self.type_ is Undefined: + self.type_ = default_value.__class__ + self.outer_type_ = self.type_ + self.annotation = self.type_ + + if self.type_ is Undefined: + raise errors_.ConfigError(f'unable to infer type for attribute "{self.name}"') + + if self.required is False and default_value is None: + self.allow_none = True + + def _type_analysis(self) -> None: # noqa: C901 (ignore complexity) + # typing interface is horrible, we have to do some ugly checks + if lenient_issubclass(self.type_, JsonWrapper): + self.type_ = self.type_.inner_type + self.parse_json = True + elif lenient_issubclass(self.type_, Json): + self.type_ = Any + self.parse_json = True + elif isinstance(self.type_, TypeVar): + if self.type_.__bound__: + self.type_ = self.type_.__bound__ + elif self.type_.__constraints__: + self.type_ = Union[self.type_.__constraints__] + else: + self.type_ = Any + elif is_new_type(self.type_): + self.type_ = new_type_supertype(self.type_) + + if self.type_ is Any or self.type_ is object: + if self.required is Undefined: + self.required = False + self.allow_none = True + return + elif self.type_ is Pattern or self.type_ is re.Pattern: + # python 3.7 only, Pattern is a typing object but without sub fields + return + elif is_literal_type(self.type_): + return + elif is_typeddict(self.type_): + return + + if is_finalvar(self.type_): + self.final = True + + if self.type_ is Final: + self.type_ = Any + else: + self.type_ = get_args(self.type_)[0] + + self._type_analysis() + return + + origin = get_origin(self.type_) + + if origin is Annotated or is_typeddict_special(origin): + self.type_ = get_args(self.type_)[0] + self._type_analysis() + return + + if self.discriminator_key is not None and not is_union(origin): + raise TypeError('`discriminator` can only be used with `Union` type with more than one variant') + + # add extra check for `collections.abc.Hashable` for python 3.10+ where origin is not `None` + if origin is None or origin is CollectionsHashable: + # field is not "typing" object eg. Union, Dict, List etc. + # allow None for virtual superclasses of NoneType, e.g. Hashable + if isinstance(self.type_, type) and isinstance(None, self.type_): + self.allow_none = True + return + elif origin is Callable: + return + elif is_union(origin): + types_ = [] + for type_ in get_args(self.type_): + if is_none_type(type_) or type_ is Any or type_ is object: + if self.required is Undefined: + self.required = False + self.allow_none = True + if is_none_type(type_): + continue + types_.append(type_) + + if len(types_) == 1: + # Optional[] + self.type_ = types_[0] + # this is the one case where the "outer type" isn't just the original type + self.outer_type_ = self.type_ + # re-run to correctly interpret the new self.type_ + self._type_analysis() + else: + self.sub_fields = [self._create_sub_type(t, f'{self.name}_{display_as_type(t)}') for t in types_] + + if self.discriminator_key is not None: + self.prepare_discriminated_union_sub_fields() + return + elif issubclass(origin, Tuple): # type: ignore + # origin == Tuple without item type + args = get_args(self.type_) + if not args: # plain tuple + self.type_ = Any + self.shape = SHAPE_TUPLE_ELLIPSIS + elif len(args) == 2 and args[1] is Ellipsis: # e.g. Tuple[int, ...] + self.type_ = args[0] + self.shape = SHAPE_TUPLE_ELLIPSIS + self.sub_fields = [self._create_sub_type(args[0], f'{self.name}_0')] + elif args == ((),): # Tuple[()] means empty tuple + self.shape = SHAPE_TUPLE + self.type_ = Any + self.sub_fields = [] + else: + self.shape = SHAPE_TUPLE + self.sub_fields = [self._create_sub_type(t, f'{self.name}_{i}') for i, t in enumerate(args)] + return + elif issubclass(origin, List): + # Create self validators + get_validators = getattr(self.type_, '__get_validators__', None) + if get_validators: + self.class_validators.update( + {f'list_{i}': Validator(validator, pre=True) for i, validator in enumerate(get_validators())} + ) + + self.type_ = get_args(self.type_)[0] + self.shape = SHAPE_LIST + elif issubclass(origin, Set): + # Create self validators + get_validators = getattr(self.type_, '__get_validators__', None) + if get_validators: + self.class_validators.update( + {f'set_{i}': Validator(validator, pre=True) for i, validator in enumerate(get_validators())} + ) + + self.type_ = get_args(self.type_)[0] + self.shape = SHAPE_SET + elif issubclass(origin, FrozenSet): + # Create self validators + get_validators = getattr(self.type_, '__get_validators__', None) + if get_validators: + self.class_validators.update( + {f'frozenset_{i}': Validator(validator, pre=True) for i, validator in enumerate(get_validators())} + ) + + self.type_ = get_args(self.type_)[0] + self.shape = SHAPE_FROZENSET + elif issubclass(origin, Deque): + self.type_ = get_args(self.type_)[0] + self.shape = SHAPE_DEQUE + elif issubclass(origin, Sequence): + self.type_ = get_args(self.type_)[0] + self.shape = SHAPE_SEQUENCE + # priority to most common mapping: dict + elif origin is dict or origin is Dict: + self.key_field = self._create_sub_type(get_args(self.type_)[0], 'key_' + self.name, for_keys=True) + self.type_ = get_args(self.type_)[1] + self.shape = SHAPE_DICT + elif issubclass(origin, DefaultDict): + self.key_field = self._create_sub_type(get_args(self.type_)[0], 'key_' + self.name, for_keys=True) + self.type_ = get_args(self.type_)[1] + self.shape = SHAPE_DEFAULTDICT + elif issubclass(origin, Counter): + self.key_field = self._create_sub_type(get_args(self.type_)[0], 'key_' + self.name, for_keys=True) + self.type_ = int + self.shape = SHAPE_COUNTER + elif issubclass(origin, Mapping): + self.key_field = self._create_sub_type(get_args(self.type_)[0], 'key_' + self.name, for_keys=True) + self.type_ = get_args(self.type_)[1] + self.shape = SHAPE_MAPPING + # Equality check as almost everything inherits form Iterable, including str + # check for Iterable and CollectionsIterable, as it could receive one even when declared with the other + elif origin in {Iterable, CollectionsIterable}: + self.type_ = get_args(self.type_)[0] + self.shape = SHAPE_ITERABLE + self.sub_fields = [self._create_sub_type(self.type_, f'{self.name}_type')] + elif issubclass(origin, Type): # type: ignore + return + elif hasattr(origin, '__get_validators__') or self.model_config.arbitrary_types_allowed: + # Is a Pydantic-compatible generic that handles itself + # or we have arbitrary_types_allowed = True + self.shape = SHAPE_GENERIC + self.sub_fields = [self._create_sub_type(t, f'{self.name}_{i}') for i, t in enumerate(get_args(self.type_))] + self.type_ = origin + return + else: + raise TypeError(f'Fields of type "{origin}" are not supported.') + + # type_ has been refined eg. as the type of a List and sub_fields needs to be populated + self.sub_fields = [self._create_sub_type(self.type_, '_' + self.name)] + + def prepare_discriminated_union_sub_fields(self) -> None: + """ + Prepare the mapping -> and update `sub_fields` + Note that this process can be aborted if a `ForwardRef` is encountered + """ + assert self.discriminator_key is not None + + if self.type_.__class__ is DeferredType: + return + + assert self.sub_fields is not None + sub_fields_mapping: Dict[str, 'ModelField'] = {} + all_aliases: Set[str] = set() + + for sub_field in self.sub_fields: + t = sub_field.type_ + if t.__class__ is ForwardRef: + # Stopping everything...will need to call `update_forward_refs` + return + + alias, discriminator_values = get_discriminator_alias_and_values(t, self.discriminator_key) + all_aliases.add(alias) + for discriminator_value in discriminator_values: + sub_fields_mapping[discriminator_value] = sub_field + + self.sub_fields_mapping = sub_fields_mapping + self.discriminator_alias = get_unique_discriminator_alias(all_aliases, self.discriminator_key) + + def _create_sub_type(self, type_: Type[Any], name: str, *, for_keys: bool = False) -> 'ModelField': + if for_keys: + class_validators = None + else: + # validators for sub items should not have `each_item` as we want to check only the first sublevel + class_validators = { + k: Validator( + func=v.func, + pre=v.pre, + each_item=False, + always=v.always, + check_fields=v.check_fields, + skip_on_failure=v.skip_on_failure, + ) + for k, v in self.class_validators.items() + if v.each_item + } + + field_info, _ = self._get_field_info(name, type_, None, self.model_config) + + return self.__class__( + type_=type_, + name=name, + class_validators=class_validators, + model_config=self.model_config, + field_info=field_info, + ) + + def populate_validators(self) -> None: + """ + Prepare self.pre_validators, self.validators, and self.post_validators based on self.type_'s __get_validators__ + and class validators. This method should be idempotent, e.g. it should be safe to call multiple times + without mis-configuring the field. + """ + self.validate_always = getattr(self.type_, 'validate_always', False) or any( + v.always for v in self.class_validators.values() + ) + + class_validators_ = self.class_validators.values() + if not self.sub_fields or self.shape == SHAPE_GENERIC: + get_validators = getattr(self.type_, '__get_validators__', None) + v_funcs = ( + *[v.func for v in class_validators_ if v.each_item and v.pre], + *(get_validators() if get_validators else list(find_validators(self.type_, self.model_config))), + *[v.func for v in class_validators_ if v.each_item and not v.pre], + ) + self.validators = prep_validators(v_funcs) + + self.pre_validators = [] + self.post_validators = [] + + if self.field_info and self.field_info.const: + self.post_validators.append(make_generic_validator(constant_validator)) + + if class_validators_: + self.pre_validators += prep_validators(v.func for v in class_validators_ if not v.each_item and v.pre) + self.post_validators += prep_validators(v.func for v in class_validators_ if not v.each_item and not v.pre) + + if self.parse_json: + self.pre_validators.append(make_generic_validator(validate_json)) + + self.pre_validators = self.pre_validators or None + self.post_validators = self.post_validators or None + + def validate( + self, v: Any, values: Dict[str, Any], *, loc: 'LocStr', cls: Optional['ModelOrDc'] = None + ) -> 'ValidateReturn': + assert self.type_.__class__ is not DeferredType + + if self.type_.__class__ is ForwardRef: + assert cls is not None + raise ConfigError( + f'field "{self.name}" not yet prepared so type is still a ForwardRef, ' + f'you might need to call {cls.__name__}.update_forward_refs().' + ) + + errors: Optional['ErrorList'] + if self.pre_validators: + v, errors = self._apply_validators(v, values, loc, cls, self.pre_validators) + if errors: + return v, errors + + if v is None: + if is_none_type(self.type_): + # keep validating + pass + elif self.allow_none: + if self.post_validators: + return self._apply_validators(v, values, loc, cls, self.post_validators) + else: + return None, None + else: + return v, ErrorWrapper(NoneIsNotAllowedError(), loc) + + if self.shape == SHAPE_SINGLETON: + v, errors = self._validate_singleton(v, values, loc, cls) + elif self.shape in MAPPING_LIKE_SHAPES: + v, errors = self._validate_mapping_like(v, values, loc, cls) + elif self.shape == SHAPE_TUPLE: + v, errors = self._validate_tuple(v, values, loc, cls) + elif self.shape == SHAPE_ITERABLE: + v, errors = self._validate_iterable(v, values, loc, cls) + elif self.shape == SHAPE_GENERIC: + v, errors = self._apply_validators(v, values, loc, cls, self.validators) + else: + # sequence, list, set, generator, tuple with ellipsis, frozen set + v, errors = self._validate_sequence_like(v, values, loc, cls) + + if not errors and self.post_validators: + v, errors = self._apply_validators(v, values, loc, cls, self.post_validators) + return v, errors + + def _validate_sequence_like( # noqa: C901 (ignore complexity) + self, v: Any, values: Dict[str, Any], loc: 'LocStr', cls: Optional['ModelOrDc'] + ) -> 'ValidateReturn': + """ + Validate sequence-like containers: lists, tuples, sets and generators + Note that large if-else blocks are necessary to enable Cython + optimization, which is why we disable the complexity check above. + """ + if not sequence_like(v): + e: errors_.PydanticTypeError + if self.shape == SHAPE_LIST: + e = errors_.ListError() + elif self.shape in (SHAPE_TUPLE, SHAPE_TUPLE_ELLIPSIS): + e = errors_.TupleError() + elif self.shape == SHAPE_SET: + e = errors_.SetError() + elif self.shape == SHAPE_FROZENSET: + e = errors_.FrozenSetError() + else: + e = errors_.SequenceError() + return v, ErrorWrapper(e, loc) + + loc = loc if isinstance(loc, tuple) else (loc,) + result = [] + errors: List[ErrorList] = [] + for i, v_ in enumerate(v): + v_loc = *loc, i + r, ee = self._validate_singleton(v_, values, v_loc, cls) + if ee: + errors.append(ee) + else: + result.append(r) + + if errors: + return v, errors + + converted: Union[List[Any], Set[Any], FrozenSet[Any], Tuple[Any, ...], Iterator[Any], Deque[Any]] = result + + if self.shape == SHAPE_SET: + converted = set(result) + elif self.shape == SHAPE_FROZENSET: + converted = frozenset(result) + elif self.shape == SHAPE_TUPLE_ELLIPSIS: + converted = tuple(result) + elif self.shape == SHAPE_DEQUE: + converted = deque(result, maxlen=getattr(v, 'maxlen', None)) + elif self.shape == SHAPE_SEQUENCE: + if isinstance(v, tuple): + converted = tuple(result) + elif isinstance(v, set): + converted = set(result) + elif isinstance(v, Generator): + converted = iter(result) + elif isinstance(v, deque): + converted = deque(result, maxlen=getattr(v, 'maxlen', None)) + return converted, None + + def _validate_iterable( + self, v: Any, values: Dict[str, Any], loc: 'LocStr', cls: Optional['ModelOrDc'] + ) -> 'ValidateReturn': + """ + Validate Iterables. + + This intentionally doesn't validate values to allow infinite generators. + """ + + try: + iterable = iter(v) + except TypeError: + return v, ErrorWrapper(errors_.IterableError(), loc) + return iterable, None + + def _validate_tuple( + self, v: Any, values: Dict[str, Any], loc: 'LocStr', cls: Optional['ModelOrDc'] + ) -> 'ValidateReturn': + e: Optional[Exception] = None + if not sequence_like(v): + e = errors_.TupleError() + else: + actual_length, expected_length = len(v), len(self.sub_fields) # type: ignore + if actual_length != expected_length: + e = errors_.TupleLengthError(actual_length=actual_length, expected_length=expected_length) + + if e: + return v, ErrorWrapper(e, loc) + + loc = loc if isinstance(loc, tuple) else (loc,) + result = [] + errors: List[ErrorList] = [] + for i, (v_, field) in enumerate(zip(v, self.sub_fields)): # type: ignore + v_loc = *loc, i + r, ee = field.validate(v_, values, loc=v_loc, cls=cls) + if ee: + errors.append(ee) + else: + result.append(r) + + if errors: + return v, errors + else: + return tuple(result), None + + def _validate_mapping_like( + self, v: Any, values: Dict[str, Any], loc: 'LocStr', cls: Optional['ModelOrDc'] + ) -> 'ValidateReturn': + try: + v_iter = dict_validator(v) + except TypeError as exc: + return v, ErrorWrapper(exc, loc) + + loc = loc if isinstance(loc, tuple) else (loc,) + result, errors = {}, [] + for k, v_ in v_iter.items(): + v_loc = *loc, '__key__' + key_result, key_errors = self.key_field.validate(k, values, loc=v_loc, cls=cls) # type: ignore + if key_errors: + errors.append(key_errors) + continue + + v_loc = *loc, k + value_result, value_errors = self._validate_singleton(v_, values, v_loc, cls) + if value_errors: + errors.append(value_errors) + continue + + result[key_result] = value_result + if errors: + return v, errors + elif self.shape == SHAPE_DICT: + return result, None + elif self.shape == SHAPE_DEFAULTDICT: + return defaultdict(self.type_, result), None + elif self.shape == SHAPE_COUNTER: + return CollectionCounter(result), None + else: + return self._get_mapping_value(v, result), None + + def _get_mapping_value(self, original: T, converted: Dict[Any, Any]) -> Union[T, Dict[Any, Any]]: + """ + When type is `Mapping[KT, KV]` (or another unsupported mapping), we try to avoid + coercing to `dict` unwillingly. + """ + original_cls = original.__class__ + + if original_cls == dict or original_cls == Dict: + return converted + elif original_cls in {defaultdict, DefaultDict}: + return defaultdict(self.type_, converted) + else: + try: + # Counter, OrderedDict, UserDict, ... + return original_cls(converted) # type: ignore + except TypeError: + raise RuntimeError(f'Could not convert dictionary to {original_cls.__name__!r}') from None + + def _validate_singleton( + self, v: Any, values: Dict[str, Any], loc: 'LocStr', cls: Optional['ModelOrDc'] + ) -> 'ValidateReturn': + if self.sub_fields: + if self.discriminator_key is not None: + return self._validate_discriminated_union(v, values, loc, cls) + + errors = [] + + if self.model_config.smart_union and is_union(get_origin(self.type_)): + # 1st pass: check if the value is an exact instance of one of the Union types + # (e.g. to avoid coercing a bool into an int) + for field in self.sub_fields: + if v.__class__ is field.outer_type_: + return v, None + + # 2nd pass: check if the value is an instance of any subclass of the Union types + for field in self.sub_fields: + # This whole logic will be improved later on to support more complex `isinstance` checks + # It will probably be done once a strict mode is added and be something like: + # ``` + # value, error = field.validate(v, values, strict=True) + # if error is None: + # return value, None + # ``` + try: + if isinstance(v, field.outer_type_): + return v, None + except TypeError: + # compound type + if lenient_isinstance(v, get_origin(field.outer_type_)): + value, error = field.validate(v, values, loc=loc, cls=cls) + if not error: + return value, None + + # 1st pass by default or 3rd pass with `smart_union` enabled: + # check if the value can be coerced into one of the Union types + for field in self.sub_fields: + value, error = field.validate(v, values, loc=loc, cls=cls) + if error: + errors.append(error) + else: + return value, None + return v, errors + else: + return self._apply_validators(v, values, loc, cls, self.validators) + + def _validate_discriminated_union( + self, v: Any, values: Dict[str, Any], loc: 'LocStr', cls: Optional['ModelOrDc'] + ) -> 'ValidateReturn': + assert self.discriminator_key is not None + assert self.discriminator_alias is not None + + try: + try: + discriminator_value = v[self.discriminator_alias] + except KeyError: + if self.model_config.allow_population_by_field_name: + discriminator_value = v[self.discriminator_key] + else: + raise + except KeyError: + return v, ErrorWrapper(MissingDiscriminator(discriminator_key=self.discriminator_key), loc) + except TypeError: + try: + # BaseModel or dataclass + discriminator_value = getattr(v, self.discriminator_key) + except (AttributeError, TypeError): + return v, ErrorWrapper(MissingDiscriminator(discriminator_key=self.discriminator_key), loc) + + if self.sub_fields_mapping is None: + assert cls is not None + raise ConfigError( + f'field "{self.name}" not yet prepared so type is still a ForwardRef, ' + f'you might need to call {cls.__name__}.update_forward_refs().' + ) + + try: + sub_field = self.sub_fields_mapping[discriminator_value] + except (KeyError, TypeError): + # KeyError: `discriminator_value` is not in the dictionary. + # TypeError: `discriminator_value` is unhashable. + assert self.sub_fields_mapping is not None + return v, ErrorWrapper( + InvalidDiscriminator( + discriminator_key=self.discriminator_key, + discriminator_value=discriminator_value, + allowed_values=list(self.sub_fields_mapping), + ), + loc, + ) + else: + if not isinstance(loc, tuple): + loc = (loc,) + return sub_field.validate(v, values, loc=(*loc, display_as_type(sub_field.type_)), cls=cls) + + def _apply_validators( + self, v: Any, values: Dict[str, Any], loc: 'LocStr', cls: Optional['ModelOrDc'], validators: 'ValidatorsList' + ) -> 'ValidateReturn': + for validator in validators: + try: + v = validator(cls, v, values, self, self.model_config) + except (ValueError, TypeError, AssertionError) as exc: + return v, ErrorWrapper(exc, loc) + return v, None + + def is_complex(self) -> bool: + """ + Whether the field is "complex" eg. env variables should be parsed as JSON. + """ + from pydantic.v1.main import BaseModel + + return ( + self.shape != SHAPE_SINGLETON + or hasattr(self.type_, '__pydantic_model__') + or lenient_issubclass(self.type_, (BaseModel, list, set, frozenset, dict)) + ) + + def _type_display(self) -> PyObjectStr: + t = display_as_type(self.type_) + + if self.shape in MAPPING_LIKE_SHAPES: + t = f'Mapping[{display_as_type(self.key_field.type_)}, {t}]' # type: ignore + elif self.shape == SHAPE_TUPLE: + t = 'Tuple[{}]'.format(', '.join(display_as_type(f.type_) for f in self.sub_fields)) # type: ignore + elif self.shape == SHAPE_GENERIC: + assert self.sub_fields + t = '{}[{}]'.format( + display_as_type(self.type_), ', '.join(display_as_type(f.type_) for f in self.sub_fields) + ) + elif self.shape != SHAPE_SINGLETON: + t = SHAPE_NAME_LOOKUP[self.shape].format(t) + + if self.allow_none and (self.shape != SHAPE_SINGLETON or not self.sub_fields): + t = f'Optional[{t}]' + return PyObjectStr(t) + + def __repr_args__(self) -> 'ReprArgs': + args = [('name', self.name), ('type', self._type_display()), ('required', self.required)] + + if not self.required: + if self.default_factory is not None: + args.append(('default_factory', f'')) + else: + args.append(('default', self.default)) + + if self.alt_alias: + args.append(('alias', self.alias)) + return args + + +class ModelPrivateAttr(Representation): + __slots__ = ('default', 'default_factory') + + def __init__(self, default: Any = Undefined, *, default_factory: Optional[NoArgAnyCallable] = None) -> None: + self.default = default + self.default_factory = default_factory + + def get_default(self) -> Any: + return smart_deepcopy(self.default) if self.default_factory is None else self.default_factory() + + def __eq__(self, other: Any) -> bool: + return isinstance(other, self.__class__) and (self.default, self.default_factory) == ( + other.default, + other.default_factory, + ) + + +def PrivateAttr( + default: Any = Undefined, + *, + default_factory: Optional[NoArgAnyCallable] = None, +) -> Any: + """ + Indicates that attribute is only used internally and never mixed with regular fields. + + Types or values of private attrs are not checked by pydantic and it's up to you to keep them relevant. + + Private attrs are stored in model __slots__. + + :param default: the attribute’s default value + :param default_factory: callable that will be called when a default value is needed for this attribute + If both `default` and `default_factory` are set, an error is raised. + """ + if default is not Undefined and default_factory is not None: + raise ValueError('cannot specify both default and default_factory') + + return ModelPrivateAttr( + default, + default_factory=default_factory, + ) + + +class DeferredType: + """ + Used to postpone field preparation, while creating recursive generic models. + """ + + +def is_finalvar_with_default_val(type_: Type[Any], val: Any) -> bool: + return is_finalvar(type_) and val is not Undefined and not isinstance(val, FieldInfo) diff --git a/lib/python3.12/site-packages/pydantic/v1/generics.py b/lib/python3.12/site-packages/pydantic/v1/generics.py new file mode 100644 index 0000000000000000000000000000000000000000..9a69f2b39390cd1221f96b644b8d494597ce205c --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/generics.py @@ -0,0 +1,400 @@ +import sys +import types +import typing +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Dict, + ForwardRef, + Generic, + Iterator, + List, + Mapping, + Optional, + Tuple, + Type, + TypeVar, + Union, + cast, +) +from weakref import WeakKeyDictionary, WeakValueDictionary + +from typing_extensions import Annotated, Literal as ExtLiteral + +from pydantic.v1.class_validators import gather_all_validators +from pydantic.v1.fields import DeferredType +from pydantic.v1.main import BaseModel, create_model +from pydantic.v1.types import JsonWrapper +from pydantic.v1.typing import display_as_type, get_all_type_hints, get_args, get_origin, typing_base +from pydantic.v1.utils import all_identical, lenient_issubclass + +if sys.version_info >= (3, 10): + from typing import _UnionGenericAlias +if sys.version_info >= (3, 8): + from typing import Literal + +GenericModelT = TypeVar('GenericModelT', bound='GenericModel') +TypeVarType = Any # since mypy doesn't allow the use of TypeVar as a type + +CacheKey = Tuple[Type[Any], Any, Tuple[Any, ...]] +Parametrization = Mapping[TypeVarType, Type[Any]] + +# weak dictionaries allow the dynamically created parametrized versions of generic models to get collected +# once they are no longer referenced by the caller. +if sys.version_info >= (3, 9): # Typing for weak dictionaries available at 3.9 + GenericTypesCache = WeakValueDictionary[CacheKey, Type[BaseModel]] + AssignedParameters = WeakKeyDictionary[Type[BaseModel], Parametrization] +else: + GenericTypesCache = WeakValueDictionary + AssignedParameters = WeakKeyDictionary + +# _generic_types_cache is a Mapping from __class_getitem__ arguments to the parametrized version of generic models. +# This ensures multiple calls of e.g. A[B] return always the same class. +_generic_types_cache = GenericTypesCache() + +# _assigned_parameters is a Mapping from parametrized version of generic models to assigned types of parametrizations +# as captured during construction of the class (not instances). +# E.g., for generic model `Model[A, B]`, when parametrized model `Model[int, str]` is created, +# `Model[int, str]`: {A: int, B: str}` will be stored in `_assigned_parameters`. +# (This information is only otherwise available after creation from the class name string). +_assigned_parameters = AssignedParameters() + + +class GenericModel(BaseModel): + __slots__ = () + __concrete__: ClassVar[bool] = False + + if TYPE_CHECKING: + # Putting this in a TYPE_CHECKING block allows us to replace `if Generic not in cls.__bases__` with + # `not hasattr(cls, "__parameters__")`. This means we don't need to force non-concrete subclasses of + # `GenericModel` to also inherit from `Generic`, which would require changes to the use of `create_model` below. + __parameters__: ClassVar[Tuple[TypeVarType, ...]] + + # Setting the return type as Type[Any] instead of Type[BaseModel] prevents PyCharm warnings + def __class_getitem__(cls: Type[GenericModelT], params: Union[Type[Any], Tuple[Type[Any], ...]]) -> Type[Any]: + """Instantiates a new class from a generic class `cls` and type variables `params`. + + :param params: Tuple of types the class . Given a generic class + `Model` with 2 type variables and a concrete model `Model[str, int]`, + the value `(str, int)` would be passed to `params`. + :return: New model class inheriting from `cls` with instantiated + types described by `params`. If no parameters are given, `cls` is + returned as is. + + """ + + def _cache_key(_params: Any) -> CacheKey: + args = get_args(_params) + # python returns a list for Callables, which is not hashable + if len(args) == 2 and isinstance(args[0], list): + args = (tuple(args[0]), args[1]) + return cls, _params, args + + cached = _generic_types_cache.get(_cache_key(params)) + if cached is not None: + return cached + if cls.__concrete__ and Generic not in cls.__bases__: + raise TypeError('Cannot parameterize a concrete instantiation of a generic model') + if not isinstance(params, tuple): + params = (params,) + if cls is GenericModel and any(isinstance(param, TypeVar) for param in params): + raise TypeError('Type parameters should be placed on typing.Generic, not GenericModel') + if not hasattr(cls, '__parameters__'): + raise TypeError(f'Type {cls.__name__} must inherit from typing.Generic before being parameterized') + + check_parameters_count(cls, params) + # Build map from generic typevars to passed params + typevars_map: Dict[TypeVarType, Type[Any]] = dict(zip(cls.__parameters__, params)) + if all_identical(typevars_map.keys(), typevars_map.values()) and typevars_map: + return cls # if arguments are equal to parameters it's the same object + + # Create new model with original model as parent inserting fields with DeferredType. + model_name = cls.__concrete_name__(params) + validators = gather_all_validators(cls) + + type_hints = get_all_type_hints(cls).items() + instance_type_hints = {k: v for k, v in type_hints if get_origin(v) is not ClassVar} + + fields = {k: (DeferredType(), cls.__fields__[k].field_info) for k in instance_type_hints if k in cls.__fields__} + + model_module, called_globally = get_caller_frame_info() + created_model = cast( + Type[GenericModel], # casting ensures mypy is aware of the __concrete__ and __parameters__ attributes + create_model( + model_name, + __module__=model_module or cls.__module__, + __base__=(cls,) + tuple(cls.__parameterized_bases__(typevars_map)), + __config__=None, + __validators__=validators, + __cls_kwargs__=None, + **fields, + ), + ) + + _assigned_parameters[created_model] = typevars_map + + if called_globally: # create global reference and therefore allow pickling + object_by_reference = None + reference_name = model_name + reference_module_globals = sys.modules[created_model.__module__].__dict__ + while object_by_reference is not created_model: + object_by_reference = reference_module_globals.setdefault(reference_name, created_model) + reference_name += '_' + + created_model.Config = cls.Config + + # Find any typevars that are still present in the model. + # If none are left, the model is fully "concrete", otherwise the new + # class is a generic class as well taking the found typevars as + # parameters. + new_params = tuple( + {param: None for param in iter_contained_typevars(typevars_map.values())} + ) # use dict as ordered set + created_model.__concrete__ = not new_params + if new_params: + created_model.__parameters__ = new_params + + # Save created model in cache so we don't end up creating duplicate + # models that should be identical. + _generic_types_cache[_cache_key(params)] = created_model + if len(params) == 1: + _generic_types_cache[_cache_key(params[0])] = created_model + + # Recursively walk class type hints and replace generic typevars + # with concrete types that were passed. + _prepare_model_fields(created_model, fields, instance_type_hints, typevars_map) + + return created_model + + @classmethod + def __concrete_name__(cls: Type[Any], params: Tuple[Type[Any], ...]) -> str: + """Compute class name for child classes. + + :param params: Tuple of types the class . Given a generic class + `Model` with 2 type variables and a concrete model `Model[str, int]`, + the value `(str, int)` would be passed to `params`. + :return: String representing a the new class where `params` are + passed to `cls` as type variables. + + This method can be overridden to achieve a custom naming scheme for GenericModels. + """ + param_names = [display_as_type(param) for param in params] + params_component = ', '.join(param_names) + return f'{cls.__name__}[{params_component}]' + + @classmethod + def __parameterized_bases__(cls, typevars_map: Parametrization) -> Iterator[Type[Any]]: + """ + Returns unbound bases of cls parameterised to given type variables + + :param typevars_map: Dictionary of type applications for binding subclasses. + Given a generic class `Model` with 2 type variables [S, T] + and a concrete model `Model[str, int]`, + the value `{S: str, T: int}` would be passed to `typevars_map`. + :return: an iterator of generic sub classes, parameterised by `typevars_map` + and other assigned parameters of `cls` + + e.g.: + ``` + class A(GenericModel, Generic[T]): + ... + + class B(A[V], Generic[V]): + ... + + assert A[int] in B.__parameterized_bases__({V: int}) + ``` + """ + + def build_base_model( + base_model: Type[GenericModel], mapped_types: Parametrization + ) -> Iterator[Type[GenericModel]]: + base_parameters = tuple(mapped_types[param] for param in base_model.__parameters__) + parameterized_base = base_model.__class_getitem__(base_parameters) + if parameterized_base is base_model or parameterized_base is cls: + # Avoid duplication in MRO + return + yield parameterized_base + + for base_model in cls.__bases__: + if not issubclass(base_model, GenericModel): + # not a class that can be meaningfully parameterized + continue + elif not getattr(base_model, '__parameters__', None): + # base_model is "GenericModel" (and has no __parameters__) + # or + # base_model is already concrete, and will be included transitively via cls. + continue + elif cls in _assigned_parameters: + if base_model in _assigned_parameters: + # cls is partially parameterised but not from base_model + # e.g. cls = B[S], base_model = A[S] + # B[S][int] should subclass A[int], (and will be transitively via B[int]) + # but it's not viable to consistently subclass types with arbitrary construction + # So don't attempt to include A[S][int] + continue + else: # base_model not in _assigned_parameters: + # cls is partially parameterized, base_model is original generic + # e.g. cls = B[str, T], base_model = B[S, T] + # Need to determine the mapping for the base_model parameters + mapped_types: Parametrization = { + key: typevars_map.get(value, value) for key, value in _assigned_parameters[cls].items() + } + yield from build_base_model(base_model, mapped_types) + else: + # cls is base generic, so base_class has a distinct base + # can construct the Parameterised base model using typevars_map directly + yield from build_base_model(base_model, typevars_map) + + +def replace_types(type_: Any, type_map: Mapping[Any, Any]) -> Any: + """Return type with all occurrences of `type_map` keys recursively replaced with their values. + + :param type_: Any type, class or generic alias + :param type_map: Mapping from `TypeVar` instance to concrete types. + :return: New type representing the basic structure of `type_` with all + `typevar_map` keys recursively replaced. + + >>> replace_types(Tuple[str, Union[List[str], float]], {str: int}) + Tuple[int, Union[List[int], float]] + + """ + if not type_map: + return type_ + + type_args = get_args(type_) + origin_type = get_origin(type_) + + if origin_type is Annotated: + annotated_type, *annotations = type_args + return Annotated[replace_types(annotated_type, type_map), tuple(annotations)] + + if (origin_type is ExtLiteral) or (sys.version_info >= (3, 8) and origin_type is Literal): + return type_map.get(type_, type_) + # Having type args is a good indicator that this is a typing module + # class instantiation or a generic alias of some sort. + if type_args: + resolved_type_args = tuple(replace_types(arg, type_map) for arg in type_args) + if all_identical(type_args, resolved_type_args): + # If all arguments are the same, there is no need to modify the + # type or create a new object at all + return type_ + if ( + origin_type is not None + and isinstance(type_, typing_base) + and not isinstance(origin_type, typing_base) + and getattr(type_, '_name', None) is not None + ): + # In python < 3.9 generic aliases don't exist so any of these like `list`, + # `type` or `collections.abc.Callable` need to be translated. + # See: https://www.python.org/dev/peps/pep-0585 + origin_type = getattr(typing, type_._name) + assert origin_type is not None + # PEP-604 syntax (Ex.: list | str) is represented with a types.UnionType object that does not have __getitem__. + # We also cannot use isinstance() since we have to compare types. + if sys.version_info >= (3, 10) and origin_type is types.UnionType: # noqa: E721 + return _UnionGenericAlias(origin_type, resolved_type_args) + return origin_type[resolved_type_args] + + # We handle pydantic generic models separately as they don't have the same + # semantics as "typing" classes or generic aliases + if not origin_type and lenient_issubclass(type_, GenericModel) and not type_.__concrete__: + type_args = type_.__parameters__ + resolved_type_args = tuple(replace_types(t, type_map) for t in type_args) + if all_identical(type_args, resolved_type_args): + return type_ + return type_[resolved_type_args] + + # Handle special case for typehints that can have lists as arguments. + # `typing.Callable[[int, str], int]` is an example for this. + if isinstance(type_, (List, list)): + resolved_list = list(replace_types(element, type_map) for element in type_) + if all_identical(type_, resolved_list): + return type_ + return resolved_list + + # For JsonWrapperValue, need to handle its inner type to allow correct parsing + # of generic Json arguments like Json[T] + if not origin_type and lenient_issubclass(type_, JsonWrapper): + type_.inner_type = replace_types(type_.inner_type, type_map) + return type_ + + # If all else fails, we try to resolve the type directly and otherwise just + # return the input with no modifications. + new_type = type_map.get(type_, type_) + # Convert string to ForwardRef + if isinstance(new_type, str): + return ForwardRef(new_type) + else: + return new_type + + +def check_parameters_count(cls: Type[GenericModel], parameters: Tuple[Any, ...]) -> None: + actual = len(parameters) + expected = len(cls.__parameters__) + if actual != expected: + description = 'many' if actual > expected else 'few' + raise TypeError(f'Too {description} parameters for {cls.__name__}; actual {actual}, expected {expected}') + + +DictValues: Type[Any] = {}.values().__class__ + + +def iter_contained_typevars(v: Any) -> Iterator[TypeVarType]: + """Recursively iterate through all subtypes and type args of `v` and yield any typevars that are found.""" + if isinstance(v, TypeVar): + yield v + elif hasattr(v, '__parameters__') and not get_origin(v) and lenient_issubclass(v, GenericModel): + yield from v.__parameters__ + elif isinstance(v, (DictValues, list)): + for var in v: + yield from iter_contained_typevars(var) + else: + args = get_args(v) + for arg in args: + yield from iter_contained_typevars(arg) + + +def get_caller_frame_info() -> Tuple[Optional[str], bool]: + """ + Used inside a function to check whether it was called globally + + Will only work against non-compiled code, therefore used only in pydantic.generics + + :returns Tuple[module_name, called_globally] + """ + try: + previous_caller_frame = sys._getframe(2) + except ValueError as e: + raise RuntimeError('This function must be used inside another function') from e + except AttributeError: # sys module does not have _getframe function, so there's nothing we can do about it + return None, False + frame_globals = previous_caller_frame.f_globals + return frame_globals.get('__name__'), previous_caller_frame.f_locals is frame_globals + + +def _prepare_model_fields( + created_model: Type[GenericModel], + fields: Mapping[str, Any], + instance_type_hints: Mapping[str, type], + typevars_map: Mapping[Any, type], +) -> None: + """ + Replace DeferredType fields with concrete type hints and prepare them. + """ + + for key, field in created_model.__fields__.items(): + if key not in fields: + assert field.type_.__class__ is not DeferredType + # https://github.com/nedbat/coveragepy/issues/198 + continue # pragma: no cover + + assert field.type_.__class__ is DeferredType, field.type_.__class__ + + field_type_hint = instance_type_hints[key] + concrete_type = replace_types(field_type_hint, typevars_map) + field.type_ = concrete_type + field.outer_type_ = concrete_type + field.prepare() + created_model.__annotations__[key] = concrete_type diff --git a/lib/python3.12/site-packages/pydantic/v1/json.py b/lib/python3.12/site-packages/pydantic/v1/json.py new file mode 100644 index 0000000000000000000000000000000000000000..41d0d5fcae0d946401c2d089b69143f540618b90 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/json.py @@ -0,0 +1,112 @@ +import datetime +from collections import deque +from decimal import Decimal +from enum import Enum +from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network +from pathlib import Path +from re import Pattern +from types import GeneratorType +from typing import Any, Callable, Dict, Type, Union +from uuid import UUID + +from pydantic.v1.color import Color +from pydantic.v1.networks import NameEmail +from pydantic.v1.types import SecretBytes, SecretStr + +__all__ = 'pydantic_encoder', 'custom_pydantic_encoder', 'timedelta_isoformat' + + +def isoformat(o: Union[datetime.date, datetime.time]) -> str: + return o.isoformat() + + +def decimal_encoder(dec_value: Decimal) -> Union[int, float]: + """ + Encodes a Decimal as int of there's no exponent, otherwise float + + This is useful when we use ConstrainedDecimal to represent Numeric(x,0) + where a integer (but not int typed) is used. Encoding this as a float + results in failed round-tripping between encode and parse. + Our Id type is a prime example of this. + + >>> decimal_encoder(Decimal("1.0")) + 1.0 + + >>> decimal_encoder(Decimal("1")) + 1 + """ + if dec_value.as_tuple().exponent >= 0: + return int(dec_value) + else: + return float(dec_value) + + +ENCODERS_BY_TYPE: Dict[Type[Any], Callable[[Any], Any]] = { + bytes: lambda o: o.decode(), + Color: str, + datetime.date: isoformat, + datetime.datetime: isoformat, + datetime.time: isoformat, + datetime.timedelta: lambda td: td.total_seconds(), + Decimal: decimal_encoder, + Enum: lambda o: o.value, + frozenset: list, + deque: list, + GeneratorType: list, + IPv4Address: str, + IPv4Interface: str, + IPv4Network: str, + IPv6Address: str, + IPv6Interface: str, + IPv6Network: str, + NameEmail: str, + Path: str, + Pattern: lambda o: o.pattern, + SecretBytes: str, + SecretStr: str, + set: list, + UUID: str, +} + + +def pydantic_encoder(obj: Any) -> Any: + from dataclasses import asdict, is_dataclass + + from pydantic.v1.main import BaseModel + + if isinstance(obj, BaseModel): + return obj.dict() + elif is_dataclass(obj): + return asdict(obj) + + # Check the class type and its superclasses for a matching encoder + for base in obj.__class__.__mro__[:-1]: + try: + encoder = ENCODERS_BY_TYPE[base] + except KeyError: + continue + return encoder(obj) + else: # We have exited the for loop without finding a suitable encoder + raise TypeError(f"Object of type '{obj.__class__.__name__}' is not JSON serializable") + + +def custom_pydantic_encoder(type_encoders: Dict[Any, Callable[[Type[Any]], Any]], obj: Any) -> Any: + # Check the class type and its superclasses for a matching encoder + for base in obj.__class__.__mro__[:-1]: + try: + encoder = type_encoders[base] + except KeyError: + continue + + return encoder(obj) + else: # We have exited the for loop without finding a suitable encoder + return pydantic_encoder(obj) + + +def timedelta_isoformat(td: datetime.timedelta) -> str: + """ + ISO 8601 encoding for Python timedelta object. + """ + minutes, seconds = divmod(td.seconds, 60) + hours, minutes = divmod(minutes, 60) + return f'{"-" if td.days < 0 else ""}P{abs(td.days)}DT{hours:d}H{minutes:d}M{seconds:d}.{td.microseconds:06d}S' diff --git a/lib/python3.12/site-packages/pydantic/v1/main.py b/lib/python3.12/site-packages/pydantic/v1/main.py new file mode 100644 index 0000000000000000000000000000000000000000..8000967eafc52526b43e451106a15996c46969c9 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/main.py @@ -0,0 +1,1113 @@ +import warnings +from abc import ABCMeta +from copy import deepcopy +from enum import Enum +from functools import partial +from pathlib import Path +from types import FunctionType, prepare_class, resolve_bases +from typing import ( + TYPE_CHECKING, + AbstractSet, + Any, + Callable, + ClassVar, + Dict, + List, + Mapping, + Optional, + Tuple, + Type, + TypeVar, + Union, + cast, + no_type_check, + overload, +) + +from typing_extensions import dataclass_transform + +from pydantic.v1.class_validators import ValidatorGroup, extract_root_validators, extract_validators, inherit_validators +from pydantic.v1.config import BaseConfig, Extra, inherit_config, prepare_config +from pydantic.v1.error_wrappers import ErrorWrapper, ValidationError +from pydantic.v1.errors import ConfigError, DictError, ExtraError, MissingError +from pydantic.v1.fields import ( + MAPPING_LIKE_SHAPES, + Field, + ModelField, + ModelPrivateAttr, + PrivateAttr, + Undefined, + is_finalvar_with_default_val, +) +from pydantic.v1.json import custom_pydantic_encoder, pydantic_encoder +from pydantic.v1.parse import Protocol, load_file, load_str_bytes +from pydantic.v1.schema import default_ref_template, model_schema +from pydantic.v1.types import PyObject, StrBytes +from pydantic.v1.typing import ( + AnyCallable, + get_args, + get_origin, + is_classvar, + is_namedtuple, + is_union, + resolve_annotations, + update_model_forward_refs, +) +from pydantic.v1.utils import ( + DUNDER_ATTRIBUTES, + ROOT_KEY, + ClassAttribute, + GetterDict, + Representation, + ValueItems, + generate_model_signature, + is_valid_field, + is_valid_private_name, + lenient_issubclass, + sequence_like, + smart_deepcopy, + unique_list, + validate_field_name, +) + +if TYPE_CHECKING: + from inspect import Signature + + from pydantic.v1.class_validators import ValidatorListDict + from pydantic.v1.types import ModelOrDc + from pydantic.v1.typing import ( + AbstractSetIntStr, + AnyClassMethod, + CallableGenerator, + DictAny, + DictStrAny, + MappingIntStrAny, + ReprArgs, + SetStr, + TupleGenerator, + ) + + Model = TypeVar('Model', bound='BaseModel') + +__all__ = 'BaseModel', 'create_model', 'validate_model' + +_T = TypeVar('_T') + + +def validate_custom_root_type(fields: Dict[str, ModelField]) -> None: + if len(fields) > 1: + raise ValueError(f'{ROOT_KEY} cannot be mixed with other fields') + + +def generate_hash_function(frozen: bool) -> Optional[Callable[[Any], int]]: + def hash_function(self_: Any) -> int: + return hash(self_.__class__) + hash(tuple(self_.__dict__.values())) + + return hash_function if frozen else None + + +# If a field is of type `Callable`, its default value should be a function and cannot to ignored. +ANNOTATED_FIELD_UNTOUCHED_TYPES: Tuple[Any, ...] = (property, type, classmethod, staticmethod) +# When creating a `BaseModel` instance, we bypass all the methods, properties... added to the model +UNTOUCHED_TYPES: Tuple[Any, ...] = (FunctionType,) + ANNOTATED_FIELD_UNTOUCHED_TYPES +# Note `ModelMetaclass` refers to `BaseModel`, but is also used to *create* `BaseModel`, so we need to add this extra +# (somewhat hacky) boolean to keep track of whether we've created the `BaseModel` class yet, and therefore whether it's +# safe to refer to it. If it *hasn't* been created, we assume that the `__new__` call we're in the middle of is for +# the `BaseModel` class, since that's defined immediately after the metaclass. +_is_base_model_class_defined = False + + +@dataclass_transform(kw_only_default=True, field_specifiers=(Field,)) +class ModelMetaclass(ABCMeta): + @no_type_check # noqa C901 + def __new__(mcs, name, bases, namespace, **kwargs): # noqa C901 + fields: Dict[str, ModelField] = {} + config = BaseConfig + validators: 'ValidatorListDict' = {} + + pre_root_validators, post_root_validators = [], [] + private_attributes: Dict[str, ModelPrivateAttr] = {} + base_private_attributes: Dict[str, ModelPrivateAttr] = {} + slots: SetStr = namespace.get('__slots__', ()) + slots = {slots} if isinstance(slots, str) else set(slots) + class_vars: SetStr = set() + hash_func: Optional[Callable[[Any], int]] = None + + for base in reversed(bases): + if _is_base_model_class_defined and issubclass(base, BaseModel) and base != BaseModel: + fields.update(smart_deepcopy(base.__fields__)) + config = inherit_config(base.__config__, config) + validators = inherit_validators(base.__validators__, validators) + pre_root_validators += base.__pre_root_validators__ + post_root_validators += base.__post_root_validators__ + base_private_attributes.update(base.__private_attributes__) + class_vars.update(base.__class_vars__) + hash_func = base.__hash__ + + resolve_forward_refs = kwargs.pop('__resolve_forward_refs__', True) + allowed_config_kwargs: SetStr = { + key + for key in dir(config) + if not (key.startswith('__') and key.endswith('__')) # skip dunder methods and attributes + } + config_kwargs = {key: kwargs.pop(key) for key in kwargs.keys() & allowed_config_kwargs} + config_from_namespace = namespace.get('Config') + if config_kwargs and config_from_namespace: + raise TypeError('Specifying config in two places is ambiguous, use either Config attribute or class kwargs') + config = inherit_config(config_from_namespace, config, **config_kwargs) + + validators = inherit_validators(extract_validators(namespace), validators) + vg = ValidatorGroup(validators) + + for f in fields.values(): + f.set_config(config) + extra_validators = vg.get_validators(f.name) + if extra_validators: + f.class_validators.update(extra_validators) + # re-run prepare to add extra validators + f.populate_validators() + + prepare_config(config, name) + + untouched_types = ANNOTATED_FIELD_UNTOUCHED_TYPES + + def is_untouched(v: Any) -> bool: + return isinstance(v, untouched_types) or v.__class__.__name__ == 'cython_function_or_method' + + if (namespace.get('__module__'), namespace.get('__qualname__')) != ('pydantic.main', 'BaseModel'): + annotations = resolve_annotations(namespace.get('__annotations__', {}), namespace.get('__module__', None)) + # annotation only fields need to come first in fields + for ann_name, ann_type in annotations.items(): + if is_classvar(ann_type): + class_vars.add(ann_name) + elif is_finalvar_with_default_val(ann_type, namespace.get(ann_name, Undefined)): + class_vars.add(ann_name) + elif is_valid_field(ann_name): + validate_field_name(bases, ann_name) + value = namespace.get(ann_name, Undefined) + allowed_types = get_args(ann_type) if is_union(get_origin(ann_type)) else (ann_type,) + if ( + is_untouched(value) + and ann_type != PyObject + and not any( + lenient_issubclass(get_origin(allowed_type), Type) for allowed_type in allowed_types + ) + ): + continue + fields[ann_name] = ModelField.infer( + name=ann_name, + value=value, + annotation=ann_type, + class_validators=vg.get_validators(ann_name), + config=config, + ) + elif ann_name not in namespace and config.underscore_attrs_are_private: + private_attributes[ann_name] = PrivateAttr() + + untouched_types = UNTOUCHED_TYPES + config.keep_untouched + for var_name, value in namespace.items(): + can_be_changed = var_name not in class_vars and not is_untouched(value) + if isinstance(value, ModelPrivateAttr): + if not is_valid_private_name(var_name): + raise NameError( + f'Private attributes "{var_name}" must not be a valid field name; ' + f'Use sunder or dunder names, e. g. "_{var_name}" or "__{var_name}__"' + ) + private_attributes[var_name] = value + elif config.underscore_attrs_are_private and is_valid_private_name(var_name) and can_be_changed: + private_attributes[var_name] = PrivateAttr(default=value) + elif is_valid_field(var_name) and var_name not in annotations and can_be_changed: + validate_field_name(bases, var_name) + inferred = ModelField.infer( + name=var_name, + value=value, + annotation=annotations.get(var_name, Undefined), + class_validators=vg.get_validators(var_name), + config=config, + ) + if var_name in fields: + if lenient_issubclass(inferred.type_, fields[var_name].type_): + inferred.type_ = fields[var_name].type_ + else: + raise TypeError( + f'The type of {name}.{var_name} differs from the new default value; ' + f'if you wish to change the type of this field, please use a type annotation' + ) + fields[var_name] = inferred + + _custom_root_type = ROOT_KEY in fields + if _custom_root_type: + validate_custom_root_type(fields) + vg.check_for_unused() + if config.json_encoders: + json_encoder = partial(custom_pydantic_encoder, config.json_encoders) + else: + json_encoder = pydantic_encoder + pre_rv_new, post_rv_new = extract_root_validators(namespace) + + if hash_func is None: + hash_func = generate_hash_function(config.frozen) + + exclude_from_namespace = fields | private_attributes.keys() | {'__slots__'} + new_namespace = { + '__config__': config, + '__fields__': fields, + '__exclude_fields__': { + name: field.field_info.exclude for name, field in fields.items() if field.field_info.exclude is not None + } + or None, + '__include_fields__': { + name: field.field_info.include for name, field in fields.items() if field.field_info.include is not None + } + or None, + '__validators__': vg.validators, + '__pre_root_validators__': unique_list( + pre_root_validators + pre_rv_new, + name_factory=lambda v: v.__name__, + ), + '__post_root_validators__': unique_list( + post_root_validators + post_rv_new, + name_factory=lambda skip_on_failure_and_v: skip_on_failure_and_v[1].__name__, + ), + '__schema_cache__': {}, + '__json_encoder__': staticmethod(json_encoder), + '__custom_root_type__': _custom_root_type, + '__private_attributes__': {**base_private_attributes, **private_attributes}, + '__slots__': slots | private_attributes.keys(), + '__hash__': hash_func, + '__class_vars__': class_vars, + **{n: v for n, v in namespace.items() if n not in exclude_from_namespace}, + } + + cls = super().__new__(mcs, name, bases, new_namespace, **kwargs) + # set __signature__ attr only for model class, but not for its instances + cls.__signature__ = ClassAttribute('__signature__', generate_model_signature(cls.__init__, fields, config)) + + if not _is_base_model_class_defined: + # Cython does not understand the `if TYPE_CHECKING:` condition in the + # BaseModel's body (where annotations are set), so clear them manually: + getattr(cls, '__annotations__', {}).clear() + + if resolve_forward_refs: + cls.__try_update_forward_refs__() + + # preserve `__set_name__` protocol defined in https://peps.python.org/pep-0487 + # for attributes not in `new_namespace` (e.g. private attributes) + for name, obj in namespace.items(): + if name not in new_namespace: + set_name = getattr(obj, '__set_name__', None) + if callable(set_name): + set_name(cls, name) + + return cls + + def __instancecheck__(self, instance: Any) -> bool: + """ + Avoid calling ABC _abc_subclasscheck unless we're pretty sure. + + See #3829 and python/cpython#92810 + """ + return hasattr(instance, '__post_root_validators__') and super().__instancecheck__(instance) + + +object_setattr = object.__setattr__ + + +class BaseModel(Representation, metaclass=ModelMetaclass): + if TYPE_CHECKING: + # populated by the metaclass, defined here to help IDEs only + __fields__: ClassVar[Dict[str, ModelField]] = {} + __include_fields__: ClassVar[Optional[Mapping[str, Any]]] = None + __exclude_fields__: ClassVar[Optional[Mapping[str, Any]]] = None + __validators__: ClassVar[Dict[str, AnyCallable]] = {} + __pre_root_validators__: ClassVar[List[AnyCallable]] + __post_root_validators__: ClassVar[List[Tuple[bool, AnyCallable]]] + __config__: ClassVar[Type[BaseConfig]] = BaseConfig + __json_encoder__: ClassVar[Callable[[Any], Any]] = lambda x: x + __schema_cache__: ClassVar['DictAny'] = {} + __custom_root_type__: ClassVar[bool] = False + __signature__: ClassVar['Signature'] + __private_attributes__: ClassVar[Dict[str, ModelPrivateAttr]] + __class_vars__: ClassVar[SetStr] + __fields_set__: ClassVar[SetStr] = set() + + Config = BaseConfig + __slots__ = ('__dict__', '__fields_set__') + __doc__ = '' # Null out the Representation docstring + + def __init__(__pydantic_self__, **data: Any) -> None: + """ + Create a new model by parsing and validating input data from keyword arguments. + + Raises ValidationError if the input data cannot be parsed to form a valid model. + """ + # Uses something other than `self` the first arg to allow "self" as a settable attribute + values, fields_set, validation_error = validate_model(__pydantic_self__.__class__, data) + if validation_error: + raise validation_error + try: + object_setattr(__pydantic_self__, '__dict__', values) + except TypeError as e: + raise TypeError( + 'Model values must be a dict; you may not have returned a dictionary from a root validator' + ) from e + object_setattr(__pydantic_self__, '__fields_set__', fields_set) + __pydantic_self__._init_private_attributes() + + @no_type_check + def __setattr__(self, name, value): # noqa: C901 (ignore complexity) + if name in self.__private_attributes__ or name in DUNDER_ATTRIBUTES: + return object_setattr(self, name, value) + + if self.__config__.extra is not Extra.allow and name not in self.__fields__: + raise ValueError(f'"{self.__class__.__name__}" object has no field "{name}"') + elif not self.__config__.allow_mutation or self.__config__.frozen: + raise TypeError(f'"{self.__class__.__name__}" is immutable and does not support item assignment') + elif name in self.__fields__ and self.__fields__[name].final: + raise TypeError( + f'"{self.__class__.__name__}" object "{name}" field is final and does not support reassignment' + ) + elif self.__config__.validate_assignment: + new_values = {**self.__dict__, name: value} + + for validator in self.__pre_root_validators__: + try: + new_values = validator(self.__class__, new_values) + except (ValueError, TypeError, AssertionError) as exc: + raise ValidationError([ErrorWrapper(exc, loc=ROOT_KEY)], self.__class__) + + known_field = self.__fields__.get(name, None) + if known_field: + # We want to + # - make sure validators are called without the current value for this field inside `values` + # - keep other values (e.g. submodels) untouched (using `BaseModel.dict()` will change them into dicts) + # - keep the order of the fields + if not known_field.field_info.allow_mutation: + raise TypeError(f'"{known_field.name}" has allow_mutation set to False and cannot be assigned') + dict_without_original_value = {k: v for k, v in self.__dict__.items() if k != name} + value, error_ = known_field.validate(value, dict_without_original_value, loc=name, cls=self.__class__) + if error_: + raise ValidationError([error_], self.__class__) + else: + new_values[name] = value + + errors = [] + for skip_on_failure, validator in self.__post_root_validators__: + if skip_on_failure and errors: + continue + try: + new_values = validator(self.__class__, new_values) + except (ValueError, TypeError, AssertionError) as exc: + errors.append(ErrorWrapper(exc, loc=ROOT_KEY)) + if errors: + raise ValidationError(errors, self.__class__) + + # update the whole __dict__ as other values than just `value` + # may be changed (e.g. with `root_validator`) + object_setattr(self, '__dict__', new_values) + else: + self.__dict__[name] = value + + self.__fields_set__.add(name) + + def __getstate__(self) -> 'DictAny': + private_attrs = ((k, getattr(self, k, Undefined)) for k in self.__private_attributes__) + return { + '__dict__': self.__dict__, + '__fields_set__': self.__fields_set__, + '__private_attribute_values__': {k: v for k, v in private_attrs if v is not Undefined}, + } + + def __setstate__(self, state: 'DictAny') -> None: + object_setattr(self, '__dict__', state['__dict__']) + object_setattr(self, '__fields_set__', state['__fields_set__']) + for name, value in state.get('__private_attribute_values__', {}).items(): + object_setattr(self, name, value) + + def _init_private_attributes(self) -> None: + for name, private_attr in self.__private_attributes__.items(): + default = private_attr.get_default() + if default is not Undefined: + object_setattr(self, name, default) + + def dict( + self, + *, + include: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']] = None, + exclude: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']] = None, + by_alias: bool = False, + skip_defaults: Optional[bool] = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + ) -> 'DictStrAny': + """ + Generate a dictionary representation of the model, optionally specifying which fields to include or exclude. + + """ + if skip_defaults is not None: + warnings.warn( + f'{self.__class__.__name__}.dict(): "skip_defaults" is deprecated and replaced by "exclude_unset"', + DeprecationWarning, + ) + exclude_unset = skip_defaults + + return dict( + self._iter( + to_dict=True, + by_alias=by_alias, + include=include, + exclude=exclude, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + ) + ) + + def json( + self, + *, + include: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']] = None, + exclude: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']] = None, + by_alias: bool = False, + skip_defaults: Optional[bool] = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + encoder: Optional[Callable[[Any], Any]] = None, + models_as_dict: bool = True, + **dumps_kwargs: Any, + ) -> str: + """ + Generate a JSON representation of the model, `include` and `exclude` arguments as per `dict()`. + + `encoder` is an optional function to supply as `default` to json.dumps(), other arguments as per `json.dumps()`. + """ + if skip_defaults is not None: + warnings.warn( + f'{self.__class__.__name__}.json(): "skip_defaults" is deprecated and replaced by "exclude_unset"', + DeprecationWarning, + ) + exclude_unset = skip_defaults + encoder = cast(Callable[[Any], Any], encoder or self.__json_encoder__) + + # We don't directly call `self.dict()`, which does exactly this with `to_dict=True` + # because we want to be able to keep raw `BaseModel` instances and not as `dict`. + # This allows users to write custom JSON encoders for given `BaseModel` classes. + data = dict( + self._iter( + to_dict=models_as_dict, + by_alias=by_alias, + include=include, + exclude=exclude, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + ) + ) + if self.__custom_root_type__: + data = data[ROOT_KEY] + return self.__config__.json_dumps(data, default=encoder, **dumps_kwargs) + + @classmethod + def _enforce_dict_if_root(cls, obj: Any) -> Any: + if cls.__custom_root_type__ and ( + not (isinstance(obj, dict) and obj.keys() == {ROOT_KEY}) + and not (isinstance(obj, BaseModel) and obj.__fields__.keys() == {ROOT_KEY}) + or cls.__fields__[ROOT_KEY].shape in MAPPING_LIKE_SHAPES + ): + return {ROOT_KEY: obj} + else: + return obj + + @classmethod + def parse_obj(cls: Type['Model'], obj: Any) -> 'Model': + obj = cls._enforce_dict_if_root(obj) + if not isinstance(obj, dict): + try: + obj = dict(obj) + except (TypeError, ValueError) as e: + exc = TypeError(f'{cls.__name__} expected dict not {obj.__class__.__name__}') + raise ValidationError([ErrorWrapper(exc, loc=ROOT_KEY)], cls) from e + return cls(**obj) + + @classmethod + def parse_raw( + cls: Type['Model'], + b: StrBytes, + *, + content_type: str = None, + encoding: str = 'utf8', + proto: Protocol = None, + allow_pickle: bool = False, + ) -> 'Model': + try: + obj = load_str_bytes( + b, + proto=proto, + content_type=content_type, + encoding=encoding, + allow_pickle=allow_pickle, + json_loads=cls.__config__.json_loads, + ) + except (ValueError, TypeError, UnicodeDecodeError) as e: + raise ValidationError([ErrorWrapper(e, loc=ROOT_KEY)], cls) + return cls.parse_obj(obj) + + @classmethod + def parse_file( + cls: Type['Model'], + path: Union[str, Path], + *, + content_type: str = None, + encoding: str = 'utf8', + proto: Protocol = None, + allow_pickle: bool = False, + ) -> 'Model': + obj = load_file( + path, + proto=proto, + content_type=content_type, + encoding=encoding, + allow_pickle=allow_pickle, + json_loads=cls.__config__.json_loads, + ) + return cls.parse_obj(obj) + + @classmethod + def from_orm(cls: Type['Model'], obj: Any) -> 'Model': + if not cls.__config__.orm_mode: + raise ConfigError('You must have the config attribute orm_mode=True to use from_orm') + obj = {ROOT_KEY: obj} if cls.__custom_root_type__ else cls._decompose_class(obj) + m = cls.__new__(cls) + values, fields_set, validation_error = validate_model(cls, obj) + if validation_error: + raise validation_error + object_setattr(m, '__dict__', values) + object_setattr(m, '__fields_set__', fields_set) + m._init_private_attributes() + return m + + @classmethod + def construct(cls: Type['Model'], _fields_set: Optional['SetStr'] = None, **values: Any) -> 'Model': + """ + Creates a new model setting __dict__ and __fields_set__ from trusted or pre-validated data. + Default values are respected, but no other validation is performed. + Behaves as if `Config.extra = 'allow'` was set since it adds all passed values + """ + m = cls.__new__(cls) + fields_values: Dict[str, Any] = {} + for name, field in cls.__fields__.items(): + if field.alt_alias and field.alias in values: + fields_values[name] = values[field.alias] + elif name in values: + fields_values[name] = values[name] + elif not field.required: + fields_values[name] = field.get_default() + fields_values.update(values) + object_setattr(m, '__dict__', fields_values) + if _fields_set is None: + _fields_set = set(values.keys()) + object_setattr(m, '__fields_set__', _fields_set) + m._init_private_attributes() + return m + + def _copy_and_set_values(self: 'Model', values: 'DictStrAny', fields_set: 'SetStr', *, deep: bool) -> 'Model': + if deep: + # chances of having empty dict here are quite low for using smart_deepcopy + values = deepcopy(values) + + cls = self.__class__ + m = cls.__new__(cls) + object_setattr(m, '__dict__', values) + object_setattr(m, '__fields_set__', fields_set) + for name in self.__private_attributes__: + value = getattr(self, name, Undefined) + if value is not Undefined: + if deep: + value = deepcopy(value) + object_setattr(m, name, value) + + return m + + def copy( + self: 'Model', + *, + include: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']] = None, + exclude: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']] = None, + update: Optional['DictStrAny'] = None, + deep: bool = False, + ) -> 'Model': + """ + Duplicate a model, optionally choose which fields to include, exclude and change. + + :param include: fields to include in new model + :param exclude: fields to exclude from new model, as with values this takes precedence over include + :param update: values to change/add in the new model. Note: the data is not validated before creating + the new model: you should trust this data + :param deep: set to `True` to make a deep copy of the model + :return: new model instance + """ + + values = dict( + self._iter(to_dict=False, by_alias=False, include=include, exclude=exclude, exclude_unset=False), + **(update or {}), + ) + + # new `__fields_set__` can have unset optional fields with a set value in `update` kwarg + if update: + fields_set = self.__fields_set__ | update.keys() + else: + fields_set = set(self.__fields_set__) + + return self._copy_and_set_values(values, fields_set, deep=deep) + + @classmethod + def schema(cls, by_alias: bool = True, ref_template: str = default_ref_template) -> 'DictStrAny': + cached = cls.__schema_cache__.get((by_alias, ref_template)) + if cached is not None: + return cached + s = model_schema(cls, by_alias=by_alias, ref_template=ref_template) + cls.__schema_cache__[(by_alias, ref_template)] = s + return s + + @classmethod + def schema_json( + cls, *, by_alias: bool = True, ref_template: str = default_ref_template, **dumps_kwargs: Any + ) -> str: + from pydantic.v1.json import pydantic_encoder + + return cls.__config__.json_dumps( + cls.schema(by_alias=by_alias, ref_template=ref_template), default=pydantic_encoder, **dumps_kwargs + ) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + + @classmethod + def validate(cls: Type['Model'], value: Any) -> 'Model': + if isinstance(value, cls): + copy_on_model_validation = cls.__config__.copy_on_model_validation + # whether to deep or shallow copy the model on validation, None means do not copy + deep_copy: Optional[bool] = None + if copy_on_model_validation not in {'deep', 'shallow', 'none'}: + # Warn about deprecated behavior + warnings.warn( + "`copy_on_model_validation` should be a string: 'deep', 'shallow' or 'none'", DeprecationWarning + ) + if copy_on_model_validation: + deep_copy = False + + if copy_on_model_validation == 'shallow': + # shallow copy + deep_copy = False + elif copy_on_model_validation == 'deep': + # deep copy + deep_copy = True + + if deep_copy is None: + return value + else: + return value._copy_and_set_values(value.__dict__, value.__fields_set__, deep=deep_copy) + + value = cls._enforce_dict_if_root(value) + + if isinstance(value, dict): + return cls(**value) + elif cls.__config__.orm_mode: + return cls.from_orm(value) + else: + try: + value_as_dict = dict(value) + except (TypeError, ValueError) as e: + raise DictError() from e + return cls(**value_as_dict) + + @classmethod + def _decompose_class(cls: Type['Model'], obj: Any) -> GetterDict: + if isinstance(obj, GetterDict): + return obj + return cls.__config__.getter_dict(obj) + + @classmethod + @no_type_check + def _get_value( + cls, + v: Any, + to_dict: bool, + by_alias: bool, + include: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']], + exclude: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']], + exclude_unset: bool, + exclude_defaults: bool, + exclude_none: bool, + ) -> Any: + if isinstance(v, BaseModel): + if to_dict: + v_dict = v.dict( + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + include=include, + exclude=exclude, + exclude_none=exclude_none, + ) + if ROOT_KEY in v_dict: + return v_dict[ROOT_KEY] + return v_dict + else: + return v.copy(include=include, exclude=exclude) + + value_exclude = ValueItems(v, exclude) if exclude else None + value_include = ValueItems(v, include) if include else None + + if isinstance(v, dict): + return { + k_: cls._get_value( + v_, + to_dict=to_dict, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + include=value_include and value_include.for_element(k_), + exclude=value_exclude and value_exclude.for_element(k_), + exclude_none=exclude_none, + ) + for k_, v_ in v.items() + if (not value_exclude or not value_exclude.is_excluded(k_)) + and (not value_include or value_include.is_included(k_)) + } + + elif sequence_like(v): + seq_args = ( + cls._get_value( + v_, + to_dict=to_dict, + by_alias=by_alias, + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + include=value_include and value_include.for_element(i), + exclude=value_exclude and value_exclude.for_element(i), + exclude_none=exclude_none, + ) + for i, v_ in enumerate(v) + if (not value_exclude or not value_exclude.is_excluded(i)) + and (not value_include or value_include.is_included(i)) + ) + + return v.__class__(*seq_args) if is_namedtuple(v.__class__) else v.__class__(seq_args) + + elif isinstance(v, Enum) and getattr(cls.Config, 'use_enum_values', False): + return v.value + + else: + return v + + @classmethod + def __try_update_forward_refs__(cls, **localns: Any) -> None: + """ + Same as update_forward_refs but will not raise exception + when forward references are not defined. + """ + update_model_forward_refs(cls, cls.__fields__.values(), cls.__config__.json_encoders, localns, (NameError,)) + + @classmethod + def update_forward_refs(cls, **localns: Any) -> None: + """ + Try to update ForwardRefs on fields based on this Model, globalns and localns. + """ + update_model_forward_refs(cls, cls.__fields__.values(), cls.__config__.json_encoders, localns) + + def __iter__(self) -> 'TupleGenerator': + """ + so `dict(model)` works + """ + yield from self.__dict__.items() + + def _iter( + self, + to_dict: bool = False, + by_alias: bool = False, + include: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']] = None, + exclude: Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']] = None, + exclude_unset: bool = False, + exclude_defaults: bool = False, + exclude_none: bool = False, + ) -> 'TupleGenerator': + # Merge field set excludes with explicit exclude parameter with explicit overriding field set options. + # The extra "is not None" guards are not logically necessary but optimizes performance for the simple case. + if exclude is not None or self.__exclude_fields__ is not None: + exclude = ValueItems.merge(self.__exclude_fields__, exclude) + + if include is not None or self.__include_fields__ is not None: + include = ValueItems.merge(self.__include_fields__, include, intersect=True) + + allowed_keys = self._calculate_keys( + include=include, exclude=exclude, exclude_unset=exclude_unset # type: ignore + ) + if allowed_keys is None and not (to_dict or by_alias or exclude_unset or exclude_defaults or exclude_none): + # huge boost for plain _iter() + yield from self.__dict__.items() + return + + value_exclude = ValueItems(self, exclude) if exclude is not None else None + value_include = ValueItems(self, include) if include is not None else None + + for field_key, v in self.__dict__.items(): + if (allowed_keys is not None and field_key not in allowed_keys) or (exclude_none and v is None): + continue + + if exclude_defaults: + model_field = self.__fields__.get(field_key) + if not getattr(model_field, 'required', True) and getattr(model_field, 'default', _missing) == v: + continue + + if by_alias and field_key in self.__fields__: + dict_key = self.__fields__[field_key].alias + else: + dict_key = field_key + + if to_dict or value_include or value_exclude: + v = self._get_value( + v, + to_dict=to_dict, + by_alias=by_alias, + include=value_include and value_include.for_element(field_key), + exclude=value_exclude and value_exclude.for_element(field_key), + exclude_unset=exclude_unset, + exclude_defaults=exclude_defaults, + exclude_none=exclude_none, + ) + yield dict_key, v + + def _calculate_keys( + self, + include: Optional['MappingIntStrAny'], + exclude: Optional['MappingIntStrAny'], + exclude_unset: bool, + update: Optional['DictStrAny'] = None, + ) -> Optional[AbstractSet[str]]: + if include is None and exclude is None and exclude_unset is False: + return None + + keys: AbstractSet[str] + if exclude_unset: + keys = self.__fields_set__.copy() + else: + keys = self.__dict__.keys() + + if include is not None: + keys &= include.keys() + + if update: + keys -= update.keys() + + if exclude: + keys -= {k for k, v in exclude.items() if ValueItems.is_true(v)} + + return keys + + def __eq__(self, other: Any) -> bool: + if isinstance(other, BaseModel): + return self.dict() == other.dict() + else: + return self.dict() == other + + def __repr_args__(self) -> 'ReprArgs': + return [ + (k, v) + for k, v in self.__dict__.items() + if k not in DUNDER_ATTRIBUTES and (k not in self.__fields__ or self.__fields__[k].field_info.repr) + ] + + +_is_base_model_class_defined = True + + +@overload +def create_model( + __model_name: str, + *, + __config__: Optional[Type[BaseConfig]] = None, + __base__: None = None, + __module__: str = __name__, + __validators__: Dict[str, 'AnyClassMethod'] = None, + __cls_kwargs__: Dict[str, Any] = None, + **field_definitions: Any, +) -> Type['BaseModel']: + ... + + +@overload +def create_model( + __model_name: str, + *, + __config__: Optional[Type[BaseConfig]] = None, + __base__: Union[Type['Model'], Tuple[Type['Model'], ...]], + __module__: str = __name__, + __validators__: Dict[str, 'AnyClassMethod'] = None, + __cls_kwargs__: Dict[str, Any] = None, + **field_definitions: Any, +) -> Type['Model']: + ... + + +def create_model( + __model_name: str, + *, + __config__: Optional[Type[BaseConfig]] = None, + __base__: Union[None, Type['Model'], Tuple[Type['Model'], ...]] = None, + __module__: str = __name__, + __validators__: Dict[str, 'AnyClassMethod'] = None, + __cls_kwargs__: Dict[str, Any] = None, + __slots__: Optional[Tuple[str, ...]] = None, + **field_definitions: Any, +) -> Type['Model']: + """ + Dynamically create a model. + :param __model_name: name of the created model + :param __config__: config class to use for the new model + :param __base__: base class for the new model to inherit from + :param __module__: module of the created model + :param __validators__: a dict of method names and @validator class methods + :param __cls_kwargs__: a dict for class creation + :param __slots__: Deprecated, `__slots__` should not be passed to `create_model` + :param field_definitions: fields of the model (or extra fields if a base is supplied) + in the format `=(, )` or `=, e.g. + `foobar=(str, ...)` or `foobar=123`, or, for complex use-cases, in the format + `=` or `=(, )`, e.g. + `foo=Field(datetime, default_factory=datetime.utcnow, alias='bar')` or + `foo=(str, FieldInfo(title='Foo'))` + """ + if __slots__ is not None: + # __slots__ will be ignored from here on + warnings.warn('__slots__ should not be passed to create_model', RuntimeWarning) + + if __base__ is not None: + if __config__ is not None: + raise ConfigError('to avoid confusion __config__ and __base__ cannot be used together') + if not isinstance(__base__, tuple): + __base__ = (__base__,) + else: + __base__ = (cast(Type['Model'], BaseModel),) + + __cls_kwargs__ = __cls_kwargs__ or {} + + fields = {} + annotations = {} + + for f_name, f_def in field_definitions.items(): + if not is_valid_field(f_name): + warnings.warn(f'fields may not start with an underscore, ignoring "{f_name}"', RuntimeWarning) + if isinstance(f_def, tuple): + try: + f_annotation, f_value = f_def + except ValueError as e: + raise ConfigError( + 'field definitions should either be a tuple of (, ) or just a ' + 'default value, unfortunately this means tuples as ' + 'default values are not allowed' + ) from e + else: + f_annotation, f_value = None, f_def + + if f_annotation: + annotations[f_name] = f_annotation + fields[f_name] = f_value + + namespace: 'DictStrAny' = {'__annotations__': annotations, '__module__': __module__} + if __validators__: + namespace.update(__validators__) + namespace.update(fields) + if __config__: + namespace['Config'] = inherit_config(__config__, BaseConfig) + resolved_bases = resolve_bases(__base__) + meta, ns, kwds = prepare_class(__model_name, resolved_bases, kwds=__cls_kwargs__) + if resolved_bases is not __base__: + ns['__orig_bases__'] = __base__ + namespace.update(ns) + return meta(__model_name, resolved_bases, namespace, **kwds) + + +_missing = object() + + +def validate_model( # noqa: C901 (ignore complexity) + model: Type[BaseModel], input_data: 'DictStrAny', cls: 'ModelOrDc' = None +) -> Tuple['DictStrAny', 'SetStr', Optional[ValidationError]]: + """ + validate data against a model. + """ + values = {} + errors = [] + # input_data names, possibly alias + names_used = set() + # field names, never aliases + fields_set = set() + config = model.__config__ + check_extra = config.extra is not Extra.ignore + cls_ = cls or model + + for validator in model.__pre_root_validators__: + try: + input_data = validator(cls_, input_data) + except (ValueError, TypeError, AssertionError) as exc: + return {}, set(), ValidationError([ErrorWrapper(exc, loc=ROOT_KEY)], cls_) + + for name, field in model.__fields__.items(): + value = input_data.get(field.alias, _missing) + using_name = False + if value is _missing and config.allow_population_by_field_name and field.alt_alias: + value = input_data.get(field.name, _missing) + using_name = True + + if value is _missing: + if field.required: + errors.append(ErrorWrapper(MissingError(), loc=field.alias)) + continue + + value = field.get_default() + + if not config.validate_all and not field.validate_always: + values[name] = value + continue + else: + fields_set.add(name) + if check_extra: + names_used.add(field.name if using_name else field.alias) + + v_, errors_ = field.validate(value, values, loc=field.alias, cls=cls_) + if isinstance(errors_, ErrorWrapper): + errors.append(errors_) + elif isinstance(errors_, list): + errors.extend(errors_) + else: + values[name] = v_ + + if check_extra: + if isinstance(input_data, GetterDict): + extra = input_data.extra_keys() - names_used + else: + extra = input_data.keys() - names_used + if extra: + fields_set |= extra + if config.extra is Extra.allow: + for f in extra: + values[f] = input_data[f] + else: + for f in sorted(extra): + errors.append(ErrorWrapper(ExtraError(), loc=f)) + + for skip_on_failure, validator in model.__post_root_validators__: + if skip_on_failure and errors: + continue + try: + values = validator(cls_, values) + except (ValueError, TypeError, AssertionError) as exc: + errors.append(ErrorWrapper(exc, loc=ROOT_KEY)) + + if errors: + return values, fields_set, ValidationError(errors, cls_) + else: + return values, fields_set, None diff --git a/lib/python3.12/site-packages/pydantic/v1/mypy.py b/lib/python3.12/site-packages/pydantic/v1/mypy.py new file mode 100644 index 0000000000000000000000000000000000000000..0a77569218ce6f832e7221b1f557d811318402e7 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/mypy.py @@ -0,0 +1,949 @@ +import sys +from configparser import ConfigParser +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type as TypingType, Union + +from mypy.errorcodes import ErrorCode +from mypy.nodes import ( + ARG_NAMED, + ARG_NAMED_OPT, + ARG_OPT, + ARG_POS, + ARG_STAR2, + MDEF, + Argument, + AssignmentStmt, + Block, + CallExpr, + ClassDef, + Context, + Decorator, + EllipsisExpr, + FuncBase, + FuncDef, + JsonDict, + MemberExpr, + NameExpr, + PassStmt, + PlaceholderNode, + RefExpr, + StrExpr, + SymbolNode, + SymbolTableNode, + TempNode, + TypeInfo, + TypeVarExpr, + Var, +) +from mypy.options import Options +from mypy.plugin import ( + CheckerPluginInterface, + ClassDefContext, + FunctionContext, + MethodContext, + Plugin, + ReportConfigContext, + SemanticAnalyzerPluginInterface, +) +from mypy.plugins import dataclasses +from mypy.semanal import set_callable_name # type: ignore +from mypy.server.trigger import make_wildcard_trigger +from mypy.types import ( + AnyType, + CallableType, + Instance, + NoneType, + Overloaded, + ProperType, + Type, + TypeOfAny, + TypeType, + TypeVarId, + TypeVarType, + UnionType, + get_proper_type, +) +from mypy.typevars import fill_typevars +from mypy.util import get_unique_redefinition_name +from mypy.version import __version__ as mypy_version + +from pydantic.v1.utils import is_valid_field + +try: + from mypy.types import TypeVarDef # type: ignore[attr-defined] +except ImportError: # pragma: no cover + # Backward-compatible with TypeVarDef from Mypy 0.910. + from mypy.types import TypeVarType as TypeVarDef + +CONFIGFILE_KEY = 'pydantic-mypy' +METADATA_KEY = 'pydantic-mypy-metadata' +_NAMESPACE = __name__[:-5] # 'pydantic' in 1.10.X, 'pydantic.v1' in v2.X +BASEMODEL_FULLNAME = f'{_NAMESPACE}.main.BaseModel' +BASESETTINGS_FULLNAME = f'{_NAMESPACE}.env_settings.BaseSettings' +MODEL_METACLASS_FULLNAME = f'{_NAMESPACE}.main.ModelMetaclass' +FIELD_FULLNAME = f'{_NAMESPACE}.fields.Field' +DATACLASS_FULLNAME = f'{_NAMESPACE}.dataclasses.dataclass' + + +def parse_mypy_version(version: str) -> Tuple[int, ...]: + return tuple(map(int, version.partition('+')[0].split('.'))) + + +MYPY_VERSION_TUPLE = parse_mypy_version(mypy_version) +BUILTINS_NAME = 'builtins' if MYPY_VERSION_TUPLE >= (0, 930) else '__builtins__' + +# Increment version if plugin changes and mypy caches should be invalidated +__version__ = 2 + + +def plugin(version: str) -> 'TypingType[Plugin]': + """ + `version` is the mypy version string + + We might want to use this to print a warning if the mypy version being used is + newer, or especially older, than we expect (or need). + """ + return PydanticPlugin + + +class PydanticPlugin(Plugin): + def __init__(self, options: Options) -> None: + self.plugin_config = PydanticPluginConfig(options) + self._plugin_data = self.plugin_config.to_data() + super().__init__(options) + + def get_base_class_hook(self, fullname: str) -> 'Optional[Callable[[ClassDefContext], None]]': + sym = self.lookup_fully_qualified(fullname) + if sym and isinstance(sym.node, TypeInfo): # pragma: no branch + # No branching may occur if the mypy cache has not been cleared + if any(get_fullname(base) == BASEMODEL_FULLNAME for base in sym.node.mro): + return self._pydantic_model_class_maker_callback + return None + + def get_metaclass_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]: + if fullname == MODEL_METACLASS_FULLNAME: + return self._pydantic_model_metaclass_marker_callback + return None + + def get_function_hook(self, fullname: str) -> 'Optional[Callable[[FunctionContext], Type]]': + sym = self.lookup_fully_qualified(fullname) + if sym and sym.fullname == FIELD_FULLNAME: + return self._pydantic_field_callback + return None + + def get_method_hook(self, fullname: str) -> Optional[Callable[[MethodContext], Type]]: + if fullname.endswith('.from_orm'): + return from_orm_callback + return None + + def get_class_decorator_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]: + """Mark pydantic.dataclasses as dataclass. + + Mypy version 1.1.1 added support for `@dataclass_transform` decorator. + """ + if fullname == DATACLASS_FULLNAME and MYPY_VERSION_TUPLE < (1, 1): + return dataclasses.dataclass_class_maker_callback # type: ignore[return-value] + return None + + def report_config_data(self, ctx: ReportConfigContext) -> Dict[str, Any]: + """Return all plugin config data. + + Used by mypy to determine if cache needs to be discarded. + """ + return self._plugin_data + + def _pydantic_model_class_maker_callback(self, ctx: ClassDefContext) -> None: + transformer = PydanticModelTransformer(ctx, self.plugin_config) + transformer.transform() + + def _pydantic_model_metaclass_marker_callback(self, ctx: ClassDefContext) -> None: + """Reset dataclass_transform_spec attribute of ModelMetaclass. + + Let the plugin handle it. This behavior can be disabled + if 'debug_dataclass_transform' is set to True', for testing purposes. + """ + if self.plugin_config.debug_dataclass_transform: + return + info_metaclass = ctx.cls.info.declared_metaclass + assert info_metaclass, "callback not passed from 'get_metaclass_hook'" + if getattr(info_metaclass.type, 'dataclass_transform_spec', None): + info_metaclass.type.dataclass_transform_spec = None # type: ignore[attr-defined] + + def _pydantic_field_callback(self, ctx: FunctionContext) -> 'Type': + """ + Extract the type of the `default` argument from the Field function, and use it as the return type. + + In particular: + * Check whether the default and default_factory argument is specified. + * Output an error if both are specified. + * Retrieve the type of the argument which is specified, and use it as return type for the function. + """ + default_any_type = ctx.default_return_type + + assert ctx.callee_arg_names[0] == 'default', '"default" is no longer first argument in Field()' + assert ctx.callee_arg_names[1] == 'default_factory', '"default_factory" is no longer second argument in Field()' + default_args = ctx.args[0] + default_factory_args = ctx.args[1] + + if default_args and default_factory_args: + error_default_and_default_factory_specified(ctx.api, ctx.context) + return default_any_type + + if default_args: + default_type = ctx.arg_types[0][0] + default_arg = default_args[0] + + # Fallback to default Any type if the field is required + if not isinstance(default_arg, EllipsisExpr): + return default_type + + elif default_factory_args: + default_factory_type = ctx.arg_types[1][0] + + # Functions which use `ParamSpec` can be overloaded, exposing the callable's types as a parameter + # Pydantic calls the default factory without any argument, so we retrieve the first item + if isinstance(default_factory_type, Overloaded): + if MYPY_VERSION_TUPLE > (0, 910): + default_factory_type = default_factory_type.items[0] + else: + # Mypy0.910 exposes the items of overloaded types in a function + default_factory_type = default_factory_type.items()[0] # type: ignore[operator] + + if isinstance(default_factory_type, CallableType): + ret_type = get_proper_type(default_factory_type.ret_type) + if ( + isinstance(ret_type, Instance) + and ret_type.args + and all(isinstance(arg, TypeVarType) for arg in ret_type.args) + ): + # Looks like the default factory is a type like `list` or `dict`, replace all args with `Any` + ret_type = ret_type.copy_modified(args=[default_any_type] * len(ret_type.args)) + return ret_type + + return default_any_type + + +class PydanticPluginConfig: + __slots__ = ( + 'init_forbid_extra', + 'init_typed', + 'warn_required_dynamic_aliases', + 'warn_untyped_fields', + 'debug_dataclass_transform', + ) + init_forbid_extra: bool + init_typed: bool + warn_required_dynamic_aliases: bool + warn_untyped_fields: bool + debug_dataclass_transform: bool # undocumented + + def __init__(self, options: Options) -> None: + if options.config_file is None: # pragma: no cover + return + + toml_config = parse_toml(options.config_file) + if toml_config is not None: + config = toml_config.get('tool', {}).get('pydantic-mypy', {}) + for key in self.__slots__: + setting = config.get(key, False) + if not isinstance(setting, bool): + raise ValueError(f'Configuration value must be a boolean for key: {key}') + setattr(self, key, setting) + else: + plugin_config = ConfigParser() + plugin_config.read(options.config_file) + for key in self.__slots__: + setting = plugin_config.getboolean(CONFIGFILE_KEY, key, fallback=False) + setattr(self, key, setting) + + def to_data(self) -> Dict[str, Any]: + return {key: getattr(self, key) for key in self.__slots__} + + +def from_orm_callback(ctx: MethodContext) -> Type: + """ + Raise an error if orm_mode is not enabled + """ + model_type: Instance + ctx_type = ctx.type + if isinstance(ctx_type, TypeType): + ctx_type = ctx_type.item + if isinstance(ctx_type, CallableType) and isinstance(ctx_type.ret_type, Instance): + model_type = ctx_type.ret_type # called on the class + elif isinstance(ctx_type, Instance): + model_type = ctx_type # called on an instance (unusual, but still valid) + else: # pragma: no cover + detail = f'ctx.type: {ctx_type} (of type {ctx_type.__class__.__name__})' + error_unexpected_behavior(detail, ctx.api, ctx.context) + return ctx.default_return_type + pydantic_metadata = model_type.type.metadata.get(METADATA_KEY) + if pydantic_metadata is None: + return ctx.default_return_type + orm_mode = pydantic_metadata.get('config', {}).get('orm_mode') + if orm_mode is not True: + error_from_orm(get_name(model_type.type), ctx.api, ctx.context) + return ctx.default_return_type + + +class PydanticModelTransformer: + tracked_config_fields: Set[str] = { + 'extra', + 'allow_mutation', + 'frozen', + 'orm_mode', + 'allow_population_by_field_name', + 'alias_generator', + } + + def __init__(self, ctx: ClassDefContext, plugin_config: PydanticPluginConfig) -> None: + self._ctx = ctx + self.plugin_config = plugin_config + + def transform(self) -> None: + """ + Configures the BaseModel subclass according to the plugin settings. + + In particular: + * determines the model config and fields, + * adds a fields-aware signature for the initializer and construct methods + * freezes the class if allow_mutation = False or frozen = True + * stores the fields, config, and if the class is settings in the mypy metadata for access by subclasses + """ + ctx = self._ctx + info = ctx.cls.info + + self.adjust_validator_signatures() + config = self.collect_config() + fields = self.collect_fields(config) + is_settings = any(get_fullname(base) == BASESETTINGS_FULLNAME for base in info.mro[:-1]) + self.add_initializer(fields, config, is_settings) + self.add_construct_method(fields) + self.set_frozen(fields, frozen=config.allow_mutation is False or config.frozen is True) + info.metadata[METADATA_KEY] = { + 'fields': {field.name: field.serialize() for field in fields}, + 'config': config.set_values_dict(), + } + + def adjust_validator_signatures(self) -> None: + """When we decorate a function `f` with `pydantic.validator(...), mypy sees + `f` as a regular method taking a `self` instance, even though pydantic + internally wraps `f` with `classmethod` if necessary. + + Teach mypy this by marking any function whose outermost decorator is a + `validator()` call as a classmethod. + """ + for name, sym in self._ctx.cls.info.names.items(): + if isinstance(sym.node, Decorator): + first_dec = sym.node.original_decorators[0] + if ( + isinstance(first_dec, CallExpr) + and isinstance(first_dec.callee, NameExpr) + and first_dec.callee.fullname == f'{_NAMESPACE}.class_validators.validator' + ): + sym.node.func.is_class = True + + def collect_config(self) -> 'ModelConfigData': + """ + Collects the values of the config attributes that are used by the plugin, accounting for parent classes. + """ + ctx = self._ctx + cls = ctx.cls + config = ModelConfigData() + for stmt in cls.defs.body: + if not isinstance(stmt, ClassDef): + continue + if stmt.name == 'Config': + for substmt in stmt.defs.body: + if not isinstance(substmt, AssignmentStmt): + continue + config.update(self.get_config_update(substmt)) + if ( + config.has_alias_generator + and not config.allow_population_by_field_name + and self.plugin_config.warn_required_dynamic_aliases + ): + error_required_dynamic_aliases(ctx.api, stmt) + for info in cls.info.mro[1:]: # 0 is the current class + if METADATA_KEY not in info.metadata: + continue + + # Each class depends on the set of fields in its ancestors + ctx.api.add_plugin_dependency(make_wildcard_trigger(get_fullname(info))) + for name, value in info.metadata[METADATA_KEY]['config'].items(): + config.setdefault(name, value) + return config + + def collect_fields(self, model_config: 'ModelConfigData') -> List['PydanticModelField']: + """ + Collects the fields for the model, accounting for parent classes + """ + # First, collect fields belonging to the current class. + ctx = self._ctx + cls = self._ctx.cls + fields = [] # type: List[PydanticModelField] + known_fields = set() # type: Set[str] + for stmt in cls.defs.body: + if not isinstance(stmt, AssignmentStmt): # `and stmt.new_syntax` to require annotation + continue + + lhs = stmt.lvalues[0] + if not isinstance(lhs, NameExpr) or not is_valid_field(lhs.name): + continue + + if not stmt.new_syntax and self.plugin_config.warn_untyped_fields: + error_untyped_fields(ctx.api, stmt) + + # if lhs.name == '__config__': # BaseConfig not well handled; I'm not sure why yet + # continue + + sym = cls.info.names.get(lhs.name) + if sym is None: # pragma: no cover + # This is likely due to a star import (see the dataclasses plugin for a more detailed explanation) + # This is the same logic used in the dataclasses plugin + continue + + node = sym.node + if isinstance(node, PlaceholderNode): # pragma: no cover + # See the PlaceholderNode docstring for more detail about how this can occur + # Basically, it is an edge case when dealing with complex import logic + # This is the same logic used in the dataclasses plugin + continue + if not isinstance(node, Var): # pragma: no cover + # Don't know if this edge case still happens with the `is_valid_field` check above + # but better safe than sorry + continue + + # x: ClassVar[int] is ignored by dataclasses. + if node.is_classvar: + continue + + is_required = self.get_is_required(cls, stmt, lhs) + alias, has_dynamic_alias = self.get_alias_info(stmt) + if ( + has_dynamic_alias + and not model_config.allow_population_by_field_name + and self.plugin_config.warn_required_dynamic_aliases + ): + error_required_dynamic_aliases(ctx.api, stmt) + fields.append( + PydanticModelField( + name=lhs.name, + is_required=is_required, + alias=alias, + has_dynamic_alias=has_dynamic_alias, + line=stmt.line, + column=stmt.column, + ) + ) + known_fields.add(lhs.name) + all_fields = fields.copy() + for info in cls.info.mro[1:]: # 0 is the current class, -2 is BaseModel, -1 is object + if METADATA_KEY not in info.metadata: + continue + + superclass_fields = [] + # Each class depends on the set of fields in its ancestors + ctx.api.add_plugin_dependency(make_wildcard_trigger(get_fullname(info))) + + for name, data in info.metadata[METADATA_KEY]['fields'].items(): + if name not in known_fields: + field = PydanticModelField.deserialize(info, data) + known_fields.add(name) + superclass_fields.append(field) + else: + (field,) = (a for a in all_fields if a.name == name) + all_fields.remove(field) + superclass_fields.append(field) + all_fields = superclass_fields + all_fields + return all_fields + + def add_initializer(self, fields: List['PydanticModelField'], config: 'ModelConfigData', is_settings: bool) -> None: + """ + Adds a fields-aware `__init__` method to the class. + + The added `__init__` will be annotated with types vs. all `Any` depending on the plugin settings. + """ + ctx = self._ctx + typed = self.plugin_config.init_typed + use_alias = config.allow_population_by_field_name is not True + force_all_optional = is_settings or bool( + config.has_alias_generator and not config.allow_population_by_field_name + ) + init_arguments = self.get_field_arguments( + fields, typed=typed, force_all_optional=force_all_optional, use_alias=use_alias + ) + if not self.should_init_forbid_extra(fields, config): + var = Var('kwargs') + init_arguments.append(Argument(var, AnyType(TypeOfAny.explicit), None, ARG_STAR2)) + + if '__init__' not in ctx.cls.info.names: + add_method(ctx, '__init__', init_arguments, NoneType()) + + def add_construct_method(self, fields: List['PydanticModelField']) -> None: + """ + Adds a fully typed `construct` classmethod to the class. + + Similar to the fields-aware __init__ method, but always uses the field names (not aliases), + and does not treat settings fields as optional. + """ + ctx = self._ctx + set_str = ctx.api.named_type(f'{BUILTINS_NAME}.set', [ctx.api.named_type(f'{BUILTINS_NAME}.str')]) + optional_set_str = UnionType([set_str, NoneType()]) + fields_set_argument = Argument(Var('_fields_set', optional_set_str), optional_set_str, None, ARG_OPT) + construct_arguments = self.get_field_arguments(fields, typed=True, force_all_optional=False, use_alias=False) + construct_arguments = [fields_set_argument] + construct_arguments + + obj_type = ctx.api.named_type(f'{BUILTINS_NAME}.object') + self_tvar_name = '_PydanticBaseModel' # Make sure it does not conflict with other names in the class + tvar_fullname = ctx.cls.fullname + '.' + self_tvar_name + if MYPY_VERSION_TUPLE >= (1, 4): + tvd = TypeVarType( + self_tvar_name, + tvar_fullname, + ( + TypeVarId(-1, namespace=ctx.cls.fullname + '.construct') + if MYPY_VERSION_TUPLE >= (1, 11) + else TypeVarId(-1) + ), + [], + obj_type, + AnyType(TypeOfAny.from_omitted_generics), # type: ignore[arg-type] + ) + self_tvar_expr = TypeVarExpr( + self_tvar_name, + tvar_fullname, + [], + obj_type, + AnyType(TypeOfAny.from_omitted_generics), # type: ignore[arg-type] + ) + else: + tvd = TypeVarDef(self_tvar_name, tvar_fullname, -1, [], obj_type) + self_tvar_expr = TypeVarExpr(self_tvar_name, tvar_fullname, [], obj_type) + ctx.cls.info.names[self_tvar_name] = SymbolTableNode(MDEF, self_tvar_expr) + + # Backward-compatible with TypeVarDef from Mypy 0.910. + if isinstance(tvd, TypeVarType): + self_type = tvd + else: + self_type = TypeVarType(tvd) + + add_method( + ctx, + 'construct', + construct_arguments, + return_type=self_type, + self_type=self_type, + tvar_def=tvd, + is_classmethod=True, + ) + + def set_frozen(self, fields: List['PydanticModelField'], frozen: bool) -> None: + """ + Marks all fields as properties so that attempts to set them trigger mypy errors. + + This is the same approach used by the attrs and dataclasses plugins. + """ + ctx = self._ctx + info = ctx.cls.info + for field in fields: + sym_node = info.names.get(field.name) + if sym_node is not None: + var = sym_node.node + if isinstance(var, Var): + var.is_property = frozen + elif isinstance(var, PlaceholderNode) and not ctx.api.final_iteration: + # See https://github.com/pydantic/pydantic/issues/5191 to hit this branch for test coverage + ctx.api.defer() + else: # pragma: no cover + # I don't know whether it's possible to hit this branch, but I've added it for safety + try: + var_str = str(var) + except TypeError: + # This happens for PlaceholderNode; perhaps it will happen for other types in the future.. + var_str = repr(var) + detail = f'sym_node.node: {var_str} (of type {var.__class__})' + error_unexpected_behavior(detail, ctx.api, ctx.cls) + else: + var = field.to_var(info, use_alias=False) + var.info = info + var.is_property = frozen + var._fullname = get_fullname(info) + '.' + get_name(var) + info.names[get_name(var)] = SymbolTableNode(MDEF, var) + + def get_config_update(self, substmt: AssignmentStmt) -> Optional['ModelConfigData']: + """ + Determines the config update due to a single statement in the Config class definition. + + Warns if a tracked config attribute is set to a value the plugin doesn't know how to interpret (e.g., an int) + """ + lhs = substmt.lvalues[0] + if not (isinstance(lhs, NameExpr) and lhs.name in self.tracked_config_fields): + return None + if lhs.name == 'extra': + if isinstance(substmt.rvalue, StrExpr): + forbid_extra = substmt.rvalue.value == 'forbid' + elif isinstance(substmt.rvalue, MemberExpr): + forbid_extra = substmt.rvalue.name == 'forbid' + else: + error_invalid_config_value(lhs.name, self._ctx.api, substmt) + return None + return ModelConfigData(forbid_extra=forbid_extra) + if lhs.name == 'alias_generator': + has_alias_generator = True + if isinstance(substmt.rvalue, NameExpr) and substmt.rvalue.fullname == 'builtins.None': + has_alias_generator = False + return ModelConfigData(has_alias_generator=has_alias_generator) + if isinstance(substmt.rvalue, NameExpr) and substmt.rvalue.fullname in ('builtins.True', 'builtins.False'): + return ModelConfigData(**{lhs.name: substmt.rvalue.fullname == 'builtins.True'}) + error_invalid_config_value(lhs.name, self._ctx.api, substmt) + return None + + @staticmethod + def get_is_required(cls: ClassDef, stmt: AssignmentStmt, lhs: NameExpr) -> bool: + """ + Returns a boolean indicating whether the field defined in `stmt` is a required field. + """ + expr = stmt.rvalue + if isinstance(expr, TempNode): + # TempNode means annotation-only, so only non-required if Optional + value_type = get_proper_type(cls.info[lhs.name].type) + return not PydanticModelTransformer.type_has_implicit_default(value_type) + if isinstance(expr, CallExpr) and isinstance(expr.callee, RefExpr) and expr.callee.fullname == FIELD_FULLNAME: + # The "default value" is a call to `Field`; at this point, the field is + # only required if default is Ellipsis (i.e., `field_name: Annotation = Field(...)`) or if default_factory + # is specified. + for arg, name in zip(expr.args, expr.arg_names): + # If name is None, then this arg is the default because it is the only positional argument. + if name is None or name == 'default': + return arg.__class__ is EllipsisExpr + if name == 'default_factory': + return False + # In this case, default and default_factory are not specified, so we need to look at the annotation + value_type = get_proper_type(cls.info[lhs.name].type) + return not PydanticModelTransformer.type_has_implicit_default(value_type) + # Only required if the "default value" is Ellipsis (i.e., `field_name: Annotation = ...`) + return isinstance(expr, EllipsisExpr) + + @staticmethod + def type_has_implicit_default(type_: Optional[ProperType]) -> bool: + """ + Returns True if the passed type will be given an implicit default value. + + In pydantic v1, this is the case for Optional types and Any (with default value None). + """ + if isinstance(type_, AnyType): + # Annotated as Any + return True + if isinstance(type_, UnionType) and any( + isinstance(item, NoneType) or isinstance(item, AnyType) for item in type_.items + ): + # Annotated as Optional, or otherwise having NoneType or AnyType in the union + return True + return False + + @staticmethod + def get_alias_info(stmt: AssignmentStmt) -> Tuple[Optional[str], bool]: + """ + Returns a pair (alias, has_dynamic_alias), extracted from the declaration of the field defined in `stmt`. + + `has_dynamic_alias` is True if and only if an alias is provided, but not as a string literal. + If `has_dynamic_alias` is True, `alias` will be None. + """ + expr = stmt.rvalue + if isinstance(expr, TempNode): + # TempNode means annotation-only + return None, False + + if not ( + isinstance(expr, CallExpr) and isinstance(expr.callee, RefExpr) and expr.callee.fullname == FIELD_FULLNAME + ): + # Assigned value is not a call to pydantic.fields.Field + return None, False + + for i, arg_name in enumerate(expr.arg_names): + if arg_name != 'alias': + continue + arg = expr.args[i] + if isinstance(arg, StrExpr): + return arg.value, False + else: + return None, True + return None, False + + def get_field_arguments( + self, fields: List['PydanticModelField'], typed: bool, force_all_optional: bool, use_alias: bool + ) -> List[Argument]: + """ + Helper function used during the construction of the `__init__` and `construct` method signatures. + + Returns a list of mypy Argument instances for use in the generated signatures. + """ + info = self._ctx.cls.info + arguments = [ + field.to_argument(info, typed=typed, force_optional=force_all_optional, use_alias=use_alias) + for field in fields + if not (use_alias and field.has_dynamic_alias) + ] + return arguments + + def should_init_forbid_extra(self, fields: List['PydanticModelField'], config: 'ModelConfigData') -> bool: + """ + Indicates whether the generated `__init__` should get a `**kwargs` at the end of its signature + + We disallow arbitrary kwargs if the extra config setting is "forbid", or if the plugin config says to, + *unless* a required dynamic alias is present (since then we can't determine a valid signature). + """ + if not config.allow_population_by_field_name: + if self.is_dynamic_alias_present(fields, bool(config.has_alias_generator)): + return False + if config.forbid_extra: + return True + return self.plugin_config.init_forbid_extra + + @staticmethod + def is_dynamic_alias_present(fields: List['PydanticModelField'], has_alias_generator: bool) -> bool: + """ + Returns whether any fields on the model have a "dynamic alias", i.e., an alias that cannot be + determined during static analysis. + """ + for field in fields: + if field.has_dynamic_alias: + return True + if has_alias_generator: + for field in fields: + if field.alias is None: + return True + return False + + +class PydanticModelField: + def __init__( + self, name: str, is_required: bool, alias: Optional[str], has_dynamic_alias: bool, line: int, column: int + ): + self.name = name + self.is_required = is_required + self.alias = alias + self.has_dynamic_alias = has_dynamic_alias + self.line = line + self.column = column + + def to_var(self, info: TypeInfo, use_alias: bool) -> Var: + name = self.name + if use_alias and self.alias is not None: + name = self.alias + return Var(name, info[self.name].type) + + def to_argument(self, info: TypeInfo, typed: bool, force_optional: bool, use_alias: bool) -> Argument: + if typed and info[self.name].type is not None: + type_annotation = info[self.name].type + else: + type_annotation = AnyType(TypeOfAny.explicit) + return Argument( + variable=self.to_var(info, use_alias), + type_annotation=type_annotation, + initializer=None, + kind=ARG_NAMED_OPT if force_optional or not self.is_required else ARG_NAMED, + ) + + def serialize(self) -> JsonDict: + return self.__dict__ + + @classmethod + def deserialize(cls, info: TypeInfo, data: JsonDict) -> 'PydanticModelField': + return cls(**data) + + +class ModelConfigData: + def __init__( + self, + forbid_extra: Optional[bool] = None, + allow_mutation: Optional[bool] = None, + frozen: Optional[bool] = None, + orm_mode: Optional[bool] = None, + allow_population_by_field_name: Optional[bool] = None, + has_alias_generator: Optional[bool] = None, + ): + self.forbid_extra = forbid_extra + self.allow_mutation = allow_mutation + self.frozen = frozen + self.orm_mode = orm_mode + self.allow_population_by_field_name = allow_population_by_field_name + self.has_alias_generator = has_alias_generator + + def set_values_dict(self) -> Dict[str, Any]: + return {k: v for k, v in self.__dict__.items() if v is not None} + + def update(self, config: Optional['ModelConfigData']) -> None: + if config is None: + return + for k, v in config.set_values_dict().items(): + setattr(self, k, v) + + def setdefault(self, key: str, value: Any) -> None: + if getattr(self, key) is None: + setattr(self, key, value) + + +ERROR_ORM = ErrorCode('pydantic-orm', 'Invalid from_orm call', 'Pydantic') +ERROR_CONFIG = ErrorCode('pydantic-config', 'Invalid config value', 'Pydantic') +ERROR_ALIAS = ErrorCode('pydantic-alias', 'Dynamic alias disallowed', 'Pydantic') +ERROR_UNEXPECTED = ErrorCode('pydantic-unexpected', 'Unexpected behavior', 'Pydantic') +ERROR_UNTYPED = ErrorCode('pydantic-field', 'Untyped field disallowed', 'Pydantic') +ERROR_FIELD_DEFAULTS = ErrorCode('pydantic-field', 'Invalid Field defaults', 'Pydantic') + + +def error_from_orm(model_name: str, api: CheckerPluginInterface, context: Context) -> None: + api.fail(f'"{model_name}" does not have orm_mode=True', context, code=ERROR_ORM) + + +def error_invalid_config_value(name: str, api: SemanticAnalyzerPluginInterface, context: Context) -> None: + api.fail(f'Invalid value for "Config.{name}"', context, code=ERROR_CONFIG) + + +def error_required_dynamic_aliases(api: SemanticAnalyzerPluginInterface, context: Context) -> None: + api.fail('Required dynamic aliases disallowed', context, code=ERROR_ALIAS) + + +def error_unexpected_behavior( + detail: str, api: Union[CheckerPluginInterface, SemanticAnalyzerPluginInterface], context: Context +) -> None: # pragma: no cover + # Can't think of a good way to test this, but I confirmed it renders as desired by adding to a non-error path + link = 'https://github.com/pydantic/pydantic/issues/new/choose' + full_message = f'The pydantic mypy plugin ran into unexpected behavior: {detail}\n' + full_message += f'Please consider reporting this bug at {link} so we can try to fix it!' + api.fail(full_message, context, code=ERROR_UNEXPECTED) + + +def error_untyped_fields(api: SemanticAnalyzerPluginInterface, context: Context) -> None: + api.fail('Untyped fields disallowed', context, code=ERROR_UNTYPED) + + +def error_default_and_default_factory_specified(api: CheckerPluginInterface, context: Context) -> None: + api.fail('Field default and default_factory cannot be specified together', context, code=ERROR_FIELD_DEFAULTS) + + +def add_method( + ctx: ClassDefContext, + name: str, + args: List[Argument], + return_type: Type, + self_type: Optional[Type] = None, + tvar_def: Optional[TypeVarDef] = None, + is_classmethod: bool = False, + is_new: bool = False, + # is_staticmethod: bool = False, +) -> None: + """ + Adds a new method to a class. + + This can be dropped if/when https://github.com/python/mypy/issues/7301 is merged + """ + info = ctx.cls.info + + # First remove any previously generated methods with the same name + # to avoid clashes and problems in the semantic analyzer. + if name in info.names: + sym = info.names[name] + if sym.plugin_generated and isinstance(sym.node, FuncDef): + ctx.cls.defs.body.remove(sym.node) # pragma: no cover + + self_type = self_type or fill_typevars(info) + if is_classmethod or is_new: + first = [Argument(Var('_cls'), TypeType.make_normalized(self_type), None, ARG_POS)] + # elif is_staticmethod: + # first = [] + else: + self_type = self_type or fill_typevars(info) + first = [Argument(Var('__pydantic_self__'), self_type, None, ARG_POS)] + args = first + args + arg_types, arg_names, arg_kinds = [], [], [] + for arg in args: + assert arg.type_annotation, 'All arguments must be fully typed.' + arg_types.append(arg.type_annotation) + arg_names.append(get_name(arg.variable)) + arg_kinds.append(arg.kind) + + function_type = ctx.api.named_type(f'{BUILTINS_NAME}.function') + signature = CallableType( + arg_types, arg_kinds, arg_names, return_type, function_type, variables=[tvar_def] if tvar_def else None + ) + + func = FuncDef(name, args, Block([PassStmt()])) + func.info = info + func.type = set_callable_name(signature, func) + func.is_class = is_classmethod + # func.is_static = is_staticmethod + func._fullname = get_fullname(info) + '.' + name + func.line = info.line + + # NOTE: we would like the plugin generated node to dominate, but we still + # need to keep any existing definitions so they get semantically analyzed. + if name in info.names: + # Get a nice unique name instead. + r_name = get_unique_redefinition_name(name, info.names) + info.names[r_name] = info.names[name] + + if is_classmethod: # or is_staticmethod: + func.is_decorated = True + v = Var(name, func.type) + v.info = info + v._fullname = func._fullname + # if is_classmethod: + v.is_classmethod = True + dec = Decorator(func, [NameExpr('classmethod')], v) + # else: + # v.is_staticmethod = True + # dec = Decorator(func, [NameExpr('staticmethod')], v) + + dec.line = info.line + sym = SymbolTableNode(MDEF, dec) + else: + sym = SymbolTableNode(MDEF, func) + sym.plugin_generated = True + + info.names[name] = sym + info.defn.defs.body.append(func) + + +def get_fullname(x: Union[FuncBase, SymbolNode]) -> str: + """ + Used for compatibility with mypy 0.740; can be dropped once support for 0.740 is dropped. + """ + fn = x.fullname + if callable(fn): # pragma: no cover + return fn() + return fn + + +def get_name(x: Union[FuncBase, SymbolNode]) -> str: + """ + Used for compatibility with mypy 0.740; can be dropped once support for 0.740 is dropped. + """ + fn = x.name + if callable(fn): # pragma: no cover + return fn() + return fn + + +def parse_toml(config_file: str) -> Optional[Dict[str, Any]]: + if not config_file.endswith('.toml'): + return None + + read_mode = 'rb' + if sys.version_info >= (3, 11): + import tomllib as toml_ + else: + try: + import tomli as toml_ + except ImportError: + # older versions of mypy have toml as a dependency, not tomli + read_mode = 'r' + try: + import toml as toml_ # type: ignore[no-redef] + except ImportError: # pragma: no cover + import warnings + + warnings.warn('No TOML parser installed, cannot read configuration from `pyproject.toml`.') + return None + + with open(config_file, read_mode) as rf: + return toml_.load(rf) # type: ignore[arg-type] diff --git a/lib/python3.12/site-packages/pydantic/v1/networks.py b/lib/python3.12/site-packages/pydantic/v1/networks.py new file mode 100644 index 0000000000000000000000000000000000000000..ba07b74867b83beaaabc3afe2e19d77d71544338 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/networks.py @@ -0,0 +1,747 @@ +import re +from ipaddress import ( + IPv4Address, + IPv4Interface, + IPv4Network, + IPv6Address, + IPv6Interface, + IPv6Network, + _BaseAddress, + _BaseNetwork, +) +from typing import ( + TYPE_CHECKING, + Any, + Collection, + Dict, + Generator, + List, + Match, + Optional, + Pattern, + Set, + Tuple, + Type, + Union, + cast, + no_type_check, +) + +from pydantic.v1 import errors +from pydantic.v1.utils import Representation, update_not_none +from pydantic.v1.validators import constr_length_validator, str_validator + +if TYPE_CHECKING: + import email_validator + from typing_extensions import TypedDict + + from pydantic.v1.config import BaseConfig + from pydantic.v1.fields import ModelField + from pydantic.v1.typing import AnyCallable + + CallableGenerator = Generator[AnyCallable, None, None] + + class Parts(TypedDict, total=False): + scheme: str + user: Optional[str] + password: Optional[str] + ipv4: Optional[str] + ipv6: Optional[str] + domain: Optional[str] + port: Optional[str] + path: Optional[str] + query: Optional[str] + fragment: Optional[str] + + class HostParts(TypedDict, total=False): + host: str + tld: Optional[str] + host_type: Optional[str] + port: Optional[str] + rebuild: bool + +else: + email_validator = None + + class Parts(dict): + pass + + +NetworkType = Union[str, bytes, int, Tuple[Union[str, bytes, int], Union[str, int]]] + +__all__ = [ + 'AnyUrl', + 'AnyHttpUrl', + 'FileUrl', + 'HttpUrl', + 'stricturl', + 'EmailStr', + 'NameEmail', + 'IPvAnyAddress', + 'IPvAnyInterface', + 'IPvAnyNetwork', + 'PostgresDsn', + 'CockroachDsn', + 'AmqpDsn', + 'RedisDsn', + 'MongoDsn', + 'KafkaDsn', + 'validate_email', +] + +_url_regex_cache = None +_multi_host_url_regex_cache = None +_ascii_domain_regex_cache = None +_int_domain_regex_cache = None +_host_regex_cache = None + +_host_regex = ( + r'(?:' + r'(?P(?:\d{1,3}\.){3}\d{1,3})(?=$|[/:#?])|' # ipv4 + r'(?P\[[A-F0-9]*:[A-F0-9:]+\])(?=$|[/:#?])|' # ipv6 + r'(?P[^\s/:?#]+)' # domain, validation occurs later + r')?' + r'(?::(?P\d+))?' # port +) +_scheme_regex = r'(?:(?P[a-z][a-z0-9+\-.]+)://)?' # scheme https://tools.ietf.org/html/rfc3986#appendix-A +_user_info_regex = r'(?:(?P[^\s:/]*)(?::(?P[^\s/]*))?@)?' +_path_regex = r'(?P/[^\s?#]*)?' +_query_regex = r'(?:\?(?P[^\s#]*))?' +_fragment_regex = r'(?:#(?P[^\s#]*))?' + + +def url_regex() -> Pattern[str]: + global _url_regex_cache + if _url_regex_cache is None: + _url_regex_cache = re.compile( + rf'{_scheme_regex}{_user_info_regex}{_host_regex}{_path_regex}{_query_regex}{_fragment_regex}', + re.IGNORECASE, + ) + return _url_regex_cache + + +def multi_host_url_regex() -> Pattern[str]: + """ + Compiled multi host url regex. + + Additionally to `url_regex` it allows to match multiple hosts. + E.g. host1.db.net,host2.db.net + """ + global _multi_host_url_regex_cache + if _multi_host_url_regex_cache is None: + _multi_host_url_regex_cache = re.compile( + rf'{_scheme_regex}{_user_info_regex}' + r'(?P([^/]*))' # validation occurs later + rf'{_path_regex}{_query_regex}{_fragment_regex}', + re.IGNORECASE, + ) + return _multi_host_url_regex_cache + + +def ascii_domain_regex() -> Pattern[str]: + global _ascii_domain_regex_cache + if _ascii_domain_regex_cache is None: + ascii_chunk = r'[_0-9a-z](?:[-_0-9a-z]{0,61}[_0-9a-z])?' + ascii_domain_ending = r'(?P\.[a-z]{2,63})?\.?' + _ascii_domain_regex_cache = re.compile( + fr'(?:{ascii_chunk}\.)*?{ascii_chunk}{ascii_domain_ending}', re.IGNORECASE + ) + return _ascii_domain_regex_cache + + +def int_domain_regex() -> Pattern[str]: + global _int_domain_regex_cache + if _int_domain_regex_cache is None: + int_chunk = r'[_0-9a-\U00040000](?:[-_0-9a-\U00040000]{0,61}[_0-9a-\U00040000])?' + int_domain_ending = r'(?P(\.[^\W\d_]{2,63})|(\.(?:xn--)[_0-9a-z-]{2,63}))?\.?' + _int_domain_regex_cache = re.compile(fr'(?:{int_chunk}\.)*?{int_chunk}{int_domain_ending}', re.IGNORECASE) + return _int_domain_regex_cache + + +def host_regex() -> Pattern[str]: + global _host_regex_cache + if _host_regex_cache is None: + _host_regex_cache = re.compile( + _host_regex, + re.IGNORECASE, + ) + return _host_regex_cache + + +class AnyUrl(str): + strip_whitespace = True + min_length = 1 + max_length = 2**16 + allowed_schemes: Optional[Collection[str]] = None + tld_required: bool = False + user_required: bool = False + host_required: bool = True + hidden_parts: Set[str] = set() + + __slots__ = ('scheme', 'user', 'password', 'host', 'tld', 'host_type', 'port', 'path', 'query', 'fragment') + + @no_type_check + def __new__(cls, url: Optional[str], **kwargs) -> object: + return str.__new__(cls, cls.build(**kwargs) if url is None else url) + + def __init__( + self, + url: str, + *, + scheme: str, + user: Optional[str] = None, + password: Optional[str] = None, + host: Optional[str] = None, + tld: Optional[str] = None, + host_type: str = 'domain', + port: Optional[str] = None, + path: Optional[str] = None, + query: Optional[str] = None, + fragment: Optional[str] = None, + ) -> None: + str.__init__(url) + self.scheme = scheme + self.user = user + self.password = password + self.host = host + self.tld = tld + self.host_type = host_type + self.port = port + self.path = path + self.query = query + self.fragment = fragment + + @classmethod + def build( + cls, + *, + scheme: str, + user: Optional[str] = None, + password: Optional[str] = None, + host: str, + port: Optional[str] = None, + path: Optional[str] = None, + query: Optional[str] = None, + fragment: Optional[str] = None, + **_kwargs: str, + ) -> str: + parts = Parts( + scheme=scheme, + user=user, + password=password, + host=host, + port=port, + path=path, + query=query, + fragment=fragment, + **_kwargs, # type: ignore[misc] + ) + + url = scheme + '://' + if user: + url += user + if password: + url += ':' + password + if user or password: + url += '@' + url += host + if port and ('port' not in cls.hidden_parts or cls.get_default_parts(parts).get('port') != port): + url += ':' + port + if path: + url += path + if query: + url += '?' + query + if fragment: + url += '#' + fragment + return url + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none(field_schema, minLength=cls.min_length, maxLength=cls.max_length, format='uri') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + + @classmethod + def validate(cls, value: Any, field: 'ModelField', config: 'BaseConfig') -> 'AnyUrl': + if value.__class__ == cls: + return value + value = str_validator(value) + if cls.strip_whitespace: + value = value.strip() + url: str = cast(str, constr_length_validator(value, field, config)) + + m = cls._match_url(url) + # the regex should always match, if it doesn't please report with details of the URL tried + assert m, 'URL regex failed unexpectedly' + + original_parts = cast('Parts', m.groupdict()) + parts = cls.apply_default_parts(original_parts) + parts = cls.validate_parts(parts) + + if m.end() != len(url): + raise errors.UrlExtraError(extra=url[m.end() :]) + + return cls._build_url(m, url, parts) + + @classmethod + def _build_url(cls, m: Match[str], url: str, parts: 'Parts') -> 'AnyUrl': + """ + Validate hosts and build the AnyUrl object. Split from `validate` so this method + can be altered in `MultiHostDsn`. + """ + host, tld, host_type, rebuild = cls.validate_host(parts) + + return cls( + None if rebuild else url, + scheme=parts['scheme'], + user=parts['user'], + password=parts['password'], + host=host, + tld=tld, + host_type=host_type, + port=parts['port'], + path=parts['path'], + query=parts['query'], + fragment=parts['fragment'], + ) + + @staticmethod + def _match_url(url: str) -> Optional[Match[str]]: + return url_regex().match(url) + + @staticmethod + def _validate_port(port: Optional[str]) -> None: + if port is not None and int(port) > 65_535: + raise errors.UrlPortError() + + @classmethod + def validate_parts(cls, parts: 'Parts', validate_port: bool = True) -> 'Parts': + """ + A method used to validate parts of a URL. + Could be overridden to set default values for parts if missing + """ + scheme = parts['scheme'] + if scheme is None: + raise errors.UrlSchemeError() + + if cls.allowed_schemes and scheme.lower() not in cls.allowed_schemes: + raise errors.UrlSchemePermittedError(set(cls.allowed_schemes)) + + if validate_port: + cls._validate_port(parts['port']) + + user = parts['user'] + if cls.user_required and user is None: + raise errors.UrlUserInfoError() + + return parts + + @classmethod + def validate_host(cls, parts: 'Parts') -> Tuple[str, Optional[str], str, bool]: + tld, host_type, rebuild = None, None, False + for f in ('domain', 'ipv4', 'ipv6'): + host = parts[f] # type: ignore[literal-required] + if host: + host_type = f + break + + if host is None: + if cls.host_required: + raise errors.UrlHostError() + elif host_type == 'domain': + is_international = False + d = ascii_domain_regex().fullmatch(host) + if d is None: + d = int_domain_regex().fullmatch(host) + if d is None: + raise errors.UrlHostError() + is_international = True + + tld = d.group('tld') + if tld is None and not is_international: + d = int_domain_regex().fullmatch(host) + assert d is not None + tld = d.group('tld') + is_international = True + + if tld is not None: + tld = tld[1:] + elif cls.tld_required: + raise errors.UrlHostTldError() + + if is_international: + host_type = 'int_domain' + rebuild = True + host = host.encode('idna').decode('ascii') + if tld is not None: + tld = tld.encode('idna').decode('ascii') + + return host, tld, host_type, rebuild # type: ignore + + @staticmethod + def get_default_parts(parts: 'Parts') -> 'Parts': + return {} + + @classmethod + def apply_default_parts(cls, parts: 'Parts') -> 'Parts': + for key, value in cls.get_default_parts(parts).items(): + if not parts[key]: # type: ignore[literal-required] + parts[key] = value # type: ignore[literal-required] + return parts + + def __repr__(self) -> str: + extra = ', '.join(f'{n}={getattr(self, n)!r}' for n in self.__slots__ if getattr(self, n) is not None) + return f'{self.__class__.__name__}({super().__repr__()}, {extra})' + + +class AnyHttpUrl(AnyUrl): + allowed_schemes = {'http', 'https'} + + __slots__ = () + + +class HttpUrl(AnyHttpUrl): + tld_required = True + # https://stackoverflow.com/questions/417142/what-is-the-maximum-length-of-a-url-in-different-browsers + max_length = 2083 + hidden_parts = {'port'} + + @staticmethod + def get_default_parts(parts: 'Parts') -> 'Parts': + return {'port': '80' if parts['scheme'] == 'http' else '443'} + + +class FileUrl(AnyUrl): + allowed_schemes = {'file'} + host_required = False + + __slots__ = () + + +class MultiHostDsn(AnyUrl): + __slots__ = AnyUrl.__slots__ + ('hosts',) + + def __init__(self, *args: Any, hosts: Optional[List['HostParts']] = None, **kwargs: Any): + super().__init__(*args, **kwargs) + self.hosts = hosts + + @staticmethod + def _match_url(url: str) -> Optional[Match[str]]: + return multi_host_url_regex().match(url) + + @classmethod + def validate_parts(cls, parts: 'Parts', validate_port: bool = True) -> 'Parts': + return super().validate_parts(parts, validate_port=False) + + @classmethod + def _build_url(cls, m: Match[str], url: str, parts: 'Parts') -> 'MultiHostDsn': + hosts_parts: List['HostParts'] = [] + host_re = host_regex() + for host in m.groupdict()['hosts'].split(','): + d: Parts = host_re.match(host).groupdict() # type: ignore + host, tld, host_type, rebuild = cls.validate_host(d) + port = d.get('port') + cls._validate_port(port) + hosts_parts.append( + { + 'host': host, + 'host_type': host_type, + 'tld': tld, + 'rebuild': rebuild, + 'port': port, + } + ) + + if len(hosts_parts) > 1: + return cls( + None if any([hp['rebuild'] for hp in hosts_parts]) else url, + scheme=parts['scheme'], + user=parts['user'], + password=parts['password'], + path=parts['path'], + query=parts['query'], + fragment=parts['fragment'], + host_type=None, + hosts=hosts_parts, + ) + else: + # backwards compatibility with single host + host_part = hosts_parts[0] + return cls( + None if host_part['rebuild'] else url, + scheme=parts['scheme'], + user=parts['user'], + password=parts['password'], + host=host_part['host'], + tld=host_part['tld'], + host_type=host_part['host_type'], + port=host_part.get('port'), + path=parts['path'], + query=parts['query'], + fragment=parts['fragment'], + ) + + +class PostgresDsn(MultiHostDsn): + allowed_schemes = { + 'postgres', + 'postgresql', + 'postgresql+asyncpg', + 'postgresql+pg8000', + 'postgresql+psycopg', + 'postgresql+psycopg2', + 'postgresql+psycopg2cffi', + 'postgresql+py-postgresql', + 'postgresql+pygresql', + } + user_required = True + + __slots__ = () + + +class CockroachDsn(AnyUrl): + allowed_schemes = { + 'cockroachdb', + 'cockroachdb+psycopg2', + 'cockroachdb+asyncpg', + } + user_required = True + + +class AmqpDsn(AnyUrl): + allowed_schemes = {'amqp', 'amqps'} + host_required = False + + +class RedisDsn(AnyUrl): + __slots__ = () + allowed_schemes = {'redis', 'rediss'} + host_required = False + + @staticmethod + def get_default_parts(parts: 'Parts') -> 'Parts': + return { + 'domain': 'localhost' if not (parts['ipv4'] or parts['ipv6']) else '', + 'port': '6379', + 'path': '/0', + } + + +class MongoDsn(AnyUrl): + allowed_schemes = {'mongodb'} + + # TODO: Needed to generic "Parts" for "Replica Set", "Sharded Cluster", and other mongodb deployment modes + @staticmethod + def get_default_parts(parts: 'Parts') -> 'Parts': + return { + 'port': '27017', + } + + +class KafkaDsn(AnyUrl): + allowed_schemes = {'kafka'} + + @staticmethod + def get_default_parts(parts: 'Parts') -> 'Parts': + return { + 'domain': 'localhost', + 'port': '9092', + } + + +def stricturl( + *, + strip_whitespace: bool = True, + min_length: int = 1, + max_length: int = 2**16, + tld_required: bool = True, + host_required: bool = True, + allowed_schemes: Optional[Collection[str]] = None, +) -> Type[AnyUrl]: + # use kwargs then define conf in a dict to aid with IDE type hinting + namespace = dict( + strip_whitespace=strip_whitespace, + min_length=min_length, + max_length=max_length, + tld_required=tld_required, + host_required=host_required, + allowed_schemes=allowed_schemes, + ) + return type('UrlValue', (AnyUrl,), namespace) + + +def import_email_validator() -> None: + global email_validator + try: + import email_validator + except ImportError as e: + raise ImportError('email-validator is not installed, run `pip install pydantic[email]`') from e + + +class EmailStr(str): + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='string', format='email') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + # included here and below so the error happens straight away + import_email_validator() + + yield str_validator + yield cls.validate + + @classmethod + def validate(cls, value: Union[str]) -> str: + return validate_email(value)[1] + + +class NameEmail(Representation): + __slots__ = 'name', 'email' + + def __init__(self, name: str, email: str): + self.name = name + self.email = email + + def __eq__(self, other: Any) -> bool: + return isinstance(other, NameEmail) and (self.name, self.email) == (other.name, other.email) + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='string', format='name-email') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + import_email_validator() + + yield cls.validate + + @classmethod + def validate(cls, value: Any) -> 'NameEmail': + if value.__class__ == cls: + return value + value = str_validator(value) + return cls(*validate_email(value)) + + def __str__(self) -> str: + return f'{self.name} <{self.email}>' + + +class IPvAnyAddress(_BaseAddress): + __slots__ = () + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='string', format='ipvanyaddress') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + + @classmethod + def validate(cls, value: Union[str, bytes, int]) -> Union[IPv4Address, IPv6Address]: + try: + return IPv4Address(value) + except ValueError: + pass + + try: + return IPv6Address(value) + except ValueError: + raise errors.IPvAnyAddressError() + + +class IPvAnyInterface(_BaseAddress): + __slots__ = () + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='string', format='ipvanyinterface') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + + @classmethod + def validate(cls, value: NetworkType) -> Union[IPv4Interface, IPv6Interface]: + try: + return IPv4Interface(value) + except ValueError: + pass + + try: + return IPv6Interface(value) + except ValueError: + raise errors.IPvAnyInterfaceError() + + +class IPvAnyNetwork(_BaseNetwork): # type: ignore + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='string', format='ipvanynetwork') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + + @classmethod + def validate(cls, value: NetworkType) -> Union[IPv4Network, IPv6Network]: + # Assume IP Network is defined with a default value for ``strict`` argument. + # Define your own class if you want to specify network address check strictness. + try: + return IPv4Network(value) + except ValueError: + pass + + try: + return IPv6Network(value) + except ValueError: + raise errors.IPvAnyNetworkError() + + +pretty_email_regex = re.compile(r'([\w ]*?) *<(.*)> *') +MAX_EMAIL_LENGTH = 2048 +"""Maximum length for an email. +A somewhat arbitrary but very generous number compared to what is allowed by most implementations. +""" + + +def validate_email(value: Union[str]) -> Tuple[str, str]: + """ + Email address validation using https://pypi.org/project/email-validator/ + Notes: + * raw ip address (literal) domain parts are not allowed. + * "John Doe " style "pretty" email addresses are processed + * spaces are striped from the beginning and end of addresses but no error is raised + """ + if email_validator is None: + import_email_validator() + + if len(value) > MAX_EMAIL_LENGTH: + raise errors.EmailError() + + m = pretty_email_regex.fullmatch(value) + name: Union[str, None] = None + if m: + name, value = m.groups() + email = value.strip() + try: + parts = email_validator.validate_email(email, check_deliverability=False) + except email_validator.EmailNotValidError as e: + raise errors.EmailError from e + + if hasattr(parts, 'normalized'): + # email-validator >= 2 + email = parts.normalized + assert email is not None + name = name or parts.local_part + return name, email + else: + # email-validator >1, <2 + at_index = email.index('@') + local_part = email[:at_index] # RFC 5321, local part must be case-sensitive. + global_part = email[at_index:].lower() + + return name or local_part, local_part + global_part diff --git a/lib/python3.12/site-packages/pydantic/v1/parse.py b/lib/python3.12/site-packages/pydantic/v1/parse.py new file mode 100644 index 0000000000000000000000000000000000000000..431d75a6406f3cf79ae4fcd26a2ad35fb5030526 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/parse.py @@ -0,0 +1,66 @@ +import json +import pickle +from enum import Enum +from pathlib import Path +from typing import Any, Callable, Union + +from pydantic.v1.types import StrBytes + + +class Protocol(str, Enum): + json = 'json' + pickle = 'pickle' + + +def load_str_bytes( + b: StrBytes, + *, + content_type: str = None, + encoding: str = 'utf8', + proto: Protocol = None, + allow_pickle: bool = False, + json_loads: Callable[[str], Any] = json.loads, +) -> Any: + if proto is None and content_type: + if content_type.endswith(('json', 'javascript')): + pass + elif allow_pickle and content_type.endswith('pickle'): + proto = Protocol.pickle + else: + raise TypeError(f'Unknown content-type: {content_type}') + + proto = proto or Protocol.json + + if proto == Protocol.json: + if isinstance(b, bytes): + b = b.decode(encoding) + return json_loads(b) + elif proto == Protocol.pickle: + if not allow_pickle: + raise RuntimeError('Trying to decode with pickle with allow_pickle=False') + bb = b if isinstance(b, bytes) else b.encode() + return pickle.loads(bb) + else: + raise TypeError(f'Unknown protocol: {proto}') + + +def load_file( + path: Union[str, Path], + *, + content_type: str = None, + encoding: str = 'utf8', + proto: Protocol = None, + allow_pickle: bool = False, + json_loads: Callable[[str], Any] = json.loads, +) -> Any: + path = Path(path) + b = path.read_bytes() + if content_type is None: + if path.suffix in ('.js', '.json'): + proto = Protocol.json + elif path.suffix == '.pkl': + proto = Protocol.pickle + + return load_str_bytes( + b, proto=proto, content_type=content_type, encoding=encoding, allow_pickle=allow_pickle, json_loads=json_loads + ) diff --git a/lib/python3.12/site-packages/pydantic/v1/py.typed b/lib/python3.12/site-packages/pydantic/v1/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/pydantic/v1/schema.py b/lib/python3.12/site-packages/pydantic/v1/schema.py new file mode 100644 index 0000000000000000000000000000000000000000..a91fe2cd4ae12e619a0ca34baf14400ded5f14cf --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/schema.py @@ -0,0 +1,1163 @@ +import re +import warnings +from collections import defaultdict +from dataclasses import is_dataclass +from datetime import date, datetime, time, timedelta +from decimal import Decimal +from enum import Enum +from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + ForwardRef, + FrozenSet, + Generic, + Iterable, + List, + Optional, + Pattern, + Sequence, + Set, + Tuple, + Type, + TypeVar, + Union, + cast, +) +from uuid import UUID + +from typing_extensions import Annotated, Literal + +from pydantic.v1.fields import ( + MAPPING_LIKE_SHAPES, + SHAPE_DEQUE, + SHAPE_FROZENSET, + SHAPE_GENERIC, + SHAPE_ITERABLE, + SHAPE_LIST, + SHAPE_SEQUENCE, + SHAPE_SET, + SHAPE_SINGLETON, + SHAPE_TUPLE, + SHAPE_TUPLE_ELLIPSIS, + FieldInfo, + ModelField, +) +from pydantic.v1.json import pydantic_encoder +from pydantic.v1.networks import AnyUrl, EmailStr +from pydantic.v1.types import ( + ConstrainedDecimal, + ConstrainedFloat, + ConstrainedFrozenSet, + ConstrainedInt, + ConstrainedList, + ConstrainedSet, + ConstrainedStr, + SecretBytes, + SecretStr, + StrictBytes, + StrictStr, + conbytes, + condecimal, + confloat, + confrozenset, + conint, + conlist, + conset, + constr, +) +from pydantic.v1.typing import ( + all_literal_values, + get_args, + get_origin, + get_sub_types, + is_callable_type, + is_literal_type, + is_namedtuple, + is_none_type, + is_union, +) +from pydantic.v1.utils import ROOT_KEY, get_model, lenient_issubclass + +if TYPE_CHECKING: + from pydantic.v1.dataclasses import Dataclass + from pydantic.v1.main import BaseModel + +default_prefix = '#/definitions/' +default_ref_template = '#/definitions/{model}' + +TypeModelOrEnum = Union[Type['BaseModel'], Type[Enum]] +TypeModelSet = Set[TypeModelOrEnum] + + +def _apply_modify_schema( + modify_schema: Callable[..., None], field: Optional[ModelField], field_schema: Dict[str, Any] +) -> None: + from inspect import signature + + sig = signature(modify_schema) + args = set(sig.parameters.keys()) + if 'field' in args or 'kwargs' in args: + modify_schema(field_schema, field=field) + else: + modify_schema(field_schema) + + +def schema( + models: Sequence[Union[Type['BaseModel'], Type['Dataclass']]], + *, + by_alias: bool = True, + title: Optional[str] = None, + description: Optional[str] = None, + ref_prefix: Optional[str] = None, + ref_template: str = default_ref_template, +) -> Dict[str, Any]: + """ + Process a list of models and generate a single JSON Schema with all of them defined in the ``definitions`` + top-level JSON key, including their sub-models. + + :param models: a list of models to include in the generated JSON Schema + :param by_alias: generate the schemas using the aliases defined, if any + :param title: title for the generated schema that includes the definitions + :param description: description for the generated schema + :param ref_prefix: the JSON Pointer prefix for schema references with ``$ref``, if None, will be set to the + default of ``#/definitions/``. Update it if you want the schemas to reference the definitions somewhere + else, e.g. for OpenAPI use ``#/components/schemas/``. The resulting generated schemas will still be at the + top-level key ``definitions``, so you can extract them from there. But all the references will have the set + prefix. + :param ref_template: Use a ``string.format()`` template for ``$ref`` instead of a prefix. This can be useful + for references that cannot be represented by ``ref_prefix`` such as a definition stored in another file. For + a sibling json file in a ``/schemas`` directory use ``"/schemas/${model}.json#"``. + :return: dict with the JSON Schema with a ``definitions`` top-level key including the schema definitions for + the models and sub-models passed in ``models``. + """ + clean_models = [get_model(model) for model in models] + flat_models = get_flat_models_from_models(clean_models) + model_name_map = get_model_name_map(flat_models) + definitions = {} + output_schema: Dict[str, Any] = {} + if title: + output_schema['title'] = title + if description: + output_schema['description'] = description + for model in clean_models: + m_schema, m_definitions, m_nested_models = model_process_schema( + model, + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + ) + definitions.update(m_definitions) + model_name = model_name_map[model] + definitions[model_name] = m_schema + if definitions: + output_schema['definitions'] = definitions + return output_schema + + +def model_schema( + model: Union[Type['BaseModel'], Type['Dataclass']], + by_alias: bool = True, + ref_prefix: Optional[str] = None, + ref_template: str = default_ref_template, +) -> Dict[str, Any]: + """ + Generate a JSON Schema for one model. With all the sub-models defined in the ``definitions`` top-level + JSON key. + + :param model: a Pydantic model (a class that inherits from BaseModel) + :param by_alias: generate the schemas using the aliases defined, if any + :param ref_prefix: the JSON Pointer prefix for schema references with ``$ref``, if None, will be set to the + default of ``#/definitions/``. Update it if you want the schemas to reference the definitions somewhere + else, e.g. for OpenAPI use ``#/components/schemas/``. The resulting generated schemas will still be at the + top-level key ``definitions``, so you can extract them from there. But all the references will have the set + prefix. + :param ref_template: Use a ``string.format()`` template for ``$ref`` instead of a prefix. This can be useful for + references that cannot be represented by ``ref_prefix`` such as a definition stored in another file. For a + sibling json file in a ``/schemas`` directory use ``"/schemas/${model}.json#"``. + :return: dict with the JSON Schema for the passed ``model`` + """ + model = get_model(model) + flat_models = get_flat_models_from_model(model) + model_name_map = get_model_name_map(flat_models) + model_name = model_name_map[model] + m_schema, m_definitions, nested_models = model_process_schema( + model, by_alias=by_alias, model_name_map=model_name_map, ref_prefix=ref_prefix, ref_template=ref_template + ) + if model_name in nested_models: + # model_name is in Nested models, it has circular references + m_definitions[model_name] = m_schema + m_schema = get_schema_ref(model_name, ref_prefix, ref_template, False) + if m_definitions: + m_schema.update({'definitions': m_definitions}) + return m_schema + + +def get_field_info_schema(field: ModelField, schema_overrides: bool = False) -> Tuple[Dict[str, Any], bool]: + # If no title is explicitly set, we don't set title in the schema for enums. + # The behaviour is the same as `BaseModel` reference, where the default title + # is in the definitions part of the schema. + schema_: Dict[str, Any] = {} + if field.field_info.title or not lenient_issubclass(field.type_, Enum): + schema_['title'] = field.field_info.title or field.alias.title().replace('_', ' ') + + if field.field_info.title: + schema_overrides = True + + if field.field_info.description: + schema_['description'] = field.field_info.description + schema_overrides = True + + if not field.required and field.default is not None and not is_callable_type(field.outer_type_): + schema_['default'] = encode_default(field.default) + schema_overrides = True + + return schema_, schema_overrides + + +def field_schema( + field: ModelField, + *, + by_alias: bool = True, + model_name_map: Dict[TypeModelOrEnum, str], + ref_prefix: Optional[str] = None, + ref_template: str = default_ref_template, + known_models: Optional[TypeModelSet] = None, +) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]: + """ + Process a Pydantic field and return a tuple with a JSON Schema for it as the first item. + Also return a dictionary of definitions with models as keys and their schemas as values. If the passed field + is a model and has sub-models, and those sub-models don't have overrides (as ``title``, ``default``, etc), they + will be included in the definitions and referenced in the schema instead of included recursively. + + :param field: a Pydantic ``ModelField`` + :param by_alias: use the defined alias (if any) in the returned schema + :param model_name_map: used to generate the JSON Schema references to other models included in the definitions + :param ref_prefix: the JSON Pointer prefix to use for references to other schemas, if None, the default of + #/definitions/ will be used + :param ref_template: Use a ``string.format()`` template for ``$ref`` instead of a prefix. This can be useful for + references that cannot be represented by ``ref_prefix`` such as a definition stored in another file. For a + sibling json file in a ``/schemas`` directory use ``"/schemas/${model}.json#"``. + :param known_models: used to solve circular references + :return: tuple of the schema for this field and additional definitions + """ + s, schema_overrides = get_field_info_schema(field) + + validation_schema = get_field_schema_validations(field) + if validation_schema: + s.update(validation_schema) + schema_overrides = True + + f_schema, f_definitions, f_nested_models = field_type_schema( + field, + by_alias=by_alias, + model_name_map=model_name_map, + schema_overrides=schema_overrides, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models or set(), + ) + + # $ref will only be returned when there are no schema_overrides + if '$ref' in f_schema: + return f_schema, f_definitions, f_nested_models + else: + s.update(f_schema) + return s, f_definitions, f_nested_models + + +numeric_types = (int, float, Decimal) +_str_types_attrs: Tuple[Tuple[str, Union[type, Tuple[type, ...]], str], ...] = ( + ('max_length', numeric_types, 'maxLength'), + ('min_length', numeric_types, 'minLength'), + ('regex', str, 'pattern'), +) + +_numeric_types_attrs: Tuple[Tuple[str, Union[type, Tuple[type, ...]], str], ...] = ( + ('gt', numeric_types, 'exclusiveMinimum'), + ('lt', numeric_types, 'exclusiveMaximum'), + ('ge', numeric_types, 'minimum'), + ('le', numeric_types, 'maximum'), + ('multiple_of', numeric_types, 'multipleOf'), +) + + +def get_field_schema_validations(field: ModelField) -> Dict[str, Any]: + """ + Get the JSON Schema validation keywords for a ``field`` with an annotation of + a Pydantic ``FieldInfo`` with validation arguments. + """ + f_schema: Dict[str, Any] = {} + + if lenient_issubclass(field.type_, Enum): + # schema is already updated by `enum_process_schema`; just update with field extra + if field.field_info.extra: + f_schema.update(field.field_info.extra) + return f_schema + + if lenient_issubclass(field.type_, (str, bytes)): + for attr_name, t, keyword in _str_types_attrs: + attr = getattr(field.field_info, attr_name, None) + if isinstance(attr, t): + f_schema[keyword] = attr + if lenient_issubclass(field.type_, numeric_types) and not issubclass(field.type_, bool): + for attr_name, t, keyword in _numeric_types_attrs: + attr = getattr(field.field_info, attr_name, None) + if isinstance(attr, t): + f_schema[keyword] = attr + if field.field_info is not None and field.field_info.const: + f_schema['const'] = field.default + if field.field_info.extra: + f_schema.update(field.field_info.extra) + modify_schema = getattr(field.outer_type_, '__modify_schema__', None) + if modify_schema: + _apply_modify_schema(modify_schema, field, f_schema) + return f_schema + + +def get_model_name_map(unique_models: TypeModelSet) -> Dict[TypeModelOrEnum, str]: + """ + Process a set of models and generate unique names for them to be used as keys in the JSON Schema + definitions. By default the names are the same as the class name. But if two models in different Python + modules have the same name (e.g. "users.Model" and "items.Model"), the generated names will be + based on the Python module path for those conflicting models to prevent name collisions. + + :param unique_models: a Python set of models + :return: dict mapping models to names + """ + name_model_map = {} + conflicting_names: Set[str] = set() + for model in unique_models: + model_name = normalize_name(model.__name__) + if model_name in conflicting_names: + model_name = get_long_model_name(model) + name_model_map[model_name] = model + elif model_name in name_model_map: + conflicting_names.add(model_name) + conflicting_model = name_model_map.pop(model_name) + name_model_map[get_long_model_name(conflicting_model)] = conflicting_model + name_model_map[get_long_model_name(model)] = model + else: + name_model_map[model_name] = model + return {v: k for k, v in name_model_map.items()} + + +def get_flat_models_from_model(model: Type['BaseModel'], known_models: Optional[TypeModelSet] = None) -> TypeModelSet: + """ + Take a single ``model`` and generate a set with itself and all the sub-models in the tree. I.e. if you pass + model ``Foo`` (subclass of Pydantic ``BaseModel``) as ``model``, and it has a field of type ``Bar`` (also + subclass of ``BaseModel``) and that model ``Bar`` has a field of type ``Baz`` (also subclass of ``BaseModel``), + the return value will be ``set([Foo, Bar, Baz])``. + + :param model: a Pydantic ``BaseModel`` subclass + :param known_models: used to solve circular references + :return: a set with the initial model and all its sub-models + """ + known_models = known_models or set() + flat_models: TypeModelSet = set() + flat_models.add(model) + known_models |= flat_models + fields = cast(Sequence[ModelField], model.__fields__.values()) + flat_models |= get_flat_models_from_fields(fields, known_models=known_models) + return flat_models + + +def get_flat_models_from_field(field: ModelField, known_models: TypeModelSet) -> TypeModelSet: + """ + Take a single Pydantic ``ModelField`` (from a model) that could have been declared as a subclass of BaseModel + (so, it could be a submodel), and generate a set with its model and all the sub-models in the tree. + I.e. if you pass a field that was declared to be of type ``Foo`` (subclass of BaseModel) as ``field``, and that + model ``Foo`` has a field of type ``Bar`` (also subclass of ``BaseModel``) and that model ``Bar`` has a field of + type ``Baz`` (also subclass of ``BaseModel``), the return value will be ``set([Foo, Bar, Baz])``. + + :param field: a Pydantic ``ModelField`` + :param known_models: used to solve circular references + :return: a set with the model used in the declaration for this field, if any, and all its sub-models + """ + from pydantic.v1.main import BaseModel + + flat_models: TypeModelSet = set() + + field_type = field.type_ + if lenient_issubclass(getattr(field_type, '__pydantic_model__', None), BaseModel): + field_type = field_type.__pydantic_model__ + + if field.sub_fields and not lenient_issubclass(field_type, BaseModel): + flat_models |= get_flat_models_from_fields(field.sub_fields, known_models=known_models) + elif lenient_issubclass(field_type, BaseModel) and field_type not in known_models: + flat_models |= get_flat_models_from_model(field_type, known_models=known_models) + elif lenient_issubclass(field_type, Enum): + flat_models.add(field_type) + return flat_models + + +def get_flat_models_from_fields(fields: Sequence[ModelField], known_models: TypeModelSet) -> TypeModelSet: + """ + Take a list of Pydantic ``ModelField``s (from a model) that could have been declared as subclasses of ``BaseModel`` + (so, any of them could be a submodel), and generate a set with their models and all the sub-models in the tree. + I.e. if you pass a the fields of a model ``Foo`` (subclass of ``BaseModel``) as ``fields``, and on of them has a + field of type ``Bar`` (also subclass of ``BaseModel``) and that model ``Bar`` has a field of type ``Baz`` (also + subclass of ``BaseModel``), the return value will be ``set([Foo, Bar, Baz])``. + + :param fields: a list of Pydantic ``ModelField``s + :param known_models: used to solve circular references + :return: a set with any model declared in the fields, and all their sub-models + """ + flat_models: TypeModelSet = set() + for field in fields: + flat_models |= get_flat_models_from_field(field, known_models=known_models) + return flat_models + + +def get_flat_models_from_models(models: Sequence[Type['BaseModel']]) -> TypeModelSet: + """ + Take a list of ``models`` and generate a set with them and all their sub-models in their trees. I.e. if you pass + a list of two models, ``Foo`` and ``Bar``, both subclasses of Pydantic ``BaseModel`` as models, and ``Bar`` has + a field of type ``Baz`` (also subclass of ``BaseModel``), the return value will be ``set([Foo, Bar, Baz])``. + """ + flat_models: TypeModelSet = set() + for model in models: + flat_models |= get_flat_models_from_model(model) + return flat_models + + +def get_long_model_name(model: TypeModelOrEnum) -> str: + return f'{model.__module__}__{model.__qualname__}'.replace('.', '__') + + +def field_type_schema( + field: ModelField, + *, + by_alias: bool, + model_name_map: Dict[TypeModelOrEnum, str], + ref_template: str, + schema_overrides: bool = False, + ref_prefix: Optional[str] = None, + known_models: TypeModelSet, +) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]: + """ + Used by ``field_schema()``, you probably should be using that function. + + Take a single ``field`` and generate the schema for its type only, not including additional + information as title, etc. Also return additional schema definitions, from sub-models. + """ + from pydantic.v1.main import BaseModel # noqa: F811 + + definitions = {} + nested_models: Set[str] = set() + f_schema: Dict[str, Any] + if field.shape in { + SHAPE_LIST, + SHAPE_TUPLE_ELLIPSIS, + SHAPE_SEQUENCE, + SHAPE_SET, + SHAPE_FROZENSET, + SHAPE_ITERABLE, + SHAPE_DEQUE, + }: + items_schema, f_definitions, f_nested_models = field_singleton_schema( + field, + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + definitions.update(f_definitions) + nested_models.update(f_nested_models) + f_schema = {'type': 'array', 'items': items_schema} + if field.shape in {SHAPE_SET, SHAPE_FROZENSET}: + f_schema['uniqueItems'] = True + + elif field.shape in MAPPING_LIKE_SHAPES: + f_schema = {'type': 'object'} + key_field = cast(ModelField, field.key_field) + regex = getattr(key_field.type_, 'regex', None) + items_schema, f_definitions, f_nested_models = field_singleton_schema( + field, + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + definitions.update(f_definitions) + nested_models.update(f_nested_models) + if regex: + # Dict keys have a regex pattern + # items_schema might be a schema or empty dict, add it either way + f_schema['patternProperties'] = {ConstrainedStr._get_pattern(regex): items_schema} + if items_schema: + # The dict values are not simply Any, so they need a schema + f_schema['additionalProperties'] = items_schema + elif field.shape == SHAPE_TUPLE or (field.shape == SHAPE_GENERIC and not issubclass(field.type_, BaseModel)): + sub_schema = [] + sub_fields = cast(List[ModelField], field.sub_fields) + for sf in sub_fields: + sf_schema, sf_definitions, sf_nested_models = field_type_schema( + sf, + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + definitions.update(sf_definitions) + nested_models.update(sf_nested_models) + sub_schema.append(sf_schema) + + sub_fields_len = len(sub_fields) + if field.shape == SHAPE_GENERIC: + all_of_schemas = sub_schema[0] if sub_fields_len == 1 else {'type': 'array', 'items': sub_schema} + f_schema = {'allOf': [all_of_schemas]} + else: + f_schema = { + 'type': 'array', + 'minItems': sub_fields_len, + 'maxItems': sub_fields_len, + } + if sub_fields_len >= 1: + f_schema['items'] = sub_schema + else: + assert field.shape in {SHAPE_SINGLETON, SHAPE_GENERIC}, field.shape + f_schema, f_definitions, f_nested_models = field_singleton_schema( + field, + by_alias=by_alias, + model_name_map=model_name_map, + schema_overrides=schema_overrides, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + definitions.update(f_definitions) + nested_models.update(f_nested_models) + + # check field type to avoid repeated calls to the same __modify_schema__ method + if field.type_ != field.outer_type_: + if field.shape == SHAPE_GENERIC: + field_type = field.type_ + else: + field_type = field.outer_type_ + modify_schema = getattr(field_type, '__modify_schema__', None) + if modify_schema: + _apply_modify_schema(modify_schema, field, f_schema) + return f_schema, definitions, nested_models + + +def model_process_schema( + model: TypeModelOrEnum, + *, + by_alias: bool = True, + model_name_map: Dict[TypeModelOrEnum, str], + ref_prefix: Optional[str] = None, + ref_template: str = default_ref_template, + known_models: Optional[TypeModelSet] = None, + field: Optional[ModelField] = None, +) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]: + """ + Used by ``model_schema()``, you probably should be using that function. + + Take a single ``model`` and generate its schema. Also return additional schema definitions, from sub-models. The + sub-models of the returned schema will be referenced, but their definitions will not be included in the schema. All + the definitions are returned as the second value. + """ + from inspect import getdoc, signature + + known_models = known_models or set() + if lenient_issubclass(model, Enum): + model = cast(Type[Enum], model) + s = enum_process_schema(model, field=field) + return s, {}, set() + model = cast(Type['BaseModel'], model) + s = {'title': model.__config__.title or model.__name__} + doc = getdoc(model) + if doc: + s['description'] = doc + known_models.add(model) + m_schema, m_definitions, nested_models = model_type_schema( + model, + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + s.update(m_schema) + schema_extra = model.__config__.schema_extra + if callable(schema_extra): + if len(signature(schema_extra).parameters) == 1: + schema_extra(s) + else: + schema_extra(s, model) + else: + s.update(schema_extra) + return s, m_definitions, nested_models + + +def model_type_schema( + model: Type['BaseModel'], + *, + by_alias: bool, + model_name_map: Dict[TypeModelOrEnum, str], + ref_template: str, + ref_prefix: Optional[str] = None, + known_models: TypeModelSet, +) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]: + """ + You probably should be using ``model_schema()``, this function is indirectly used by that function. + + Take a single ``model`` and generate the schema for its type only, not including additional + information as title, etc. Also return additional schema definitions, from sub-models. + """ + properties = {} + required = [] + definitions: Dict[str, Any] = {} + nested_models: Set[str] = set() + for k, f in model.__fields__.items(): + try: + f_schema, f_definitions, f_nested_models = field_schema( + f, + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + except SkipField as skip: + warnings.warn(skip.message, UserWarning) + continue + definitions.update(f_definitions) + nested_models.update(f_nested_models) + if by_alias: + properties[f.alias] = f_schema + if f.required: + required.append(f.alias) + else: + properties[k] = f_schema + if f.required: + required.append(k) + if ROOT_KEY in properties: + out_schema = properties[ROOT_KEY] + out_schema['title'] = model.__config__.title or model.__name__ + else: + out_schema = {'type': 'object', 'properties': properties} + if required: + out_schema['required'] = required + if model.__config__.extra == 'forbid': + out_schema['additionalProperties'] = False + return out_schema, definitions, nested_models + + +def enum_process_schema(enum: Type[Enum], *, field: Optional[ModelField] = None) -> Dict[str, Any]: + """ + Take a single `enum` and generate its schema. + + This is similar to the `model_process_schema` function, but applies to ``Enum`` objects. + """ + import inspect + + schema_: Dict[str, Any] = { + 'title': enum.__name__, + # Python assigns all enums a default docstring value of 'An enumeration', so + # all enums will have a description field even if not explicitly provided. + 'description': inspect.cleandoc(enum.__doc__ or 'An enumeration.'), + # Add enum values and the enum field type to the schema. + 'enum': [item.value for item in cast(Iterable[Enum], enum)], + } + + add_field_type_to_schema(enum, schema_) + + modify_schema = getattr(enum, '__modify_schema__', None) + if modify_schema: + _apply_modify_schema(modify_schema, field, schema_) + + return schema_ + + +def field_singleton_sub_fields_schema( + field: ModelField, + *, + by_alias: bool, + model_name_map: Dict[TypeModelOrEnum, str], + ref_template: str, + schema_overrides: bool = False, + ref_prefix: Optional[str] = None, + known_models: TypeModelSet, +) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]: + """ + This function is indirectly used by ``field_schema()``, you probably should be using that function. + + Take a list of Pydantic ``ModelField`` from the declaration of a type with parameters, and generate their + schema. I.e., fields used as "type parameters", like ``str`` and ``int`` in ``Tuple[str, int]``. + """ + sub_fields = cast(List[ModelField], field.sub_fields) + definitions = {} + nested_models: Set[str] = set() + if len(sub_fields) == 1: + return field_type_schema( + sub_fields[0], + by_alias=by_alias, + model_name_map=model_name_map, + schema_overrides=schema_overrides, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + else: + s: Dict[str, Any] = {} + # https://github.com/OAI/OpenAPI-Specification/blob/master/versions/3.0.2.md#discriminator-object + field_has_discriminator: bool = field.discriminator_key is not None + if field_has_discriminator: + assert field.sub_fields_mapping is not None + + discriminator_models_refs: Dict[str, Union[str, Dict[str, Any]]] = {} + + for discriminator_value, sub_field in field.sub_fields_mapping.items(): + if isinstance(discriminator_value, Enum): + discriminator_value = str(discriminator_value.value) + # sub_field is either a `BaseModel` or directly an `Annotated` `Union` of many + if is_union(get_origin(sub_field.type_)): + sub_models = get_sub_types(sub_field.type_) + discriminator_models_refs[discriminator_value] = { + model_name_map[sub_model]: get_schema_ref( + model_name_map[sub_model], ref_prefix, ref_template, False + ) + for sub_model in sub_models + } + else: + sub_field_type = sub_field.type_ + if hasattr(sub_field_type, '__pydantic_model__'): + sub_field_type = sub_field_type.__pydantic_model__ + + discriminator_model_name = model_name_map[sub_field_type] + discriminator_model_ref = get_schema_ref(discriminator_model_name, ref_prefix, ref_template, False) + discriminator_models_refs[discriminator_value] = discriminator_model_ref['$ref'] + + s['discriminator'] = { + 'propertyName': field.discriminator_alias if by_alias else field.discriminator_key, + 'mapping': discriminator_models_refs, + } + + sub_field_schemas = [] + for sf in sub_fields: + sub_schema, sub_definitions, sub_nested_models = field_type_schema( + sf, + by_alias=by_alias, + model_name_map=model_name_map, + schema_overrides=schema_overrides, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + definitions.update(sub_definitions) + if schema_overrides and 'allOf' in sub_schema: + # if the sub_field is a referenced schema we only need the referenced + # object. Otherwise we will end up with several allOf inside anyOf/oneOf. + # See https://github.com/pydantic/pydantic/issues/1209 + sub_schema = sub_schema['allOf'][0] + + if sub_schema.keys() == {'discriminator', 'oneOf'}: + # we don't want discriminator information inside oneOf choices, this is dealt with elsewhere + sub_schema.pop('discriminator') + sub_field_schemas.append(sub_schema) + nested_models.update(sub_nested_models) + s['oneOf' if field_has_discriminator else 'anyOf'] = sub_field_schemas + return s, definitions, nested_models + + +# Order is important, e.g. subclasses of str must go before str +# this is used only for standard library types, custom types should use __modify_schema__ instead +field_class_to_schema: Tuple[Tuple[Any, Dict[str, Any]], ...] = ( + (Path, {'type': 'string', 'format': 'path'}), + (datetime, {'type': 'string', 'format': 'date-time'}), + (date, {'type': 'string', 'format': 'date'}), + (time, {'type': 'string', 'format': 'time'}), + (timedelta, {'type': 'number', 'format': 'time-delta'}), + (IPv4Network, {'type': 'string', 'format': 'ipv4network'}), + (IPv6Network, {'type': 'string', 'format': 'ipv6network'}), + (IPv4Interface, {'type': 'string', 'format': 'ipv4interface'}), + (IPv6Interface, {'type': 'string', 'format': 'ipv6interface'}), + (IPv4Address, {'type': 'string', 'format': 'ipv4'}), + (IPv6Address, {'type': 'string', 'format': 'ipv6'}), + (Pattern, {'type': 'string', 'format': 'regex'}), + (str, {'type': 'string'}), + (bytes, {'type': 'string', 'format': 'binary'}), + (bool, {'type': 'boolean'}), + (int, {'type': 'integer'}), + (float, {'type': 'number'}), + (Decimal, {'type': 'number'}), + (UUID, {'type': 'string', 'format': 'uuid'}), + (dict, {'type': 'object'}), + (list, {'type': 'array', 'items': {}}), + (tuple, {'type': 'array', 'items': {}}), + (set, {'type': 'array', 'items': {}, 'uniqueItems': True}), + (frozenset, {'type': 'array', 'items': {}, 'uniqueItems': True}), +) + +json_scheme = {'type': 'string', 'format': 'json-string'} + + +def add_field_type_to_schema(field_type: Any, schema_: Dict[str, Any]) -> None: + """ + Update the given `schema` with the type-specific metadata for the given `field_type`. + + This function looks through `field_class_to_schema` for a class that matches the given `field_type`, + and then modifies the given `schema` with the information from that type. + """ + for type_, t_schema in field_class_to_schema: + # Fallback for `typing.Pattern` and `re.Pattern` as they are not a valid class + if lenient_issubclass(field_type, type_) or field_type is type_ is Pattern: + schema_.update(t_schema) + break + + +def get_schema_ref(name: str, ref_prefix: Optional[str], ref_template: str, schema_overrides: bool) -> Dict[str, Any]: + if ref_prefix: + schema_ref = {'$ref': ref_prefix + name} + else: + schema_ref = {'$ref': ref_template.format(model=name)} + return {'allOf': [schema_ref]} if schema_overrides else schema_ref + + +def field_singleton_schema( # noqa: C901 (ignore complexity) + field: ModelField, + *, + by_alias: bool, + model_name_map: Dict[TypeModelOrEnum, str], + ref_template: str, + schema_overrides: bool = False, + ref_prefix: Optional[str] = None, + known_models: TypeModelSet, +) -> Tuple[Dict[str, Any], Dict[str, Any], Set[str]]: + """ + This function is indirectly used by ``field_schema()``, you should probably be using that function. + + Take a single Pydantic ``ModelField``, and return its schema and any additional definitions from sub-models. + """ + from pydantic.v1.main import BaseModel + + definitions: Dict[str, Any] = {} + nested_models: Set[str] = set() + field_type = field.type_ + + # Recurse into this field if it contains sub_fields and is NOT a + # BaseModel OR that BaseModel is a const + if field.sub_fields and ( + (field.field_info and field.field_info.const) or not lenient_issubclass(field_type, BaseModel) + ): + return field_singleton_sub_fields_schema( + field, + by_alias=by_alias, + model_name_map=model_name_map, + schema_overrides=schema_overrides, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + if field_type is Any or field_type is object or field_type.__class__ == TypeVar or get_origin(field_type) is type: + return {}, definitions, nested_models # no restrictions + if is_none_type(field_type): + return {'type': 'null'}, definitions, nested_models + if is_callable_type(field_type): + raise SkipField(f'Callable {field.name} was excluded from schema since JSON schema has no equivalent type.') + f_schema: Dict[str, Any] = {} + if field.field_info is not None and field.field_info.const: + f_schema['const'] = field.default + + if is_literal_type(field_type): + values = tuple(x.value if isinstance(x, Enum) else x for x in all_literal_values(field_type)) + + if len({v.__class__ for v in values}) > 1: + return field_schema( + multitypes_literal_field_for_schema(values, field), + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + ) + + # All values have the same type + field_type = values[0].__class__ + f_schema['enum'] = list(values) + add_field_type_to_schema(field_type, f_schema) + elif lenient_issubclass(field_type, Enum): + enum_name = model_name_map[field_type] + f_schema, schema_overrides = get_field_info_schema(field, schema_overrides) + f_schema.update(get_schema_ref(enum_name, ref_prefix, ref_template, schema_overrides)) + definitions[enum_name] = enum_process_schema(field_type, field=field) + elif is_namedtuple(field_type): + sub_schema, *_ = model_process_schema( + field_type.__pydantic_model__, + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + field=field, + ) + items_schemas = list(sub_schema['properties'].values()) + f_schema.update( + { + 'type': 'array', + 'items': items_schemas, + 'minItems': len(items_schemas), + 'maxItems': len(items_schemas), + } + ) + elif not hasattr(field_type, '__pydantic_model__'): + add_field_type_to_schema(field_type, f_schema) + + modify_schema = getattr(field_type, '__modify_schema__', None) + if modify_schema: + _apply_modify_schema(modify_schema, field, f_schema) + + if f_schema: + return f_schema, definitions, nested_models + + # Handle dataclass-based models + if lenient_issubclass(getattr(field_type, '__pydantic_model__', None), BaseModel): + field_type = field_type.__pydantic_model__ + + if issubclass(field_type, BaseModel): + model_name = model_name_map[field_type] + if field_type not in known_models: + sub_schema, sub_definitions, sub_nested_models = model_process_schema( + field_type, + by_alias=by_alias, + model_name_map=model_name_map, + ref_prefix=ref_prefix, + ref_template=ref_template, + known_models=known_models, + field=field, + ) + definitions.update(sub_definitions) + definitions[model_name] = sub_schema + nested_models.update(sub_nested_models) + else: + nested_models.add(model_name) + schema_ref = get_schema_ref(model_name, ref_prefix, ref_template, schema_overrides) + return schema_ref, definitions, nested_models + + # For generics with no args + args = get_args(field_type) + if args is not None and not args and Generic in field_type.__bases__: + return f_schema, definitions, nested_models + + raise ValueError(f'Value not declarable with JSON Schema, field: {field}') + + +def multitypes_literal_field_for_schema(values: Tuple[Any, ...], field: ModelField) -> ModelField: + """ + To support `Literal` with values of different types, we split it into multiple `Literal` with same type + e.g. `Literal['qwe', 'asd', 1, 2]` becomes `Union[Literal['qwe', 'asd'], Literal[1, 2]]` + """ + literal_distinct_types = defaultdict(list) + for v in values: + literal_distinct_types[v.__class__].append(v) + distinct_literals = (Literal[tuple(same_type_values)] for same_type_values in literal_distinct_types.values()) + + return ModelField( + name=field.name, + type_=Union[tuple(distinct_literals)], # type: ignore + class_validators=field.class_validators, + model_config=field.model_config, + default=field.default, + required=field.required, + alias=field.alias, + field_info=field.field_info, + ) + + +def encode_default(dft: Any) -> Any: + from pydantic.v1.main import BaseModel + + if isinstance(dft, BaseModel) or is_dataclass(dft): + dft = cast('dict[str, Any]', pydantic_encoder(dft)) + + if isinstance(dft, dict): + return {encode_default(k): encode_default(v) for k, v in dft.items()} + elif isinstance(dft, Enum): + return dft.value + elif isinstance(dft, (int, float, str)): + return dft + elif isinstance(dft, (list, tuple)): + t = dft.__class__ + seq_args = (encode_default(v) for v in dft) + return t(*seq_args) if is_namedtuple(t) else t(seq_args) + elif dft is None: + return None + else: + return pydantic_encoder(dft) + + +_map_types_constraint: Dict[Any, Callable[..., type]] = {int: conint, float: confloat, Decimal: condecimal} + + +def get_annotation_from_field_info( + annotation: Any, field_info: FieldInfo, field_name: str, validate_assignment: bool = False +) -> Type[Any]: + """ + Get an annotation with validation implemented for numbers and strings based on the field_info. + :param annotation: an annotation from a field specification, as ``str``, ``ConstrainedStr`` + :param field_info: an instance of FieldInfo, possibly with declarations for validations and JSON Schema + :param field_name: name of the field for use in error messages + :param validate_assignment: default False, flag for BaseModel Config value of validate_assignment + :return: the same ``annotation`` if unmodified or a new annotation with validation in place + """ + constraints = field_info.get_constraints() + used_constraints: Set[str] = set() + if constraints: + annotation, used_constraints = get_annotation_with_constraints(annotation, field_info) + if validate_assignment: + used_constraints.add('allow_mutation') + + unused_constraints = constraints - used_constraints + if unused_constraints: + raise ValueError( + f'On field "{field_name}" the following field constraints are set but not enforced: ' + f'{", ".join(unused_constraints)}. ' + f'\nFor more details see https://docs.pydantic.dev/usage/schema/#unenforced-field-constraints' + ) + + return annotation + + +def get_annotation_with_constraints(annotation: Any, field_info: FieldInfo) -> Tuple[Type[Any], Set[str]]: # noqa: C901 + """ + Get an annotation with used constraints implemented for numbers and strings based on the field_info. + + :param annotation: an annotation from a field specification, as ``str``, ``ConstrainedStr`` + :param field_info: an instance of FieldInfo, possibly with declarations for validations and JSON Schema + :return: the same ``annotation`` if unmodified or a new annotation along with the used constraints. + """ + used_constraints: Set[str] = set() + + def go(type_: Any) -> Type[Any]: + if ( + is_literal_type(type_) + or isinstance(type_, ForwardRef) + or lenient_issubclass(type_, (ConstrainedList, ConstrainedSet, ConstrainedFrozenSet)) + ): + return type_ + origin = get_origin(type_) + if origin is not None: + args: Tuple[Any, ...] = get_args(type_) + if any(isinstance(a, ForwardRef) for a in args): + # forward refs cause infinite recursion below + return type_ + + if origin is Annotated: + return go(args[0]) + if is_union(origin): + return Union[tuple(go(a) for a in args)] # type: ignore + + if issubclass(origin, List) and ( + field_info.min_items is not None + or field_info.max_items is not None + or field_info.unique_items is not None + ): + used_constraints.update({'min_items', 'max_items', 'unique_items'}) + return conlist( + go(args[0]), + min_items=field_info.min_items, + max_items=field_info.max_items, + unique_items=field_info.unique_items, + ) + + if issubclass(origin, Set) and (field_info.min_items is not None or field_info.max_items is not None): + used_constraints.update({'min_items', 'max_items'}) + return conset(go(args[0]), min_items=field_info.min_items, max_items=field_info.max_items) + + if issubclass(origin, FrozenSet) and (field_info.min_items is not None or field_info.max_items is not None): + used_constraints.update({'min_items', 'max_items'}) + return confrozenset(go(args[0]), min_items=field_info.min_items, max_items=field_info.max_items) + + for t in (Tuple, List, Set, FrozenSet, Sequence): + if issubclass(origin, t): # type: ignore + return t[tuple(go(a) for a in args)] # type: ignore + + if issubclass(origin, Dict): + return Dict[args[0], go(args[1])] # type: ignore + + attrs: Optional[Tuple[str, ...]] = None + constraint_func: Optional[Callable[..., type]] = None + if isinstance(type_, type): + if issubclass(type_, (SecretStr, SecretBytes)): + attrs = ('max_length', 'min_length') + + def constraint_func(**kw: Any) -> Type[Any]: # noqa: F811 + return type(type_.__name__, (type_,), kw) + + elif issubclass(type_, str) and not issubclass(type_, (EmailStr, AnyUrl)): + attrs = ('max_length', 'min_length', 'regex') + if issubclass(type_, StrictStr): + + def constraint_func(**kw: Any) -> Type[Any]: + return type(type_.__name__, (type_,), kw) + + else: + constraint_func = constr + elif issubclass(type_, bytes): + attrs = ('max_length', 'min_length', 'regex') + if issubclass(type_, StrictBytes): + + def constraint_func(**kw: Any) -> Type[Any]: + return type(type_.__name__, (type_,), kw) + + else: + constraint_func = conbytes + elif issubclass(type_, numeric_types) and not issubclass( + type_, + ( + ConstrainedInt, + ConstrainedFloat, + ConstrainedDecimal, + ConstrainedList, + ConstrainedSet, + ConstrainedFrozenSet, + bool, + ), + ): + # Is numeric type + attrs = ('gt', 'lt', 'ge', 'le', 'multiple_of') + if issubclass(type_, float): + attrs += ('allow_inf_nan',) + if issubclass(type_, Decimal): + attrs += ('max_digits', 'decimal_places') + numeric_type = next(t for t in numeric_types if issubclass(type_, t)) # pragma: no branch + constraint_func = _map_types_constraint[numeric_type] + + if attrs: + used_constraints.update(set(attrs)) + kwargs = { + attr_name: attr + for attr_name, attr in ((attr_name, getattr(field_info, attr_name)) for attr_name in attrs) + if attr is not None + } + if kwargs: + constraint_func = cast(Callable[..., type], constraint_func) + return constraint_func(**kwargs) + return type_ + + return go(annotation), used_constraints + + +def normalize_name(name: str) -> str: + """ + Normalizes the given name. This can be applied to either a model *or* enum. + """ + return re.sub(r'[^a-zA-Z0-9.\-_]', '_', name) + + +class SkipField(Exception): + """ + Utility exception used to exclude fields from schema. + """ + + def __init__(self, message: str) -> None: + self.message = message diff --git a/lib/python3.12/site-packages/pydantic/v1/tools.py b/lib/python3.12/site-packages/pydantic/v1/tools.py new file mode 100644 index 0000000000000000000000000000000000000000..6838a23ecc7419f0478449d9d680c230a99c20e0 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/tools.py @@ -0,0 +1,92 @@ +import json +from functools import lru_cache +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Optional, Type, TypeVar, Union + +from pydantic.v1.parse import Protocol, load_file, load_str_bytes +from pydantic.v1.types import StrBytes +from pydantic.v1.typing import display_as_type + +__all__ = ('parse_file_as', 'parse_obj_as', 'parse_raw_as', 'schema_of', 'schema_json_of') + +NameFactory = Union[str, Callable[[Type[Any]], str]] + +if TYPE_CHECKING: + from pydantic.v1.typing import DictStrAny + + +def _generate_parsing_type_name(type_: Any) -> str: + return f'ParsingModel[{display_as_type(type_)}]' + + +@lru_cache(maxsize=2048) +def _get_parsing_type(type_: Any, *, type_name: Optional[NameFactory] = None) -> Any: + from pydantic.v1.main import create_model + + if type_name is None: + type_name = _generate_parsing_type_name + if not isinstance(type_name, str): + type_name = type_name(type_) + return create_model(type_name, __root__=(type_, ...)) + + +T = TypeVar('T') + + +def parse_obj_as(type_: Type[T], obj: Any, *, type_name: Optional[NameFactory] = None) -> T: + model_type = _get_parsing_type(type_, type_name=type_name) # type: ignore[arg-type] + return model_type(__root__=obj).__root__ + + +def parse_file_as( + type_: Type[T], + path: Union[str, Path], + *, + content_type: str = None, + encoding: str = 'utf8', + proto: Protocol = None, + allow_pickle: bool = False, + json_loads: Callable[[str], Any] = json.loads, + type_name: Optional[NameFactory] = None, +) -> T: + obj = load_file( + path, + proto=proto, + content_type=content_type, + encoding=encoding, + allow_pickle=allow_pickle, + json_loads=json_loads, + ) + return parse_obj_as(type_, obj, type_name=type_name) + + +def parse_raw_as( + type_: Type[T], + b: StrBytes, + *, + content_type: str = None, + encoding: str = 'utf8', + proto: Protocol = None, + allow_pickle: bool = False, + json_loads: Callable[[str], Any] = json.loads, + type_name: Optional[NameFactory] = None, +) -> T: + obj = load_str_bytes( + b, + proto=proto, + content_type=content_type, + encoding=encoding, + allow_pickle=allow_pickle, + json_loads=json_loads, + ) + return parse_obj_as(type_, obj, type_name=type_name) + + +def schema_of(type_: Any, *, title: Optional[NameFactory] = None, **schema_kwargs: Any) -> 'DictStrAny': + """Generate a JSON schema (as dict) for the passed model or dynamically generated one""" + return _get_parsing_type(type_, type_name=title).schema(**schema_kwargs) + + +def schema_json_of(type_: Any, *, title: Optional[NameFactory] = None, **schema_json_kwargs: Any) -> str: + """Generate a JSON schema (as JSON) for the passed model or dynamically generated one""" + return _get_parsing_type(type_, type_name=title).schema_json(**schema_json_kwargs) diff --git a/lib/python3.12/site-packages/pydantic/v1/types.py b/lib/python3.12/site-packages/pydantic/v1/types.py new file mode 100644 index 0000000000000000000000000000000000000000..0cd789a44561ec6885a67b5cc975bdb0cf2f55c1 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/types.py @@ -0,0 +1,1205 @@ +import abc +import math +import re +import warnings +from datetime import date +from decimal import Decimal, InvalidOperation +from enum import Enum +from pathlib import Path +from types import new_class +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ClassVar, + Dict, + FrozenSet, + List, + Optional, + Pattern, + Set, + Tuple, + Type, + TypeVar, + Union, + cast, + overload, +) +from uuid import UUID +from weakref import WeakSet + +from pydantic.v1 import errors +from pydantic.v1.datetime_parse import parse_date +from pydantic.v1.utils import import_string, update_not_none +from pydantic.v1.validators import ( + bytes_validator, + constr_length_validator, + constr_lower, + constr_strip_whitespace, + constr_upper, + decimal_validator, + float_finite_validator, + float_validator, + frozenset_validator, + int_validator, + list_validator, + number_multiple_validator, + number_size_validator, + path_exists_validator, + path_validator, + set_validator, + str_validator, + strict_bytes_validator, + strict_float_validator, + strict_int_validator, + strict_str_validator, +) + +__all__ = [ + 'NoneStr', + 'NoneBytes', + 'StrBytes', + 'NoneStrBytes', + 'StrictStr', + 'ConstrainedBytes', + 'conbytes', + 'ConstrainedList', + 'conlist', + 'ConstrainedSet', + 'conset', + 'ConstrainedFrozenSet', + 'confrozenset', + 'ConstrainedStr', + 'constr', + 'PyObject', + 'ConstrainedInt', + 'conint', + 'PositiveInt', + 'NegativeInt', + 'NonNegativeInt', + 'NonPositiveInt', + 'ConstrainedFloat', + 'confloat', + 'PositiveFloat', + 'NegativeFloat', + 'NonNegativeFloat', + 'NonPositiveFloat', + 'FiniteFloat', + 'ConstrainedDecimal', + 'condecimal', + 'UUID1', + 'UUID3', + 'UUID4', + 'UUID5', + 'FilePath', + 'DirectoryPath', + 'Json', + 'JsonWrapper', + 'SecretField', + 'SecretStr', + 'SecretBytes', + 'StrictBool', + 'StrictBytes', + 'StrictInt', + 'StrictFloat', + 'PaymentCardNumber', + 'ByteSize', + 'PastDate', + 'FutureDate', + 'ConstrainedDate', + 'condate', +] + +NoneStr = Optional[str] +NoneBytes = Optional[bytes] +StrBytes = Union[str, bytes] +NoneStrBytes = Optional[StrBytes] +OptionalInt = Optional[int] +OptionalIntFloat = Union[OptionalInt, float] +OptionalIntFloatDecimal = Union[OptionalIntFloat, Decimal] +OptionalDate = Optional[date] +StrIntFloat = Union[str, int, float] + +if TYPE_CHECKING: + from typing_extensions import Annotated + + from pydantic.v1.dataclasses import Dataclass + from pydantic.v1.main import BaseModel + from pydantic.v1.typing import CallableGenerator + + ModelOrDc = Type[Union[BaseModel, Dataclass]] + +T = TypeVar('T') +_DEFINED_TYPES: 'WeakSet[type]' = WeakSet() + + +@overload +def _registered(typ: Type[T]) -> Type[T]: + pass + + +@overload +def _registered(typ: 'ConstrainedNumberMeta') -> 'ConstrainedNumberMeta': + pass + + +def _registered(typ: Union[Type[T], 'ConstrainedNumberMeta']) -> Union[Type[T], 'ConstrainedNumberMeta']: + # In order to generate valid examples of constrained types, Hypothesis needs + # to inspect the type object - so we keep a weakref to each contype object + # until it can be registered. When (or if) our Hypothesis plugin is loaded, + # it monkeypatches this function. + # If Hypothesis is never used, the total effect is to keep a weak reference + # which has minimal memory usage and doesn't even affect garbage collection. + _DEFINED_TYPES.add(typ) + return typ + + +class ConstrainedNumberMeta(type): + def __new__(cls, name: str, bases: Any, dct: Dict[str, Any]) -> 'ConstrainedInt': # type: ignore + new_cls = cast('ConstrainedInt', type.__new__(cls, name, bases, dct)) + + if new_cls.gt is not None and new_cls.ge is not None: + raise errors.ConfigError('bounds gt and ge cannot be specified at the same time') + if new_cls.lt is not None and new_cls.le is not None: + raise errors.ConfigError('bounds lt and le cannot be specified at the same time') + + return _registered(new_cls) # type: ignore + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ BOOLEAN TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +if TYPE_CHECKING: + StrictBool = bool +else: + + class StrictBool(int): + """ + StrictBool to allow for bools which are not type-coerced. + """ + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='boolean') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + + @classmethod + def validate(cls, value: Any) -> bool: + """ + Ensure that we only allow bools. + """ + if isinstance(value, bool): + return value + + raise errors.StrictBoolError() + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ INTEGER TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class ConstrainedInt(int, metaclass=ConstrainedNumberMeta): + strict: bool = False + gt: OptionalInt = None + ge: OptionalInt = None + lt: OptionalInt = None + le: OptionalInt = None + multiple_of: OptionalInt = None + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none( + field_schema, + exclusiveMinimum=cls.gt, + exclusiveMaximum=cls.lt, + minimum=cls.ge, + maximum=cls.le, + multipleOf=cls.multiple_of, + ) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield strict_int_validator if cls.strict else int_validator + yield number_size_validator + yield number_multiple_validator + + +def conint( + *, + strict: bool = False, + gt: Optional[int] = None, + ge: Optional[int] = None, + lt: Optional[int] = None, + le: Optional[int] = None, + multiple_of: Optional[int] = None, +) -> Type[int]: + # use kwargs then define conf in a dict to aid with IDE type hinting + namespace = dict(strict=strict, gt=gt, ge=ge, lt=lt, le=le, multiple_of=multiple_of) + return type('ConstrainedIntValue', (ConstrainedInt,), namespace) + + +if TYPE_CHECKING: + PositiveInt = int + NegativeInt = int + NonPositiveInt = int + NonNegativeInt = int + StrictInt = int +else: + + class PositiveInt(ConstrainedInt): + gt = 0 + + class NegativeInt(ConstrainedInt): + lt = 0 + + class NonPositiveInt(ConstrainedInt): + le = 0 + + class NonNegativeInt(ConstrainedInt): + ge = 0 + + class StrictInt(ConstrainedInt): + strict = True + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ FLOAT TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class ConstrainedFloat(float, metaclass=ConstrainedNumberMeta): + strict: bool = False + gt: OptionalIntFloat = None + ge: OptionalIntFloat = None + lt: OptionalIntFloat = None + le: OptionalIntFloat = None + multiple_of: OptionalIntFloat = None + allow_inf_nan: Optional[bool] = None + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none( + field_schema, + exclusiveMinimum=cls.gt, + exclusiveMaximum=cls.lt, + minimum=cls.ge, + maximum=cls.le, + multipleOf=cls.multiple_of, + ) + # Modify constraints to account for differences between IEEE floats and JSON + if field_schema.get('exclusiveMinimum') == -math.inf: + del field_schema['exclusiveMinimum'] + if field_schema.get('minimum') == -math.inf: + del field_schema['minimum'] + if field_schema.get('exclusiveMaximum') == math.inf: + del field_schema['exclusiveMaximum'] + if field_schema.get('maximum') == math.inf: + del field_schema['maximum'] + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield strict_float_validator if cls.strict else float_validator + yield number_size_validator + yield number_multiple_validator + yield float_finite_validator + + +def confloat( + *, + strict: bool = False, + gt: float = None, + ge: float = None, + lt: float = None, + le: float = None, + multiple_of: float = None, + allow_inf_nan: Optional[bool] = None, +) -> Type[float]: + # use kwargs then define conf in a dict to aid with IDE type hinting + namespace = dict(strict=strict, gt=gt, ge=ge, lt=lt, le=le, multiple_of=multiple_of, allow_inf_nan=allow_inf_nan) + return type('ConstrainedFloatValue', (ConstrainedFloat,), namespace) + + +if TYPE_CHECKING: + PositiveFloat = float + NegativeFloat = float + NonPositiveFloat = float + NonNegativeFloat = float + StrictFloat = float + FiniteFloat = float +else: + + class PositiveFloat(ConstrainedFloat): + gt = 0 + + class NegativeFloat(ConstrainedFloat): + lt = 0 + + class NonPositiveFloat(ConstrainedFloat): + le = 0 + + class NonNegativeFloat(ConstrainedFloat): + ge = 0 + + class StrictFloat(ConstrainedFloat): + strict = True + + class FiniteFloat(ConstrainedFloat): + allow_inf_nan = False + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ BYTES TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class ConstrainedBytes(bytes): + strip_whitespace = False + to_upper = False + to_lower = False + min_length: OptionalInt = None + max_length: OptionalInt = None + strict: bool = False + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none(field_schema, minLength=cls.min_length, maxLength=cls.max_length) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield strict_bytes_validator if cls.strict else bytes_validator + yield constr_strip_whitespace + yield constr_upper + yield constr_lower + yield constr_length_validator + + +def conbytes( + *, + strip_whitespace: bool = False, + to_upper: bool = False, + to_lower: bool = False, + min_length: Optional[int] = None, + max_length: Optional[int] = None, + strict: bool = False, +) -> Type[bytes]: + # use kwargs then define conf in a dict to aid with IDE type hinting + namespace = dict( + strip_whitespace=strip_whitespace, + to_upper=to_upper, + to_lower=to_lower, + min_length=min_length, + max_length=max_length, + strict=strict, + ) + return _registered(type('ConstrainedBytesValue', (ConstrainedBytes,), namespace)) + + +if TYPE_CHECKING: + StrictBytes = bytes +else: + + class StrictBytes(ConstrainedBytes): + strict = True + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ STRING TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class ConstrainedStr(str): + strip_whitespace = False + to_upper = False + to_lower = False + min_length: OptionalInt = None + max_length: OptionalInt = None + curtail_length: OptionalInt = None + regex: Optional[Union[str, Pattern[str]]] = None + strict = False + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none( + field_schema, + minLength=cls.min_length, + maxLength=cls.max_length, + pattern=cls.regex and cls._get_pattern(cls.regex), + ) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield strict_str_validator if cls.strict else str_validator + yield constr_strip_whitespace + yield constr_upper + yield constr_lower + yield constr_length_validator + yield cls.validate + + @classmethod + def validate(cls, value: Union[str]) -> Union[str]: + if cls.curtail_length and len(value) > cls.curtail_length: + value = value[: cls.curtail_length] + + if cls.regex: + if not re.match(cls.regex, value): + raise errors.StrRegexError(pattern=cls._get_pattern(cls.regex)) + + return value + + @staticmethod + def _get_pattern(regex: Union[str, Pattern[str]]) -> str: + return regex if isinstance(regex, str) else regex.pattern + + +def constr( + *, + strip_whitespace: bool = False, + to_upper: bool = False, + to_lower: bool = False, + strict: bool = False, + min_length: Optional[int] = None, + max_length: Optional[int] = None, + curtail_length: Optional[int] = None, + regex: Optional[str] = None, +) -> Type[str]: + # use kwargs then define conf in a dict to aid with IDE type hinting + namespace = dict( + strip_whitespace=strip_whitespace, + to_upper=to_upper, + to_lower=to_lower, + strict=strict, + min_length=min_length, + max_length=max_length, + curtail_length=curtail_length, + regex=regex and re.compile(regex), + ) + return _registered(type('ConstrainedStrValue', (ConstrainedStr,), namespace)) + + +if TYPE_CHECKING: + StrictStr = str +else: + + class StrictStr(ConstrainedStr): + strict = True + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ SET TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +# This types superclass should be Set[T], but cython chokes on that... +class ConstrainedSet(set): # type: ignore + # Needed for pydantic to detect that this is a set + __origin__ = set + __args__: Set[Type[T]] # type: ignore + + min_items: Optional[int] = None + max_items: Optional[int] = None + item_type: Type[T] # type: ignore + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.set_length_validator + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none(field_schema, minItems=cls.min_items, maxItems=cls.max_items) + + @classmethod + def set_length_validator(cls, v: 'Optional[Set[T]]') -> 'Optional[Set[T]]': + if v is None: + return None + + v = set_validator(v) + v_len = len(v) + + if cls.min_items is not None and v_len < cls.min_items: + raise errors.SetMinLengthError(limit_value=cls.min_items) + + if cls.max_items is not None and v_len > cls.max_items: + raise errors.SetMaxLengthError(limit_value=cls.max_items) + + return v + + +def conset(item_type: Type[T], *, min_items: Optional[int] = None, max_items: Optional[int] = None) -> Type[Set[T]]: + # __args__ is needed to conform to typing generics api + namespace = {'min_items': min_items, 'max_items': max_items, 'item_type': item_type, '__args__': [item_type]} + # We use new_class to be able to deal with Generic types + return new_class('ConstrainedSetValue', (ConstrainedSet,), {}, lambda ns: ns.update(namespace)) + + +# This types superclass should be FrozenSet[T], but cython chokes on that... +class ConstrainedFrozenSet(frozenset): # type: ignore + # Needed for pydantic to detect that this is a set + __origin__ = frozenset + __args__: FrozenSet[Type[T]] # type: ignore + + min_items: Optional[int] = None + max_items: Optional[int] = None + item_type: Type[T] # type: ignore + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.frozenset_length_validator + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none(field_schema, minItems=cls.min_items, maxItems=cls.max_items) + + @classmethod + def frozenset_length_validator(cls, v: 'Optional[FrozenSet[T]]') -> 'Optional[FrozenSet[T]]': + if v is None: + return None + + v = frozenset_validator(v) + v_len = len(v) + + if cls.min_items is not None and v_len < cls.min_items: + raise errors.FrozenSetMinLengthError(limit_value=cls.min_items) + + if cls.max_items is not None and v_len > cls.max_items: + raise errors.FrozenSetMaxLengthError(limit_value=cls.max_items) + + return v + + +def confrozenset( + item_type: Type[T], *, min_items: Optional[int] = None, max_items: Optional[int] = None +) -> Type[FrozenSet[T]]: + # __args__ is needed to conform to typing generics api + namespace = {'min_items': min_items, 'max_items': max_items, 'item_type': item_type, '__args__': [item_type]} + # We use new_class to be able to deal with Generic types + return new_class('ConstrainedFrozenSetValue', (ConstrainedFrozenSet,), {}, lambda ns: ns.update(namespace)) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ LIST TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +# This types superclass should be List[T], but cython chokes on that... +class ConstrainedList(list): # type: ignore + # Needed for pydantic to detect that this is a list + __origin__ = list + __args__: Tuple[Type[T], ...] # type: ignore + + min_items: Optional[int] = None + max_items: Optional[int] = None + unique_items: Optional[bool] = None + item_type: Type[T] # type: ignore + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.list_length_validator + if cls.unique_items: + yield cls.unique_items_validator + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none(field_schema, minItems=cls.min_items, maxItems=cls.max_items, uniqueItems=cls.unique_items) + + @classmethod + def list_length_validator(cls, v: 'Optional[List[T]]') -> 'Optional[List[T]]': + if v is None: + return None + + v = list_validator(v) + v_len = len(v) + + if cls.min_items is not None and v_len < cls.min_items: + raise errors.ListMinLengthError(limit_value=cls.min_items) + + if cls.max_items is not None and v_len > cls.max_items: + raise errors.ListMaxLengthError(limit_value=cls.max_items) + + return v + + @classmethod + def unique_items_validator(cls, v: 'Optional[List[T]]') -> 'Optional[List[T]]': + if v is None: + return None + + for i, value in enumerate(v, start=1): + if value in v[i:]: + raise errors.ListUniqueItemsError() + + return v + + +def conlist( + item_type: Type[T], *, min_items: Optional[int] = None, max_items: Optional[int] = None, unique_items: bool = None +) -> Type[List[T]]: + # __args__ is needed to conform to typing generics api + namespace = dict( + min_items=min_items, max_items=max_items, unique_items=unique_items, item_type=item_type, __args__=(item_type,) + ) + # We use new_class to be able to deal with Generic types + return new_class('ConstrainedListValue', (ConstrainedList,), {}, lambda ns: ns.update(namespace)) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ PYOBJECT TYPE ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +if TYPE_CHECKING: + PyObject = Callable[..., Any] +else: + + class PyObject: + validate_always = True + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + + @classmethod + def validate(cls, value: Any) -> Any: + if isinstance(value, Callable): + return value + + try: + value = str_validator(value) + except errors.StrError: + raise errors.PyObjectError(error_message='value is neither a valid import path not a valid callable') + + try: + return import_string(value) + except ImportError as e: + raise errors.PyObjectError(error_message=str(e)) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ DECIMAL TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class ConstrainedDecimal(Decimal, metaclass=ConstrainedNumberMeta): + gt: OptionalIntFloatDecimal = None + ge: OptionalIntFloatDecimal = None + lt: OptionalIntFloatDecimal = None + le: OptionalIntFloatDecimal = None + max_digits: OptionalInt = None + decimal_places: OptionalInt = None + multiple_of: OptionalIntFloatDecimal = None + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none( + field_schema, + exclusiveMinimum=cls.gt, + exclusiveMaximum=cls.lt, + minimum=cls.ge, + maximum=cls.le, + multipleOf=cls.multiple_of, + ) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield decimal_validator + yield number_size_validator + yield number_multiple_validator + yield cls.validate + + @classmethod + def validate(cls, value: Decimal) -> Decimal: + try: + normalized_value = value.normalize() + except InvalidOperation: + normalized_value = value + digit_tuple, exponent = normalized_value.as_tuple()[1:] + if exponent in {'F', 'n', 'N'}: + raise errors.DecimalIsNotFiniteError() + + if exponent >= 0: + # A positive exponent adds that many trailing zeros. + digits = len(digit_tuple) + exponent + decimals = 0 + else: + # If the absolute value of the negative exponent is larger than the + # number of digits, then it's the same as the number of digits, + # because it'll consume all of the digits in digit_tuple and then + # add abs(exponent) - len(digit_tuple) leading zeros after the + # decimal point. + if abs(exponent) > len(digit_tuple): + digits = decimals = abs(exponent) + else: + digits = len(digit_tuple) + decimals = abs(exponent) + whole_digits = digits - decimals + + if cls.max_digits is not None and digits > cls.max_digits: + raise errors.DecimalMaxDigitsError(max_digits=cls.max_digits) + + if cls.decimal_places is not None and decimals > cls.decimal_places: + raise errors.DecimalMaxPlacesError(decimal_places=cls.decimal_places) + + if cls.max_digits is not None and cls.decimal_places is not None: + expected = cls.max_digits - cls.decimal_places + if whole_digits > expected: + raise errors.DecimalWholeDigitsError(whole_digits=expected) + + return value + + +def condecimal( + *, + gt: Decimal = None, + ge: Decimal = None, + lt: Decimal = None, + le: Decimal = None, + max_digits: Optional[int] = None, + decimal_places: Optional[int] = None, + multiple_of: Decimal = None, +) -> Type[Decimal]: + # use kwargs then define conf in a dict to aid with IDE type hinting + namespace = dict( + gt=gt, ge=ge, lt=lt, le=le, max_digits=max_digits, decimal_places=decimal_places, multiple_of=multiple_of + ) + return type('ConstrainedDecimalValue', (ConstrainedDecimal,), namespace) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ UUID TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +if TYPE_CHECKING: + UUID1 = UUID + UUID3 = UUID + UUID4 = UUID + UUID5 = UUID +else: + + class UUID1(UUID): + _required_version = 1 + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='string', format=f'uuid{cls._required_version}') + + class UUID3(UUID1): + _required_version = 3 + + class UUID4(UUID1): + _required_version = 4 + + class UUID5(UUID1): + _required_version = 5 + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ PATH TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +if TYPE_CHECKING: + FilePath = Path + DirectoryPath = Path +else: + + class FilePath(Path): + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(format='file-path') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield path_validator + yield path_exists_validator + yield cls.validate + + @classmethod + def validate(cls, value: Path) -> Path: + if not value.is_file(): + raise errors.PathNotAFileError(path=value) + + return value + + class DirectoryPath(Path): + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(format='directory-path') + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield path_validator + yield path_exists_validator + yield cls.validate + + @classmethod + def validate(cls, value: Path) -> Path: + if not value.is_dir(): + raise errors.PathNotADirectoryError(path=value) + + return value + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ JSON TYPE ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class JsonWrapper: + pass + + +class JsonMeta(type): + def __getitem__(self, t: Type[Any]) -> Type[JsonWrapper]: + if t is Any: + return Json # allow Json[Any] to replecate plain Json + return _registered(type('JsonWrapperValue', (JsonWrapper,), {'inner_type': t})) + + +if TYPE_CHECKING: + Json = Annotated[T, ...] # Json[list[str]] will be recognized by type checkers as list[str] + +else: + + class Json(metaclass=JsonMeta): + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + field_schema.update(type='string', format='json-string') + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ SECRET TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class SecretField(abc.ABC): + """ + Note: this should be implemented as a generic like `SecretField(ABC, Generic[T])`, + the `__init__()` should be part of the abstract class and the + `get_secret_value()` method should use the generic `T` type. + + However Cython doesn't support very well generics at the moment and + the generated code fails to be imported (see + https://github.com/cython/cython/issues/2753). + """ + + def __eq__(self, other: Any) -> bool: + return isinstance(other, self.__class__) and self.get_secret_value() == other.get_secret_value() + + def __str__(self) -> str: + return '**********' if self.get_secret_value() else '' + + def __hash__(self) -> int: + return hash(self.get_secret_value()) + + @abc.abstractmethod + def get_secret_value(self) -> Any: # pragma: no cover + ... + + +class SecretStr(SecretField): + min_length: OptionalInt = None + max_length: OptionalInt = None + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none( + field_schema, + type='string', + writeOnly=True, + format='password', + minLength=cls.min_length, + maxLength=cls.max_length, + ) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + yield constr_length_validator + + @classmethod + def validate(cls, value: Any) -> 'SecretStr': + if isinstance(value, cls): + return value + value = str_validator(value) + return cls(value) + + def __init__(self, value: str): + self._secret_value = value + + def __repr__(self) -> str: + return f"SecretStr('{self}')" + + def __len__(self) -> int: + return len(self._secret_value) + + def display(self) -> str: + warnings.warn('`secret_str.display()` is deprecated, use `str(secret_str)` instead', DeprecationWarning) + return str(self) + + def get_secret_value(self) -> str: + return self._secret_value + + +class SecretBytes(SecretField): + min_length: OptionalInt = None + max_length: OptionalInt = None + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none( + field_schema, + type='string', + writeOnly=True, + format='password', + minLength=cls.min_length, + maxLength=cls.max_length, + ) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + yield constr_length_validator + + @classmethod + def validate(cls, value: Any) -> 'SecretBytes': + if isinstance(value, cls): + return value + value = bytes_validator(value) + return cls(value) + + def __init__(self, value: bytes): + self._secret_value = value + + def __repr__(self) -> str: + return f"SecretBytes(b'{self}')" + + def __len__(self) -> int: + return len(self._secret_value) + + def display(self) -> str: + warnings.warn('`secret_bytes.display()` is deprecated, use `str(secret_bytes)` instead', DeprecationWarning) + return str(self) + + def get_secret_value(self) -> bytes: + return self._secret_value + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ PAYMENT CARD TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + +class PaymentCardBrand(str, Enum): + # If you add another card type, please also add it to the + # Hypothesis strategy in `pydantic._hypothesis_plugin`. + amex = 'American Express' + mastercard = 'Mastercard' + visa = 'Visa' + other = 'other' + + def __str__(self) -> str: + return self.value + + +class PaymentCardNumber(str): + """ + Based on: https://en.wikipedia.org/wiki/Payment_card_number + """ + + strip_whitespace: ClassVar[bool] = True + min_length: ClassVar[int] = 12 + max_length: ClassVar[int] = 19 + bin: str + last4: str + brand: PaymentCardBrand + + def __init__(self, card_number: str): + self.bin = card_number[:6] + self.last4 = card_number[-4:] + self.brand = self._get_brand(card_number) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield str_validator + yield constr_strip_whitespace + yield constr_length_validator + yield cls.validate_digits + yield cls.validate_luhn_check_digit + yield cls + yield cls.validate_length_for_brand + + @property + def masked(self) -> str: + num_masked = len(self) - 10 # len(bin) + len(last4) == 10 + return f'{self.bin}{"*" * num_masked}{self.last4}' + + @classmethod + def validate_digits(cls, card_number: str) -> str: + if not card_number.isdigit(): + raise errors.NotDigitError + return card_number + + @classmethod + def validate_luhn_check_digit(cls, card_number: str) -> str: + """ + Based on: https://en.wikipedia.org/wiki/Luhn_algorithm + """ + sum_ = int(card_number[-1]) + length = len(card_number) + parity = length % 2 + for i in range(length - 1): + digit = int(card_number[i]) + if i % 2 == parity: + digit *= 2 + if digit > 9: + digit -= 9 + sum_ += digit + valid = sum_ % 10 == 0 + if not valid: + raise errors.LuhnValidationError + return card_number + + @classmethod + def validate_length_for_brand(cls, card_number: 'PaymentCardNumber') -> 'PaymentCardNumber': + """ + Validate length based on BIN for major brands: + https://en.wikipedia.org/wiki/Payment_card_number#Issuer_identification_number_(IIN) + """ + required_length: Union[None, int, str] = None + if card_number.brand in PaymentCardBrand.mastercard: + required_length = 16 + valid = len(card_number) == required_length + elif card_number.brand == PaymentCardBrand.visa: + required_length = '13, 16 or 19' + valid = len(card_number) in {13, 16, 19} + elif card_number.brand == PaymentCardBrand.amex: + required_length = 15 + valid = len(card_number) == required_length + else: + valid = True + if not valid: + raise errors.InvalidLengthForBrand(brand=card_number.brand, required_length=required_length) + return card_number + + @staticmethod + def _get_brand(card_number: str) -> PaymentCardBrand: + if card_number[0] == '4': + brand = PaymentCardBrand.visa + elif 51 <= int(card_number[:2]) <= 55: + brand = PaymentCardBrand.mastercard + elif card_number[:2] in {'34', '37'}: + brand = PaymentCardBrand.amex + else: + brand = PaymentCardBrand.other + return brand + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ BYTE SIZE TYPE ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +BYTE_SIZES = { + 'b': 1, + 'kb': 10**3, + 'mb': 10**6, + 'gb': 10**9, + 'tb': 10**12, + 'pb': 10**15, + 'eb': 10**18, + 'kib': 2**10, + 'mib': 2**20, + 'gib': 2**30, + 'tib': 2**40, + 'pib': 2**50, + 'eib': 2**60, +} +BYTE_SIZES.update({k.lower()[0]: v for k, v in BYTE_SIZES.items() if 'i' not in k}) +byte_string_re = re.compile(r'^\s*(\d*\.?\d+)\s*(\w+)?', re.IGNORECASE) + + +class ByteSize(int): + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield cls.validate + + @classmethod + def validate(cls, v: StrIntFloat) -> 'ByteSize': + try: + return cls(int(v)) + except ValueError: + pass + + str_match = byte_string_re.match(str(v)) + if str_match is None: + raise errors.InvalidByteSize() + + scalar, unit = str_match.groups() + if unit is None: + unit = 'b' + + try: + unit_mult = BYTE_SIZES[unit.lower()] + except KeyError: + raise errors.InvalidByteSizeUnit(unit=unit) + + return cls(int(float(scalar) * unit_mult)) + + def human_readable(self, decimal: bool = False) -> str: + if decimal: + divisor = 1000 + units = ['B', 'KB', 'MB', 'GB', 'TB', 'PB'] + final_unit = 'EB' + else: + divisor = 1024 + units = ['B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB'] + final_unit = 'EiB' + + num = float(self) + for unit in units: + if abs(num) < divisor: + return f'{num:0.1f}{unit}' + num /= divisor + + return f'{num:0.1f}{final_unit}' + + def to(self, unit: str) -> float: + try: + unit_div = BYTE_SIZES[unit.lower()] + except KeyError: + raise errors.InvalidByteSizeUnit(unit=unit) + + return self / unit_div + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ DATE TYPES ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +if TYPE_CHECKING: + PastDate = date + FutureDate = date +else: + + class PastDate(date): + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield parse_date + yield cls.validate + + @classmethod + def validate(cls, value: date) -> date: + if value >= date.today(): + raise errors.DateNotInThePastError() + + return value + + class FutureDate(date): + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield parse_date + yield cls.validate + + @classmethod + def validate(cls, value: date) -> date: + if value <= date.today(): + raise errors.DateNotInTheFutureError() + + return value + + +class ConstrainedDate(date, metaclass=ConstrainedNumberMeta): + gt: OptionalDate = None + ge: OptionalDate = None + lt: OptionalDate = None + le: OptionalDate = None + + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + update_not_none(field_schema, exclusiveMinimum=cls.gt, exclusiveMaximum=cls.lt, minimum=cls.ge, maximum=cls.le) + + @classmethod + def __get_validators__(cls) -> 'CallableGenerator': + yield parse_date + yield number_size_validator + + +def condate( + *, + gt: date = None, + ge: date = None, + lt: date = None, + le: date = None, +) -> Type[date]: + # use kwargs then define conf in a dict to aid with IDE type hinting + namespace = dict(gt=gt, ge=ge, lt=lt, le=le) + return type('ConstrainedDateValue', (ConstrainedDate,), namespace) diff --git a/lib/python3.12/site-packages/pydantic/v1/typing.py b/lib/python3.12/site-packages/pydantic/v1/typing.py new file mode 100644 index 0000000000000000000000000000000000000000..3038ccda4c8de6ebaf1b8164a30ef7323ac07e65 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/typing.py @@ -0,0 +1,615 @@ +import sys +import typing +from collections.abc import Callable +from os import PathLike +from typing import ( # type: ignore + TYPE_CHECKING, + AbstractSet, + Any, + Callable as TypingCallable, + ClassVar, + Dict, + ForwardRef, + Generator, + Iterable, + List, + Mapping, + NewType, + Optional, + Sequence, + Set, + Tuple, + Type, + TypeVar, + Union, + _eval_type, + cast, + get_type_hints, +) + +from typing_extensions import ( + Annotated, + Final, + Literal, + NotRequired as TypedDictNotRequired, + Required as TypedDictRequired, +) + +try: + from typing import _TypingBase as typing_base # type: ignore +except ImportError: + from typing import _Final as typing_base # type: ignore + +try: + from typing import GenericAlias as TypingGenericAlias # type: ignore +except ImportError: + # python < 3.9 does not have GenericAlias (list[int], tuple[str, ...] and so on) + TypingGenericAlias = () + +try: + from types import UnionType as TypesUnionType # type: ignore +except ImportError: + # python < 3.10 does not have UnionType (str | int, byte | bool and so on) + TypesUnionType = () + + +if sys.version_info < (3, 9): + + def evaluate_forwardref(type_: ForwardRef, globalns: Any, localns: Any) -> Any: + return type_._evaluate(globalns, localns) + +elif sys.version_info < (3, 12, 4): + + def evaluate_forwardref(type_: ForwardRef, globalns: Any, localns: Any) -> Any: + # Even though it is the right signature for python 3.9, mypy complains with + # `error: Too many arguments for "_evaluate" of "ForwardRef"` hence the cast... + # Python 3.13/3.12.4+ made `recursive_guard` a kwarg, so name it explicitly to avoid: + # TypeError: ForwardRef._evaluate() missing 1 required keyword-only argument: 'recursive_guard' + return cast(Any, type_)._evaluate(globalns, localns, recursive_guard=set()) + +else: + + def evaluate_forwardref(type_: ForwardRef, globalns: Any, localns: Any) -> Any: + # Pydantic 1.x will not support PEP 695 syntax, but provide `type_params` to avoid + # warnings: + return cast(Any, type_)._evaluate(globalns, localns, type_params=(), recursive_guard=set()) + + +if sys.version_info < (3, 9): + # Ensure we always get all the whole `Annotated` hint, not just the annotated type. + # For 3.7 to 3.8, `get_type_hints` doesn't recognize `typing_extensions.Annotated`, + # so it already returns the full annotation + get_all_type_hints = get_type_hints + +else: + + def get_all_type_hints(obj: Any, globalns: Any = None, localns: Any = None) -> Any: + return get_type_hints(obj, globalns, localns, include_extras=True) + + +_T = TypeVar('_T') + +AnyCallable = TypingCallable[..., Any] +NoArgAnyCallable = TypingCallable[[], Any] + +# workaround for https://github.com/python/mypy/issues/9496 +AnyArgTCallable = TypingCallable[..., _T] + + +# Annotated[...] is implemented by returning an instance of one of these classes, depending on +# python/typing_extensions version. +AnnotatedTypeNames = {'AnnotatedMeta', '_AnnotatedAlias'} + + +LITERAL_TYPES: Set[Any] = {Literal} +if hasattr(typing, 'Literal'): + LITERAL_TYPES.add(typing.Literal) + + +if sys.version_info < (3, 8): + + def get_origin(t: Type[Any]) -> Optional[Type[Any]]: + if type(t).__name__ in AnnotatedTypeNames: + # weirdly this is a runtime requirement, as well as for mypy + return cast(Type[Any], Annotated) + return getattr(t, '__origin__', None) + +else: + from typing import get_origin as _typing_get_origin + + def get_origin(tp: Type[Any]) -> Optional[Type[Any]]: + """ + We can't directly use `typing.get_origin` since we need a fallback to support + custom generic classes like `ConstrainedList` + It should be useless once https://github.com/cython/cython/issues/3537 is + solved and https://github.com/pydantic/pydantic/pull/1753 is merged. + """ + if type(tp).__name__ in AnnotatedTypeNames: + return cast(Type[Any], Annotated) # mypy complains about _SpecialForm + return _typing_get_origin(tp) or getattr(tp, '__origin__', None) + + +if sys.version_info < (3, 8): + from typing import _GenericAlias + + def get_args(t: Type[Any]) -> Tuple[Any, ...]: + """Compatibility version of get_args for python 3.7. + + Mostly compatible with the python 3.8 `typing` module version + and able to handle almost all use cases. + """ + if type(t).__name__ in AnnotatedTypeNames: + return t.__args__ + t.__metadata__ + if isinstance(t, _GenericAlias): + res = t.__args__ + if t.__origin__ is Callable and res and res[0] is not Ellipsis: + res = (list(res[:-1]), res[-1]) + return res + return getattr(t, '__args__', ()) + +else: + from typing import get_args as _typing_get_args + + def _generic_get_args(tp: Type[Any]) -> Tuple[Any, ...]: + """ + In python 3.9, `typing.Dict`, `typing.List`, ... + do have an empty `__args__` by default (instead of the generic ~T for example). + In order to still support `Dict` for example and consider it as `Dict[Any, Any]`, + we retrieve the `_nparams` value that tells us how many parameters it needs. + """ + if hasattr(tp, '_nparams'): + return (Any,) * tp._nparams + # Special case for `tuple[()]`, which used to return ((),) with `typing.Tuple` + # in python 3.10- but now returns () for `tuple` and `Tuple`. + # This will probably be clarified in pydantic v2 + try: + if tp == Tuple[()] or sys.version_info >= (3, 9) and tp == tuple[()]: # type: ignore[misc] + return ((),) + # there is a TypeError when compiled with cython + except TypeError: # pragma: no cover + pass + return () + + def get_args(tp: Type[Any]) -> Tuple[Any, ...]: + """Get type arguments with all substitutions performed. + + For unions, basic simplifications used by Union constructor are performed. + Examples:: + get_args(Dict[str, int]) == (str, int) + get_args(int) == () + get_args(Union[int, Union[T, int], str][int]) == (int, str) + get_args(Union[int, Tuple[T, int]][str]) == (int, Tuple[str, int]) + get_args(Callable[[], T][int]) == ([], int) + """ + if type(tp).__name__ in AnnotatedTypeNames: + return tp.__args__ + tp.__metadata__ + # the fallback is needed for the same reasons as `get_origin` (see above) + return _typing_get_args(tp) or getattr(tp, '__args__', ()) or _generic_get_args(tp) + + +if sys.version_info < (3, 9): + + def convert_generics(tp: Type[Any]) -> Type[Any]: + """Python 3.9 and older only supports generics from `typing` module. + They convert strings to ForwardRef automatically. + + Examples:: + typing.List['Hero'] == typing.List[ForwardRef('Hero')] + """ + return tp + +else: + from typing import _UnionGenericAlias # type: ignore + + from typing_extensions import _AnnotatedAlias + + def convert_generics(tp: Type[Any]) -> Type[Any]: + """ + Recursively searches for `str` type hints and replaces them with ForwardRef. + + Examples:: + convert_generics(list['Hero']) == list[ForwardRef('Hero')] + convert_generics(dict['Hero', 'Team']) == dict[ForwardRef('Hero'), ForwardRef('Team')] + convert_generics(typing.Dict['Hero', 'Team']) == typing.Dict[ForwardRef('Hero'), ForwardRef('Team')] + convert_generics(list[str | 'Hero'] | int) == list[str | ForwardRef('Hero')] | int + """ + origin = get_origin(tp) + if not origin or not hasattr(tp, '__args__'): + return tp + + args = get_args(tp) + + # typing.Annotated needs special treatment + if origin is Annotated: + return _AnnotatedAlias(convert_generics(args[0]), args[1:]) + + # recursively replace `str` instances inside of `GenericAlias` with `ForwardRef(arg)` + converted = tuple( + ForwardRef(arg) if isinstance(arg, str) and isinstance(tp, TypingGenericAlias) else convert_generics(arg) + for arg in args + ) + + if converted == args: + return tp + elif isinstance(tp, TypingGenericAlias): + return TypingGenericAlias(origin, converted) + elif isinstance(tp, TypesUnionType): + # recreate types.UnionType (PEP604, Python >= 3.10) + return _UnionGenericAlias(origin, converted) + else: + try: + setattr(tp, '__args__', converted) + except AttributeError: + pass + return tp + + +if sys.version_info < (3, 10): + + def is_union(tp: Optional[Type[Any]]) -> bool: + return tp is Union + + WithArgsTypes = (TypingGenericAlias,) + +else: + import types + import typing + + def is_union(tp: Optional[Type[Any]]) -> bool: + return tp is Union or tp is types.UnionType # noqa: E721 + + WithArgsTypes = (typing._GenericAlias, types.GenericAlias, types.UnionType) + + +StrPath = Union[str, PathLike] + + +if TYPE_CHECKING: + from pydantic.v1.fields import ModelField + + TupleGenerator = Generator[Tuple[str, Any], None, None] + DictStrAny = Dict[str, Any] + DictAny = Dict[Any, Any] + SetStr = Set[str] + ListStr = List[str] + IntStr = Union[int, str] + AbstractSetIntStr = AbstractSet[IntStr] + DictIntStrAny = Dict[IntStr, Any] + MappingIntStrAny = Mapping[IntStr, Any] + CallableGenerator = Generator[AnyCallable, None, None] + ReprArgs = Sequence[Tuple[Optional[str], Any]] + + MYPY = False + if MYPY: + AnyClassMethod = classmethod[Any] + else: + # classmethod[TargetType, CallableParamSpecType, CallableReturnType] + AnyClassMethod = classmethod[Any, Any, Any] + +__all__ = ( + 'AnyCallable', + 'NoArgAnyCallable', + 'NoneType', + 'is_none_type', + 'display_as_type', + 'resolve_annotations', + 'is_callable_type', + 'is_literal_type', + 'all_literal_values', + 'is_namedtuple', + 'is_typeddict', + 'is_typeddict_special', + 'is_new_type', + 'new_type_supertype', + 'is_classvar', + 'is_finalvar', + 'update_field_forward_refs', + 'update_model_forward_refs', + 'TupleGenerator', + 'DictStrAny', + 'DictAny', + 'SetStr', + 'ListStr', + 'IntStr', + 'AbstractSetIntStr', + 'DictIntStrAny', + 'CallableGenerator', + 'ReprArgs', + 'AnyClassMethod', + 'CallableGenerator', + 'WithArgsTypes', + 'get_args', + 'get_origin', + 'get_sub_types', + 'typing_base', + 'get_all_type_hints', + 'is_union', + 'StrPath', + 'MappingIntStrAny', +) + + +NoneType = None.__class__ + + +NONE_TYPES: Tuple[Any, Any, Any] = (None, NoneType, Literal[None]) + + +if sys.version_info < (3, 8): + # Even though this implementation is slower, we need it for python 3.7: + # In python 3.7 "Literal" is not a builtin type and uses a different + # mechanism. + # for this reason `Literal[None] is Literal[None]` evaluates to `False`, + # breaking the faster implementation used for the other python versions. + + def is_none_type(type_: Any) -> bool: + return type_ in NONE_TYPES + +elif sys.version_info[:2] == (3, 8): + + def is_none_type(type_: Any) -> bool: + for none_type in NONE_TYPES: + if type_ is none_type: + return True + # With python 3.8, specifically 3.8.10, Literal "is" check sare very flakey + # can change on very subtle changes like use of types in other modules, + # hopefully this check avoids that issue. + if is_literal_type(type_): # pragma: no cover + return all_literal_values(type_) == (None,) + return False + +else: + + def is_none_type(type_: Any) -> bool: + return type_ in NONE_TYPES + + +def display_as_type(v: Type[Any]) -> str: + if not isinstance(v, typing_base) and not isinstance(v, WithArgsTypes) and not isinstance(v, type): + v = v.__class__ + + if is_union(get_origin(v)): + return f'Union[{", ".join(map(display_as_type, get_args(v)))}]' + + if isinstance(v, WithArgsTypes): + # Generic alias are constructs like `list[int]` + return str(v).replace('typing.', '') + + try: + return v.__name__ + except AttributeError: + # happens with typing objects + return str(v).replace('typing.', '') + + +def resolve_annotations(raw_annotations: Dict[str, Type[Any]], module_name: Optional[str]) -> Dict[str, Type[Any]]: + """ + Partially taken from typing.get_type_hints. + + Resolve string or ForwardRef annotations into type objects if possible. + """ + base_globals: Optional[Dict[str, Any]] = None + if module_name: + try: + module = sys.modules[module_name] + except KeyError: + # happens occasionally, see https://github.com/pydantic/pydantic/issues/2363 + pass + else: + base_globals = module.__dict__ + + annotations = {} + for name, value in raw_annotations.items(): + if isinstance(value, str): + if (3, 10) > sys.version_info >= (3, 9, 8) or sys.version_info >= (3, 10, 1): + value = ForwardRef(value, is_argument=False, is_class=True) + else: + value = ForwardRef(value, is_argument=False) + try: + if sys.version_info >= (3, 13): + value = _eval_type(value, base_globals, None, type_params=()) + else: + value = _eval_type(value, base_globals, None) + except NameError: + # this is ok, it can be fixed with update_forward_refs + pass + annotations[name] = value + return annotations + + +def is_callable_type(type_: Type[Any]) -> bool: + return type_ is Callable or get_origin(type_) is Callable + + +def is_literal_type(type_: Type[Any]) -> bool: + return Literal is not None and get_origin(type_) in LITERAL_TYPES + + +def literal_values(type_: Type[Any]) -> Tuple[Any, ...]: + return get_args(type_) + + +def all_literal_values(type_: Type[Any]) -> Tuple[Any, ...]: + """ + This method is used to retrieve all Literal values as + Literal can be used recursively (see https://www.python.org/dev/peps/pep-0586) + e.g. `Literal[Literal[Literal[1, 2, 3], "foo"], 5, None]` + """ + if not is_literal_type(type_): + return (type_,) + + values = literal_values(type_) + return tuple(x for value in values for x in all_literal_values(value)) + + +def is_namedtuple(type_: Type[Any]) -> bool: + """ + Check if a given class is a named tuple. + It can be either a `typing.NamedTuple` or `collections.namedtuple` + """ + from pydantic.v1.utils import lenient_issubclass + + return lenient_issubclass(type_, tuple) and hasattr(type_, '_fields') + + +def is_typeddict(type_: Type[Any]) -> bool: + """ + Check if a given class is a typed dict (from `typing` or `typing_extensions`) + In 3.10, there will be a public method (https://docs.python.org/3.10/library/typing.html#typing.is_typeddict) + """ + from pydantic.v1.utils import lenient_issubclass + + return lenient_issubclass(type_, dict) and hasattr(type_, '__total__') + + +def _check_typeddict_special(type_: Any) -> bool: + return type_ is TypedDictRequired or type_ is TypedDictNotRequired + + +def is_typeddict_special(type_: Any) -> bool: + """ + Check if type is a TypedDict special form (Required or NotRequired). + """ + return _check_typeddict_special(type_) or _check_typeddict_special(get_origin(type_)) + + +test_type = NewType('test_type', str) + + +def is_new_type(type_: Type[Any]) -> bool: + """ + Check whether type_ was created using typing.NewType + """ + return isinstance(type_, test_type.__class__) and hasattr(type_, '__supertype__') # type: ignore + + +def new_type_supertype(type_: Type[Any]) -> Type[Any]: + while hasattr(type_, '__supertype__'): + type_ = type_.__supertype__ + return type_ + + +def _check_classvar(v: Optional[Type[Any]]) -> bool: + if v is None: + return False + + return v.__class__ == ClassVar.__class__ and getattr(v, '_name', None) == 'ClassVar' + + +def _check_finalvar(v: Optional[Type[Any]]) -> bool: + """ + Check if a given type is a `typing.Final` type. + """ + if v is None: + return False + + return v.__class__ == Final.__class__ and (sys.version_info < (3, 8) or getattr(v, '_name', None) == 'Final') + + +def is_classvar(ann_type: Type[Any]) -> bool: + if _check_classvar(ann_type) or _check_classvar(get_origin(ann_type)): + return True + + # this is an ugly workaround for class vars that contain forward references and are therefore themselves + # forward references, see #3679 + if ann_type.__class__ == ForwardRef and ann_type.__forward_arg__.startswith('ClassVar['): + return True + + return False + + +def is_finalvar(ann_type: Type[Any]) -> bool: + return _check_finalvar(ann_type) or _check_finalvar(get_origin(ann_type)) + + +def update_field_forward_refs(field: 'ModelField', globalns: Any, localns: Any) -> None: + """ + Try to update ForwardRefs on fields based on this ModelField, globalns and localns. + """ + prepare = False + if field.type_.__class__ == ForwardRef: + prepare = True + field.type_ = evaluate_forwardref(field.type_, globalns, localns or None) + if field.outer_type_.__class__ == ForwardRef: + prepare = True + field.outer_type_ = evaluate_forwardref(field.outer_type_, globalns, localns or None) + if prepare: + field.prepare() + + if field.sub_fields: + for sub_f in field.sub_fields: + update_field_forward_refs(sub_f, globalns=globalns, localns=localns) + + if field.discriminator_key is not None: + field.prepare_discriminated_union_sub_fields() + + +def update_model_forward_refs( + model: Type[Any], + fields: Iterable['ModelField'], + json_encoders: Dict[Union[Type[Any], str, ForwardRef], AnyCallable], + localns: 'DictStrAny', + exc_to_suppress: Tuple[Type[BaseException], ...] = (), +) -> None: + """ + Try to update model fields ForwardRefs based on model and localns. + """ + if model.__module__ in sys.modules: + globalns = sys.modules[model.__module__].__dict__.copy() + else: + globalns = {} + + globalns.setdefault(model.__name__, model) + + for f in fields: + try: + update_field_forward_refs(f, globalns=globalns, localns=localns) + except exc_to_suppress: + pass + + for key in set(json_encoders.keys()): + if isinstance(key, str): + fr: ForwardRef = ForwardRef(key) + elif isinstance(key, ForwardRef): + fr = key + else: + continue + + try: + new_key = evaluate_forwardref(fr, globalns, localns or None) + except exc_to_suppress: # pragma: no cover + continue + + json_encoders[new_key] = json_encoders.pop(key) + + +def get_class(type_: Type[Any]) -> Union[None, bool, Type[Any]]: + """ + Tries to get the class of a Type[T] annotation. Returns True if Type is used + without brackets. Otherwise returns None. + """ + if type_ is type: + return True + + if get_origin(type_) is None: + return None + + args = get_args(type_) + if not args or not isinstance(args[0], type): + return True + else: + return args[0] + + +def get_sub_types(tp: Any) -> List[Any]: + """ + Return all the types that are allowed by type `tp` + `tp` can be a `Union` of allowed types or an `Annotated` type + """ + origin = get_origin(tp) + if origin is Annotated: + return get_sub_types(get_args(tp)[0]) + elif is_union(origin): + return [x for t in get_args(tp) for x in get_sub_types(t)] + else: + return [tp] diff --git a/lib/python3.12/site-packages/pydantic/v1/utils.py b/lib/python3.12/site-packages/pydantic/v1/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..02543fd1370844a58f51c9be0341310a515a8ef6 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/utils.py @@ -0,0 +1,806 @@ +import keyword +import warnings +import weakref +from collections import OrderedDict, defaultdict, deque +from copy import deepcopy +from itertools import islice, zip_longest +from types import BuiltinFunctionType, CodeType, FunctionType, GeneratorType, LambdaType, ModuleType +from typing import ( + TYPE_CHECKING, + AbstractSet, + Any, + Callable, + Collection, + Dict, + Generator, + Iterable, + Iterator, + List, + Mapping, + NoReturn, + Optional, + Set, + Tuple, + Type, + TypeVar, + Union, +) + +from typing_extensions import Annotated + +from pydantic.v1.errors import ConfigError +from pydantic.v1.typing import ( + NoneType, + WithArgsTypes, + all_literal_values, + display_as_type, + get_args, + get_origin, + is_literal_type, + is_union, +) +from pydantic.v1.version import version_info + +if TYPE_CHECKING: + from inspect import Signature + from pathlib import Path + + from pydantic.v1.config import BaseConfig + from pydantic.v1.dataclasses import Dataclass + from pydantic.v1.fields import ModelField + from pydantic.v1.main import BaseModel + from pydantic.v1.typing import AbstractSetIntStr, DictIntStrAny, IntStr, MappingIntStrAny, ReprArgs + + RichReprResult = Iterable[Union[Any, Tuple[Any], Tuple[str, Any], Tuple[str, Any, Any]]] + +__all__ = ( + 'import_string', + 'sequence_like', + 'validate_field_name', + 'lenient_isinstance', + 'lenient_issubclass', + 'in_ipython', + 'is_valid_identifier', + 'deep_update', + 'update_not_none', + 'almost_equal_floats', + 'get_model', + 'to_camel', + 'to_lower_camel', + 'is_valid_field', + 'smart_deepcopy', + 'PyObjectStr', + 'Representation', + 'GetterDict', + 'ValueItems', + 'version_info', # required here to match behaviour in v1.3 + 'ClassAttribute', + 'path_type', + 'ROOT_KEY', + 'get_unique_discriminator_alias', + 'get_discriminator_alias_and_values', + 'DUNDER_ATTRIBUTES', +) + +ROOT_KEY = '__root__' +# these are types that are returned unchanged by deepcopy +IMMUTABLE_NON_COLLECTIONS_TYPES: Set[Type[Any]] = { + int, + float, + complex, + str, + bool, + bytes, + type, + NoneType, + FunctionType, + BuiltinFunctionType, + LambdaType, + weakref.ref, + CodeType, + # note: including ModuleType will differ from behaviour of deepcopy by not producing error. + # It might be not a good idea in general, but considering that this function used only internally + # against default values of fields, this will allow to actually have a field with module as default value + ModuleType, + NotImplemented.__class__, + Ellipsis.__class__, +} + +# these are types that if empty, might be copied with simple copy() instead of deepcopy() +BUILTIN_COLLECTIONS: Set[Type[Any]] = { + list, + set, + tuple, + frozenset, + dict, + OrderedDict, + defaultdict, + deque, +} + + +def import_string(dotted_path: str) -> Any: + """ + Stolen approximately from django. Import a dotted module path and return the attribute/class designated by the + last name in the path. Raise ImportError if the import fails. + """ + from importlib import import_module + + try: + module_path, class_name = dotted_path.strip(' ').rsplit('.', 1) + except ValueError as e: + raise ImportError(f'"{dotted_path}" doesn\'t look like a module path') from e + + module = import_module(module_path) + try: + return getattr(module, class_name) + except AttributeError as e: + raise ImportError(f'Module "{module_path}" does not define a "{class_name}" attribute') from e + + +def truncate(v: Union[str], *, max_len: int = 80) -> str: + """ + Truncate a value and add a unicode ellipsis (three dots) to the end if it was too long + """ + warnings.warn('`truncate` is no-longer used by pydantic and is deprecated', DeprecationWarning) + if isinstance(v, str) and len(v) > (max_len - 2): + # -3 so quote + string + … + quote has correct length + return (v[: (max_len - 3)] + '…').__repr__() + try: + v = v.__repr__() + except TypeError: + v = v.__class__.__repr__(v) # in case v is a type + if len(v) > max_len: + v = v[: max_len - 1] + '…' + return v + + +def sequence_like(v: Any) -> bool: + return isinstance(v, (list, tuple, set, frozenset, GeneratorType, deque)) + + +def validate_field_name(bases: Iterable[Type[Any]], field_name: str) -> None: + """ + Ensure that the field's name does not shadow an existing attribute of the model. + """ + for base in bases: + if getattr(base, field_name, None): + raise NameError( + f'Field name "{field_name}" shadows a BaseModel attribute; ' + f'use a different field name with "alias=\'{field_name}\'".' + ) + + +def lenient_isinstance(o: Any, class_or_tuple: Union[Type[Any], Tuple[Type[Any], ...], None]) -> bool: + try: + return isinstance(o, class_or_tuple) # type: ignore[arg-type] + except TypeError: + return False + + +def lenient_issubclass(cls: Any, class_or_tuple: Union[Type[Any], Tuple[Type[Any], ...], None]) -> bool: + try: + return isinstance(cls, type) and issubclass(cls, class_or_tuple) # type: ignore[arg-type] + except TypeError: + if isinstance(cls, WithArgsTypes): + return False + raise # pragma: no cover + + +def in_ipython() -> bool: + """ + Check whether we're in an ipython environment, including jupyter notebooks. + """ + try: + eval('__IPYTHON__') + except NameError: + return False + else: # pragma: no cover + return True + + +def is_valid_identifier(identifier: str) -> bool: + """ + Checks that a string is a valid identifier and not a Python keyword. + :param identifier: The identifier to test. + :return: True if the identifier is valid. + """ + return identifier.isidentifier() and not keyword.iskeyword(identifier) + + +KeyType = TypeVar('KeyType') + + +def deep_update(mapping: Dict[KeyType, Any], *updating_mappings: Dict[KeyType, Any]) -> Dict[KeyType, Any]: + updated_mapping = mapping.copy() + for updating_mapping in updating_mappings: + for k, v in updating_mapping.items(): + if k in updated_mapping and isinstance(updated_mapping[k], dict) and isinstance(v, dict): + updated_mapping[k] = deep_update(updated_mapping[k], v) + else: + updated_mapping[k] = v + return updated_mapping + + +def update_not_none(mapping: Dict[Any, Any], **update: Any) -> None: + mapping.update({k: v for k, v in update.items() if v is not None}) + + +def almost_equal_floats(value_1: float, value_2: float, *, delta: float = 1e-8) -> bool: + """ + Return True if two floats are almost equal + """ + return abs(value_1 - value_2) <= delta + + +def generate_model_signature( + init: Callable[..., None], fields: Dict[str, 'ModelField'], config: Type['BaseConfig'] +) -> 'Signature': + """ + Generate signature for model based on its fields + """ + from inspect import Parameter, Signature, signature + + from pydantic.v1.config import Extra + + present_params = signature(init).parameters.values() + merged_params: Dict[str, Parameter] = {} + var_kw = None + use_var_kw = False + + for param in islice(present_params, 1, None): # skip self arg + if param.kind is param.VAR_KEYWORD: + var_kw = param + continue + merged_params[param.name] = param + + if var_kw: # if custom init has no var_kw, fields which are not declared in it cannot be passed through + allow_names = config.allow_population_by_field_name + for field_name, field in fields.items(): + param_name = field.alias + if field_name in merged_params or param_name in merged_params: + continue + elif not is_valid_identifier(param_name): + if allow_names and is_valid_identifier(field_name): + param_name = field_name + else: + use_var_kw = True + continue + + # TODO: replace annotation with actual expected types once #1055 solved + kwargs = {'default': field.default} if not field.required else {} + merged_params[param_name] = Parameter( + param_name, Parameter.KEYWORD_ONLY, annotation=field.annotation, **kwargs + ) + + if config.extra is Extra.allow: + use_var_kw = True + + if var_kw and use_var_kw: + # Make sure the parameter for extra kwargs + # does not have the same name as a field + default_model_signature = [ + ('__pydantic_self__', Parameter.POSITIONAL_OR_KEYWORD), + ('data', Parameter.VAR_KEYWORD), + ] + if [(p.name, p.kind) for p in present_params] == default_model_signature: + # if this is the standard model signature, use extra_data as the extra args name + var_kw_name = 'extra_data' + else: + # else start from var_kw + var_kw_name = var_kw.name + + # generate a name that's definitely unique + while var_kw_name in fields: + var_kw_name += '_' + merged_params[var_kw_name] = var_kw.replace(name=var_kw_name) + + return Signature(parameters=list(merged_params.values()), return_annotation=None) + + +def get_model(obj: Union[Type['BaseModel'], Type['Dataclass']]) -> Type['BaseModel']: + from pydantic.v1.main import BaseModel + + try: + model_cls = obj.__pydantic_model__ # type: ignore + except AttributeError: + model_cls = obj + + if not issubclass(model_cls, BaseModel): + raise TypeError('Unsupported type, must be either BaseModel or dataclass') + return model_cls + + +def to_camel(string: str) -> str: + return ''.join(word.capitalize() for word in string.split('_')) + + +def to_lower_camel(string: str) -> str: + if len(string) >= 1: + pascal_string = to_camel(string) + return pascal_string[0].lower() + pascal_string[1:] + return string.lower() + + +T = TypeVar('T') + + +def unique_list( + input_list: Union[List[T], Tuple[T, ...]], + *, + name_factory: Callable[[T], str] = str, +) -> List[T]: + """ + Make a list unique while maintaining order. + We update the list if another one with the same name is set + (e.g. root validator overridden in subclass) + """ + result: List[T] = [] + result_names: List[str] = [] + for v in input_list: + v_name = name_factory(v) + if v_name not in result_names: + result_names.append(v_name) + result.append(v) + else: + result[result_names.index(v_name)] = v + + return result + + +class PyObjectStr(str): + """ + String class where repr doesn't include quotes. Useful with Representation when you want to return a string + representation of something that valid (or pseudo-valid) python. + """ + + def __repr__(self) -> str: + return str(self) + + +class Representation: + """ + Mixin to provide __str__, __repr__, and __pretty__ methods. See #884 for more details. + + __pretty__ is used by [devtools](https://python-devtools.helpmanual.io/) to provide human readable representations + of objects. + """ + + __slots__: Tuple[str, ...] = tuple() + + def __repr_args__(self) -> 'ReprArgs': + """ + Returns the attributes to show in __str__, __repr__, and __pretty__ this is generally overridden. + + Can either return: + * name - value pairs, e.g.: `[('foo_name', 'foo'), ('bar_name', ['b', 'a', 'r'])]` + * or, just values, e.g.: `[(None, 'foo'), (None, ['b', 'a', 'r'])]` + """ + attrs = ((s, getattr(self, s)) for s in self.__slots__) + return [(a, v) for a, v in attrs if v is not None] + + def __repr_name__(self) -> str: + """ + Name of the instance's class, used in __repr__. + """ + return self.__class__.__name__ + + def __repr_str__(self, join_str: str) -> str: + return join_str.join(repr(v) if a is None else f'{a}={v!r}' for a, v in self.__repr_args__()) + + def __pretty__(self, fmt: Callable[[Any], Any], **kwargs: Any) -> Generator[Any, None, None]: + """ + Used by devtools (https://python-devtools.helpmanual.io/) to provide a human readable representations of objects + """ + yield self.__repr_name__() + '(' + yield 1 + for name, value in self.__repr_args__(): + if name is not None: + yield name + '=' + yield fmt(value) + yield ',' + yield 0 + yield -1 + yield ')' + + def __str__(self) -> str: + return self.__repr_str__(' ') + + def __repr__(self) -> str: + return f'{self.__repr_name__()}({self.__repr_str__(", ")})' + + def __rich_repr__(self) -> 'RichReprResult': + """Get fields for Rich library""" + for name, field_repr in self.__repr_args__(): + if name is None: + yield field_repr + else: + yield name, field_repr + + +class GetterDict(Representation): + """ + Hack to make object's smell just enough like dicts for validate_model. + + We can't inherit from Mapping[str, Any] because it upsets cython so we have to implement all methods ourselves. + """ + + __slots__ = ('_obj',) + + def __init__(self, obj: Any): + self._obj = obj + + def __getitem__(self, key: str) -> Any: + try: + return getattr(self._obj, key) + except AttributeError as e: + raise KeyError(key) from e + + def get(self, key: Any, default: Any = None) -> Any: + return getattr(self._obj, key, default) + + def extra_keys(self) -> Set[Any]: + """ + We don't want to get any other attributes of obj if the model didn't explicitly ask for them + """ + return set() + + def keys(self) -> List[Any]: + """ + Keys of the pseudo dictionary, uses a list not set so order information can be maintained like python + dictionaries. + """ + return list(self) + + def values(self) -> List[Any]: + return [self[k] for k in self] + + def items(self) -> Iterator[Tuple[str, Any]]: + for k in self: + yield k, self.get(k) + + def __iter__(self) -> Iterator[str]: + for name in dir(self._obj): + if not name.startswith('_'): + yield name + + def __len__(self) -> int: + return sum(1 for _ in self) + + def __contains__(self, item: Any) -> bool: + return item in self.keys() + + def __eq__(self, other: Any) -> bool: + return dict(self) == dict(other.items()) + + def __repr_args__(self) -> 'ReprArgs': + return [(None, dict(self))] + + def __repr_name__(self) -> str: + return f'GetterDict[{display_as_type(self._obj)}]' + + +class ValueItems(Representation): + """ + Class for more convenient calculation of excluded or included fields on values. + """ + + __slots__ = ('_items', '_type') + + def __init__(self, value: Any, items: Union['AbstractSetIntStr', 'MappingIntStrAny']) -> None: + items = self._coerce_items(items) + + if isinstance(value, (list, tuple)): + items = self._normalize_indexes(items, len(value)) + + self._items: 'MappingIntStrAny' = items + + def is_excluded(self, item: Any) -> bool: + """ + Check if item is fully excluded. + + :param item: key or index of a value + """ + return self.is_true(self._items.get(item)) + + def is_included(self, item: Any) -> bool: + """ + Check if value is contained in self._items + + :param item: key or index of value + """ + return item in self._items + + def for_element(self, e: 'IntStr') -> Optional[Union['AbstractSetIntStr', 'MappingIntStrAny']]: + """ + :param e: key or index of element on value + :return: raw values for element if self._items is dict and contain needed element + """ + + item = self._items.get(e) + return item if not self.is_true(item) else None + + def _normalize_indexes(self, items: 'MappingIntStrAny', v_length: int) -> 'DictIntStrAny': + """ + :param items: dict or set of indexes which will be normalized + :param v_length: length of sequence indexes of which will be + + >>> self._normalize_indexes({0: True, -2: True, -1: True}, 4) + {0: True, 2: True, 3: True} + >>> self._normalize_indexes({'__all__': True}, 4) + {0: True, 1: True, 2: True, 3: True} + """ + + normalized_items: 'DictIntStrAny' = {} + all_items = None + for i, v in items.items(): + if not (isinstance(v, Mapping) or isinstance(v, AbstractSet) or self.is_true(v)): + raise TypeError(f'Unexpected type of exclude value for index "{i}" {v.__class__}') + if i == '__all__': + all_items = self._coerce_value(v) + continue + if not isinstance(i, int): + raise TypeError( + 'Excluding fields from a sequence of sub-models or dicts must be performed index-wise: ' + 'expected integer keys or keyword "__all__"' + ) + normalized_i = v_length + i if i < 0 else i + normalized_items[normalized_i] = self.merge(v, normalized_items.get(normalized_i)) + + if not all_items: + return normalized_items + if self.is_true(all_items): + for i in range(v_length): + normalized_items.setdefault(i, ...) + return normalized_items + for i in range(v_length): + normalized_item = normalized_items.setdefault(i, {}) + if not self.is_true(normalized_item): + normalized_items[i] = self.merge(all_items, normalized_item) + return normalized_items + + @classmethod + def merge(cls, base: Any, override: Any, intersect: bool = False) -> Any: + """ + Merge a ``base`` item with an ``override`` item. + + Both ``base`` and ``override`` are converted to dictionaries if possible. + Sets are converted to dictionaries with the sets entries as keys and + Ellipsis as values. + + Each key-value pair existing in ``base`` is merged with ``override``, + while the rest of the key-value pairs are updated recursively with this function. + + Merging takes place based on the "union" of keys if ``intersect`` is + set to ``False`` (default) and on the intersection of keys if + ``intersect`` is set to ``True``. + """ + override = cls._coerce_value(override) + base = cls._coerce_value(base) + if override is None: + return base + if cls.is_true(base) or base is None: + return override + if cls.is_true(override): + return base if intersect else override + + # intersection or union of keys while preserving ordering: + if intersect: + merge_keys = [k for k in base if k in override] + [k for k in override if k in base] + else: + merge_keys = list(base) + [k for k in override if k not in base] + + merged: 'DictIntStrAny' = {} + for k in merge_keys: + merged_item = cls.merge(base.get(k), override.get(k), intersect=intersect) + if merged_item is not None: + merged[k] = merged_item + + return merged + + @staticmethod + def _coerce_items(items: Union['AbstractSetIntStr', 'MappingIntStrAny']) -> 'MappingIntStrAny': + if isinstance(items, Mapping): + pass + elif isinstance(items, AbstractSet): + items = dict.fromkeys(items, ...) + else: + class_name = getattr(items, '__class__', '???') + assert_never( + items, + f'Unexpected type of exclude value {class_name}', + ) + return items + + @classmethod + def _coerce_value(cls, value: Any) -> Any: + if value is None or cls.is_true(value): + return value + return cls._coerce_items(value) + + @staticmethod + def is_true(v: Any) -> bool: + return v is True or v is ... + + def __repr_args__(self) -> 'ReprArgs': + return [(None, self._items)] + + +class ClassAttribute: + """ + Hide class attribute from its instances + """ + + __slots__ = ( + 'name', + 'value', + ) + + def __init__(self, name: str, value: Any) -> None: + self.name = name + self.value = value + + def __get__(self, instance: Any, owner: Type[Any]) -> None: + if instance is None: + return self.value + raise AttributeError(f'{self.name!r} attribute of {owner.__name__!r} is class-only') + + +path_types = { + 'is_dir': 'directory', + 'is_file': 'file', + 'is_mount': 'mount point', + 'is_symlink': 'symlink', + 'is_block_device': 'block device', + 'is_char_device': 'char device', + 'is_fifo': 'FIFO', + 'is_socket': 'socket', +} + + +def path_type(p: 'Path') -> str: + """ + Find out what sort of thing a path is. + """ + assert p.exists(), 'path does not exist' + for method, name in path_types.items(): + if getattr(p, method)(): + return name + + return 'unknown' + + +Obj = TypeVar('Obj') + + +def smart_deepcopy(obj: Obj) -> Obj: + """ + Return type as is for immutable built-in types + Use obj.copy() for built-in empty collections + Use copy.deepcopy() for non-empty collections and unknown objects + """ + + obj_type = obj.__class__ + if obj_type in IMMUTABLE_NON_COLLECTIONS_TYPES: + return obj # fastest case: obj is immutable and not collection therefore will not be copied anyway + try: + if not obj and obj_type in BUILTIN_COLLECTIONS: + # faster way for empty collections, no need to copy its members + return obj if obj_type is tuple else obj.copy() # type: ignore # tuple doesn't have copy method + except (TypeError, ValueError, RuntimeError): + # do we really dare to catch ALL errors? Seems a bit risky + pass + + return deepcopy(obj) # slowest way when we actually might need a deepcopy + + +def is_valid_field(name: str) -> bool: + if not name.startswith('_'): + return True + return ROOT_KEY == name + + +DUNDER_ATTRIBUTES = { + '__annotations__', + '__classcell__', + '__doc__', + '__module__', + '__orig_bases__', + '__orig_class__', + '__qualname__', + '__firstlineno__', + '__static_attributes__', +} + + +def is_valid_private_name(name: str) -> bool: + return not is_valid_field(name) and name not in DUNDER_ATTRIBUTES + + +_EMPTY = object() + + +def all_identical(left: Iterable[Any], right: Iterable[Any]) -> bool: + """ + Check that the items of `left` are the same objects as those in `right`. + + >>> a, b = object(), object() + >>> all_identical([a, b, a], [a, b, a]) + True + >>> all_identical([a, b, [a]], [a, b, [a]]) # new list object, while "equal" is not "identical" + False + """ + for left_item, right_item in zip_longest(left, right, fillvalue=_EMPTY): + if left_item is not right_item: + return False + return True + + +def assert_never(obj: NoReturn, msg: str) -> NoReturn: + """ + Helper to make sure that we have covered all possible types. + + This is mostly useful for ``mypy``, docs: + https://mypy.readthedocs.io/en/latest/literal_types.html#exhaustive-checks + """ + raise TypeError(msg) + + +def get_unique_discriminator_alias(all_aliases: Collection[str], discriminator_key: str) -> str: + """Validate that all aliases are the same and if that's the case return the alias""" + unique_aliases = set(all_aliases) + if len(unique_aliases) > 1: + raise ConfigError( + f'Aliases for discriminator {discriminator_key!r} must be the same (got {", ".join(sorted(all_aliases))})' + ) + return unique_aliases.pop() + + +def get_discriminator_alias_and_values(tp: Any, discriminator_key: str) -> Tuple[str, Tuple[str, ...]]: + """ + Get alias and all valid values in the `Literal` type of the discriminator field + `tp` can be a `BaseModel` class or directly an `Annotated` `Union` of many. + """ + is_root_model = getattr(tp, '__custom_root_type__', False) + + if get_origin(tp) is Annotated: + tp = get_args(tp)[0] + + if hasattr(tp, '__pydantic_model__'): + tp = tp.__pydantic_model__ + + if is_union(get_origin(tp)): + alias, all_values = _get_union_alias_and_all_values(tp, discriminator_key) + return alias, tuple(v for values in all_values for v in values) + elif is_root_model: + union_type = tp.__fields__[ROOT_KEY].type_ + alias, all_values = _get_union_alias_and_all_values(union_type, discriminator_key) + + if len(set(all_values)) > 1: + raise ConfigError( + f'Field {discriminator_key!r} is not the same for all submodels of {display_as_type(tp)!r}' + ) + + return alias, all_values[0] + + else: + try: + t_discriminator_type = tp.__fields__[discriminator_key].type_ + except AttributeError as e: + raise TypeError(f'Type {tp.__name__!r} is not a valid `BaseModel` or `dataclass`') from e + except KeyError as e: + raise ConfigError(f'Model {tp.__name__!r} needs a discriminator field for key {discriminator_key!r}') from e + + if not is_literal_type(t_discriminator_type): + raise ConfigError(f'Field {discriminator_key!r} of model {tp.__name__!r} needs to be a `Literal`') + + return tp.__fields__[discriminator_key].alias, all_literal_values(t_discriminator_type) + + +def _get_union_alias_and_all_values( + union_type: Type[Any], discriminator_key: str +) -> Tuple[str, Tuple[Tuple[str, ...], ...]]: + zipped_aliases_values = [get_discriminator_alias_and_values(t, discriminator_key) for t in get_args(union_type)] + # unzip: [('alias_a',('v1', 'v2)), ('alias_b', ('v3',))] => [('alias_a', 'alias_b'), (('v1', 'v2'), ('v3',))] + all_aliases, all_values = zip(*zipped_aliases_values) + return get_unique_discriminator_alias(all_aliases, discriminator_key), all_values diff --git a/lib/python3.12/site-packages/pydantic/v1/validators.py b/lib/python3.12/site-packages/pydantic/v1/validators.py new file mode 100644 index 0000000000000000000000000000000000000000..c0940e8114a47ceacf21f4df946a6401946b9b8f --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/validators.py @@ -0,0 +1,768 @@ +import math +import re +from collections import OrderedDict, deque +from collections.abc import Hashable as CollectionsHashable +from datetime import date, datetime, time, timedelta +from decimal import Decimal, DecimalException +from enum import Enum, IntEnum +from ipaddress import IPv4Address, IPv4Interface, IPv4Network, IPv6Address, IPv6Interface, IPv6Network +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Deque, + Dict, + ForwardRef, + FrozenSet, + Generator, + Hashable, + List, + NamedTuple, + Pattern, + Set, + Tuple, + Type, + TypeVar, + Union, +) +from uuid import UUID +from warnings import warn + +from pydantic.v1 import errors +from pydantic.v1.datetime_parse import parse_date, parse_datetime, parse_duration, parse_time +from pydantic.v1.typing import ( + AnyCallable, + all_literal_values, + display_as_type, + get_class, + is_callable_type, + is_literal_type, + is_namedtuple, + is_none_type, + is_typeddict, +) +from pydantic.v1.utils import almost_equal_floats, lenient_issubclass, sequence_like + +if TYPE_CHECKING: + from typing_extensions import Literal, TypedDict + + from pydantic.v1.config import BaseConfig + from pydantic.v1.fields import ModelField + from pydantic.v1.types import ConstrainedDecimal, ConstrainedFloat, ConstrainedInt + + ConstrainedNumber = Union[ConstrainedDecimal, ConstrainedFloat, ConstrainedInt] + AnyOrderedDict = OrderedDict[Any, Any] + Number = Union[int, float, Decimal] + StrBytes = Union[str, bytes] + + +def str_validator(v: Any) -> Union[str]: + if isinstance(v, str): + if isinstance(v, Enum): + return v.value + else: + return v + elif isinstance(v, (float, int, Decimal)): + # is there anything else we want to add here? If you think so, create an issue. + return str(v) + elif isinstance(v, (bytes, bytearray)): + return v.decode() + else: + raise errors.StrError() + + +def strict_str_validator(v: Any) -> Union[str]: + if isinstance(v, str) and not isinstance(v, Enum): + return v + raise errors.StrError() + + +def bytes_validator(v: Any) -> Union[bytes]: + if isinstance(v, bytes): + return v + elif isinstance(v, bytearray): + return bytes(v) + elif isinstance(v, str): + return v.encode() + elif isinstance(v, (float, int, Decimal)): + return str(v).encode() + else: + raise errors.BytesError() + + +def strict_bytes_validator(v: Any) -> Union[bytes]: + if isinstance(v, bytes): + return v + elif isinstance(v, bytearray): + return bytes(v) + else: + raise errors.BytesError() + + +BOOL_FALSE = {0, '0', 'off', 'f', 'false', 'n', 'no'} +BOOL_TRUE = {1, '1', 'on', 't', 'true', 'y', 'yes'} + + +def bool_validator(v: Any) -> bool: + if v is True or v is False: + return v + if isinstance(v, bytes): + v = v.decode() + if isinstance(v, str): + v = v.lower() + try: + if v in BOOL_TRUE: + return True + if v in BOOL_FALSE: + return False + except TypeError: + raise errors.BoolError() + raise errors.BoolError() + + +# matches the default limit cpython, see https://github.com/python/cpython/pull/96500 +max_str_int = 4_300 + + +def int_validator(v: Any) -> int: + if isinstance(v, int) and not (v is True or v is False): + return v + + # see https://github.com/pydantic/pydantic/issues/1477 and in turn, https://github.com/python/cpython/issues/95778 + # this check should be unnecessary once patch releases are out for 3.7, 3.8, 3.9 and 3.10 + # but better to check here until then. + # NOTICE: this does not fully protect user from the DOS risk since the standard library JSON implementation + # (and other std lib modules like xml) use `int()` and are likely called before this, the best workaround is to + # 1. update to the latest patch release of python once released, 2. use a different JSON library like ujson + if isinstance(v, (str, bytes, bytearray)) and len(v) > max_str_int: + raise errors.IntegerError() + + try: + return int(v) + except (TypeError, ValueError, OverflowError): + raise errors.IntegerError() + + +def strict_int_validator(v: Any) -> int: + if isinstance(v, int) and not (v is True or v is False): + return v + raise errors.IntegerError() + + +def float_validator(v: Any) -> float: + if isinstance(v, float): + return v + + try: + return float(v) + except (TypeError, ValueError): + raise errors.FloatError() + + +def strict_float_validator(v: Any) -> float: + if isinstance(v, float): + return v + raise errors.FloatError() + + +def float_finite_validator(v: 'Number', field: 'ModelField', config: 'BaseConfig') -> 'Number': + allow_inf_nan = getattr(field.type_, 'allow_inf_nan', None) + if allow_inf_nan is None: + allow_inf_nan = config.allow_inf_nan + + if allow_inf_nan is False and (math.isnan(v) or math.isinf(v)): + raise errors.NumberNotFiniteError() + return v + + +def number_multiple_validator(v: 'Number', field: 'ModelField') -> 'Number': + field_type: ConstrainedNumber = field.type_ + if field_type.multiple_of is not None: + mod = float(v) / float(field_type.multiple_of) % 1 + if not almost_equal_floats(mod, 0.0) and not almost_equal_floats(mod, 1.0): + raise errors.NumberNotMultipleError(multiple_of=field_type.multiple_of) + return v + + +def number_size_validator(v: 'Number', field: 'ModelField') -> 'Number': + field_type: ConstrainedNumber = field.type_ + if field_type.gt is not None and not v > field_type.gt: + raise errors.NumberNotGtError(limit_value=field_type.gt) + elif field_type.ge is not None and not v >= field_type.ge: + raise errors.NumberNotGeError(limit_value=field_type.ge) + + if field_type.lt is not None and not v < field_type.lt: + raise errors.NumberNotLtError(limit_value=field_type.lt) + if field_type.le is not None and not v <= field_type.le: + raise errors.NumberNotLeError(limit_value=field_type.le) + + return v + + +def constant_validator(v: 'Any', field: 'ModelField') -> 'Any': + """Validate ``const`` fields. + + The value provided for a ``const`` field must be equal to the default value + of the field. This is to support the keyword of the same name in JSON + Schema. + """ + if v != field.default: + raise errors.WrongConstantError(given=v, permitted=[field.default]) + + return v + + +def anystr_length_validator(v: 'StrBytes', config: 'BaseConfig') -> 'StrBytes': + v_len = len(v) + + min_length = config.min_anystr_length + if v_len < min_length: + raise errors.AnyStrMinLengthError(limit_value=min_length) + + max_length = config.max_anystr_length + if max_length is not None and v_len > max_length: + raise errors.AnyStrMaxLengthError(limit_value=max_length) + + return v + + +def anystr_strip_whitespace(v: 'StrBytes') -> 'StrBytes': + return v.strip() + + +def anystr_upper(v: 'StrBytes') -> 'StrBytes': + return v.upper() + + +def anystr_lower(v: 'StrBytes') -> 'StrBytes': + return v.lower() + + +def ordered_dict_validator(v: Any) -> 'AnyOrderedDict': + if isinstance(v, OrderedDict): + return v + + try: + return OrderedDict(v) + except (TypeError, ValueError): + raise errors.DictError() + + +def dict_validator(v: Any) -> Dict[Any, Any]: + if isinstance(v, dict): + return v + + try: + return dict(v) + except (TypeError, ValueError): + raise errors.DictError() + + +def list_validator(v: Any) -> List[Any]: + if isinstance(v, list): + return v + elif sequence_like(v): + return list(v) + else: + raise errors.ListError() + + +def tuple_validator(v: Any) -> Tuple[Any, ...]: + if isinstance(v, tuple): + return v + elif sequence_like(v): + return tuple(v) + else: + raise errors.TupleError() + + +def set_validator(v: Any) -> Set[Any]: + if isinstance(v, set): + return v + elif sequence_like(v): + return set(v) + else: + raise errors.SetError() + + +def frozenset_validator(v: Any) -> FrozenSet[Any]: + if isinstance(v, frozenset): + return v + elif sequence_like(v): + return frozenset(v) + else: + raise errors.FrozenSetError() + + +def deque_validator(v: Any) -> Deque[Any]: + if isinstance(v, deque): + return v + elif sequence_like(v): + return deque(v) + else: + raise errors.DequeError() + + +def enum_member_validator(v: Any, field: 'ModelField', config: 'BaseConfig') -> Enum: + try: + enum_v = field.type_(v) + except ValueError: + # field.type_ should be an enum, so will be iterable + raise errors.EnumMemberError(enum_values=list(field.type_)) + return enum_v.value if config.use_enum_values else enum_v + + +def uuid_validator(v: Any, field: 'ModelField') -> UUID: + try: + if isinstance(v, str): + v = UUID(v) + elif isinstance(v, (bytes, bytearray)): + try: + v = UUID(v.decode()) + except ValueError: + # 16 bytes in big-endian order as the bytes argument fail + # the above check + v = UUID(bytes=v) + except ValueError: + raise errors.UUIDError() + + if not isinstance(v, UUID): + raise errors.UUIDError() + + required_version = getattr(field.type_, '_required_version', None) + if required_version and v.version != required_version: + raise errors.UUIDVersionError(required_version=required_version) + + return v + + +def decimal_validator(v: Any) -> Decimal: + if isinstance(v, Decimal): + return v + elif isinstance(v, (bytes, bytearray)): + v = v.decode() + + v = str(v).strip() + + try: + v = Decimal(v) + except DecimalException: + raise errors.DecimalError() + + if not v.is_finite(): + raise errors.DecimalIsNotFiniteError() + + return v + + +def hashable_validator(v: Any) -> Hashable: + if isinstance(v, Hashable): + return v + + raise errors.HashableError() + + +def ip_v4_address_validator(v: Any) -> IPv4Address: + if isinstance(v, IPv4Address): + return v + + try: + return IPv4Address(v) + except ValueError: + raise errors.IPv4AddressError() + + +def ip_v6_address_validator(v: Any) -> IPv6Address: + if isinstance(v, IPv6Address): + return v + + try: + return IPv6Address(v) + except ValueError: + raise errors.IPv6AddressError() + + +def ip_v4_network_validator(v: Any) -> IPv4Network: + """ + Assume IPv4Network initialised with a default ``strict`` argument + + See more: + https://docs.python.org/library/ipaddress.html#ipaddress.IPv4Network + """ + if isinstance(v, IPv4Network): + return v + + try: + return IPv4Network(v) + except ValueError: + raise errors.IPv4NetworkError() + + +def ip_v6_network_validator(v: Any) -> IPv6Network: + """ + Assume IPv6Network initialised with a default ``strict`` argument + + See more: + https://docs.python.org/library/ipaddress.html#ipaddress.IPv6Network + """ + if isinstance(v, IPv6Network): + return v + + try: + return IPv6Network(v) + except ValueError: + raise errors.IPv6NetworkError() + + +def ip_v4_interface_validator(v: Any) -> IPv4Interface: + if isinstance(v, IPv4Interface): + return v + + try: + return IPv4Interface(v) + except ValueError: + raise errors.IPv4InterfaceError() + + +def ip_v6_interface_validator(v: Any) -> IPv6Interface: + if isinstance(v, IPv6Interface): + return v + + try: + return IPv6Interface(v) + except ValueError: + raise errors.IPv6InterfaceError() + + +def path_validator(v: Any) -> Path: + if isinstance(v, Path): + return v + + try: + return Path(v) + except TypeError: + raise errors.PathError() + + +def path_exists_validator(v: Any) -> Path: + if not v.exists(): + raise errors.PathNotExistsError(path=v) + + return v + + +def callable_validator(v: Any) -> AnyCallable: + """ + Perform a simple check if the value is callable. + + Note: complete matching of argument type hints and return types is not performed + """ + if callable(v): + return v + + raise errors.CallableError(value=v) + + +def enum_validator(v: Any) -> Enum: + if isinstance(v, Enum): + return v + + raise errors.EnumError(value=v) + + +def int_enum_validator(v: Any) -> IntEnum: + if isinstance(v, IntEnum): + return v + + raise errors.IntEnumError(value=v) + + +def make_literal_validator(type_: Any) -> Callable[[Any], Any]: + permitted_choices = all_literal_values(type_) + + # To have a O(1) complexity and still return one of the values set inside the `Literal`, + # we create a dict with the set values (a set causes some problems with the way intersection works). + # In some cases the set value and checked value can indeed be different (see `test_literal_validator_str_enum`) + allowed_choices = {v: v for v in permitted_choices} + + def literal_validator(v: Any) -> Any: + try: + return allowed_choices[v] + except (KeyError, TypeError): + raise errors.WrongConstantError(given=v, permitted=permitted_choices) + + return literal_validator + + +def constr_length_validator(v: 'StrBytes', field: 'ModelField', config: 'BaseConfig') -> 'StrBytes': + v_len = len(v) + + min_length = field.type_.min_length if field.type_.min_length is not None else config.min_anystr_length + if v_len < min_length: + raise errors.AnyStrMinLengthError(limit_value=min_length) + + max_length = field.type_.max_length if field.type_.max_length is not None else config.max_anystr_length + if max_length is not None and v_len > max_length: + raise errors.AnyStrMaxLengthError(limit_value=max_length) + + return v + + +def constr_strip_whitespace(v: 'StrBytes', field: 'ModelField', config: 'BaseConfig') -> 'StrBytes': + strip_whitespace = field.type_.strip_whitespace or config.anystr_strip_whitespace + if strip_whitespace: + v = v.strip() + + return v + + +def constr_upper(v: 'StrBytes', field: 'ModelField', config: 'BaseConfig') -> 'StrBytes': + upper = field.type_.to_upper or config.anystr_upper + if upper: + v = v.upper() + + return v + + +def constr_lower(v: 'StrBytes', field: 'ModelField', config: 'BaseConfig') -> 'StrBytes': + lower = field.type_.to_lower or config.anystr_lower + if lower: + v = v.lower() + return v + + +def validate_json(v: Any, config: 'BaseConfig') -> Any: + if v is None: + # pass None through to other validators + return v + try: + return config.json_loads(v) # type: ignore + except ValueError: + raise errors.JsonError() + except TypeError: + raise errors.JsonTypeError() + + +T = TypeVar('T') + + +def make_arbitrary_type_validator(type_: Type[T]) -> Callable[[T], T]: + def arbitrary_type_validator(v: Any) -> T: + if isinstance(v, type_): + return v + raise errors.ArbitraryTypeError(expected_arbitrary_type=type_) + + return arbitrary_type_validator + + +def make_class_validator(type_: Type[T]) -> Callable[[Any], Type[T]]: + def class_validator(v: Any) -> Type[T]: + if lenient_issubclass(v, type_): + return v + raise errors.SubclassError(expected_class=type_) + + return class_validator + + +def any_class_validator(v: Any) -> Type[T]: + if isinstance(v, type): + return v + raise errors.ClassError() + + +def none_validator(v: Any) -> 'Literal[None]': + if v is None: + return v + raise errors.NotNoneError() + + +def pattern_validator(v: Any) -> Pattern[str]: + if isinstance(v, Pattern): + return v + + str_value = str_validator(v) + + try: + return re.compile(str_value) + except re.error: + raise errors.PatternError() + + +NamedTupleT = TypeVar('NamedTupleT', bound=NamedTuple) + + +def make_namedtuple_validator( + namedtuple_cls: Type[NamedTupleT], config: Type['BaseConfig'] +) -> Callable[[Tuple[Any, ...]], NamedTupleT]: + from pydantic.v1.annotated_types import create_model_from_namedtuple + + NamedTupleModel = create_model_from_namedtuple( + namedtuple_cls, + __config__=config, + __module__=namedtuple_cls.__module__, + ) + namedtuple_cls.__pydantic_model__ = NamedTupleModel # type: ignore[attr-defined] + + def namedtuple_validator(values: Tuple[Any, ...]) -> NamedTupleT: + annotations = NamedTupleModel.__annotations__ + + if len(values) > len(annotations): + raise errors.ListMaxLengthError(limit_value=len(annotations)) + + dict_values: Dict[str, Any] = dict(zip(annotations, values)) + validated_dict_values: Dict[str, Any] = dict(NamedTupleModel(**dict_values)) + return namedtuple_cls(**validated_dict_values) + + return namedtuple_validator + + +def make_typeddict_validator( + typeddict_cls: Type['TypedDict'], config: Type['BaseConfig'] # type: ignore[valid-type] +) -> Callable[[Any], Dict[str, Any]]: + from pydantic.v1.annotated_types import create_model_from_typeddict + + TypedDictModel = create_model_from_typeddict( + typeddict_cls, + __config__=config, + __module__=typeddict_cls.__module__, + ) + typeddict_cls.__pydantic_model__ = TypedDictModel # type: ignore[attr-defined] + + def typeddict_validator(values: 'TypedDict') -> Dict[str, Any]: # type: ignore[valid-type] + return TypedDictModel.parse_obj(values).dict(exclude_unset=True) + + return typeddict_validator + + +class IfConfig: + def __init__(self, validator: AnyCallable, *config_attr_names: str, ignored_value: Any = False) -> None: + self.validator = validator + self.config_attr_names = config_attr_names + self.ignored_value = ignored_value + + def check(self, config: Type['BaseConfig']) -> bool: + return any(getattr(config, name) not in {None, self.ignored_value} for name in self.config_attr_names) + + +# order is important here, for example: bool is a subclass of int so has to come first, datetime before date same, +# IPv4Interface before IPv4Address, etc +_VALIDATORS: List[Tuple[Type[Any], List[Any]]] = [ + (IntEnum, [int_validator, enum_member_validator]), + (Enum, [enum_member_validator]), + ( + str, + [ + str_validator, + IfConfig(anystr_strip_whitespace, 'anystr_strip_whitespace'), + IfConfig(anystr_upper, 'anystr_upper'), + IfConfig(anystr_lower, 'anystr_lower'), + IfConfig(anystr_length_validator, 'min_anystr_length', 'max_anystr_length'), + ], + ), + ( + bytes, + [ + bytes_validator, + IfConfig(anystr_strip_whitespace, 'anystr_strip_whitespace'), + IfConfig(anystr_upper, 'anystr_upper'), + IfConfig(anystr_lower, 'anystr_lower'), + IfConfig(anystr_length_validator, 'min_anystr_length', 'max_anystr_length'), + ], + ), + (bool, [bool_validator]), + (int, [int_validator]), + (float, [float_validator, IfConfig(float_finite_validator, 'allow_inf_nan', ignored_value=True)]), + (Path, [path_validator]), + (datetime, [parse_datetime]), + (date, [parse_date]), + (time, [parse_time]), + (timedelta, [parse_duration]), + (OrderedDict, [ordered_dict_validator]), + (dict, [dict_validator]), + (list, [list_validator]), + (tuple, [tuple_validator]), + (set, [set_validator]), + (frozenset, [frozenset_validator]), + (deque, [deque_validator]), + (UUID, [uuid_validator]), + (Decimal, [decimal_validator]), + (IPv4Interface, [ip_v4_interface_validator]), + (IPv6Interface, [ip_v6_interface_validator]), + (IPv4Address, [ip_v4_address_validator]), + (IPv6Address, [ip_v6_address_validator]), + (IPv4Network, [ip_v4_network_validator]), + (IPv6Network, [ip_v6_network_validator]), +] + + +def find_validators( # noqa: C901 (ignore complexity) + type_: Type[Any], config: Type['BaseConfig'] +) -> Generator[AnyCallable, None, None]: + from pydantic.v1.dataclasses import is_builtin_dataclass, make_dataclass_validator + + if type_ is Any or type_ is object: + return + type_type = type_.__class__ + if type_type == ForwardRef or type_type == TypeVar: + return + + if is_none_type(type_): + yield none_validator + return + if type_ is Pattern or type_ is re.Pattern: + yield pattern_validator + return + if type_ is Hashable or type_ is CollectionsHashable: + yield hashable_validator + return + if is_callable_type(type_): + yield callable_validator + return + if is_literal_type(type_): + yield make_literal_validator(type_) + return + if is_builtin_dataclass(type_): + yield from make_dataclass_validator(type_, config) + return + if type_ is Enum: + yield enum_validator + return + if type_ is IntEnum: + yield int_enum_validator + return + if is_namedtuple(type_): + yield tuple_validator + yield make_namedtuple_validator(type_, config) + return + if is_typeddict(type_): + yield make_typeddict_validator(type_, config) + return + + class_ = get_class(type_) + if class_ is not None: + if class_ is not Any and isinstance(class_, type): + yield make_class_validator(class_) + else: + yield any_class_validator + return + + for val_type, validators in _VALIDATORS: + try: + if issubclass(type_, val_type): + for v in validators: + if isinstance(v, IfConfig): + if v.check(config): + yield v.validator + else: + yield v + return + except TypeError: + raise RuntimeError(f'error checking inheritance of {type_!r} (type: {display_as_type(type_)})') + + if config.arbitrary_types_allowed: + yield make_arbitrary_type_validator(type_) + else: + if hasattr(type_, '__pydantic_core_schema__'): + warn(f'Mixing V1 and V2 models is not supported. `{type_.__name__}` is a V2 model.', UserWarning) + raise RuntimeError(f'no validator found for {type_}, see `arbitrary_types_allowed` in Config') diff --git a/lib/python3.12/site-packages/pydantic/v1/version.py b/lib/python3.12/site-packages/pydantic/v1/version.py new file mode 100644 index 0000000000000000000000000000000000000000..c77cde126f30c68223b1dc852928d5c2272608de --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/v1/version.py @@ -0,0 +1,38 @@ +__all__ = 'compiled', 'VERSION', 'version_info' + +VERSION = '1.10.21' + +try: + import cython # type: ignore +except ImportError: + compiled: bool = False +else: # pragma: no cover + try: + compiled = cython.compiled + except AttributeError: + compiled = False + + +def version_info() -> str: + import platform + import sys + from importlib import import_module + from pathlib import Path + + optional_deps = [] + for p in ('devtools', 'dotenv', 'email-validator', 'typing-extensions'): + try: + import_module(p.replace('-', '_')) + except ImportError: + continue + optional_deps.append(p) + + info = { + 'pydantic version': VERSION, + 'pydantic compiled': compiled, + 'install path': Path(__file__).resolve().parent, + 'python version': sys.version, + 'platform': platform.platform(), + 'optional deps. installed': optional_deps, + } + return '\n'.join('{:>30} {}'.format(k + ':', str(v).replace('\n', ' ')) for k, v in info.items()) diff --git a/lib/python3.12/site-packages/pydantic/validate_call_decorator.py b/lib/python3.12/site-packages/pydantic/validate_call_decorator.py new file mode 100644 index 0000000000000000000000000000000000000000..fe4d9c9b6c5ef3a98e4f2770d853d92d8a0124b0 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/validate_call_decorator.py @@ -0,0 +1,116 @@ +"""Decorator for validating function calls.""" + +from __future__ import annotations as _annotations + +import inspect +from functools import partial +from types import BuiltinFunctionType +from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast, overload + +from ._internal import _generate_schema, _typing_extra, _validate_call +from .errors import PydanticUserError + +__all__ = ('validate_call',) + +if TYPE_CHECKING: + from .config import ConfigDict + + AnyCallableT = TypeVar('AnyCallableT', bound=Callable[..., Any]) + + +_INVALID_TYPE_ERROR_CODE = 'validate-call-type' + + +def _check_function_type(function: object) -> None: + """Check if the input function is a supported type for `validate_call`.""" + if isinstance(function, _generate_schema.VALIDATE_CALL_SUPPORTED_TYPES): + try: + inspect.signature(cast(_generate_schema.ValidateCallSupportedTypes, function)) + except ValueError: + raise PydanticUserError( + f"Input function `{function}` doesn't have a valid signature", code=_INVALID_TYPE_ERROR_CODE + ) + + if isinstance(function, partial): + try: + assert not isinstance(partial.func, partial), 'Partial of partial' + _check_function_type(function.func) + except PydanticUserError as e: + raise PydanticUserError( + f'Partial of `{function.func}` is invalid because the type of `{function.func}` is not supported by `validate_call`', + code=_INVALID_TYPE_ERROR_CODE, + ) from e + + return + + if isinstance(function, BuiltinFunctionType): + raise PydanticUserError(f'Input built-in function `{function}` is not supported', code=_INVALID_TYPE_ERROR_CODE) + if isinstance(function, (classmethod, staticmethod, property)): + name = type(function).__name__ + raise PydanticUserError( + f'The `@{name}` decorator should be applied after `@validate_call` (put `@{name}` on top)', + code=_INVALID_TYPE_ERROR_CODE, + ) + + if inspect.isclass(function): + raise PydanticUserError( + f'Unable to validate {function}: `validate_call` should be applied to functions, not classes (put `@validate_call` on top of `__init__` or `__new__` instead)', + code=_INVALID_TYPE_ERROR_CODE, + ) + if callable(function): + raise PydanticUserError( + f'Unable to validate {function}: `validate_call` should be applied to functions, not instances or other callables. Use `validate_call` explicitly on `__call__` instead.', + code=_INVALID_TYPE_ERROR_CODE, + ) + + raise PydanticUserError( + f'Unable to validate {function}: `validate_call` should be applied to one of the following: function, method, partial, or lambda', + code=_INVALID_TYPE_ERROR_CODE, + ) + + +@overload +def validate_call( + *, config: ConfigDict | None = None, validate_return: bool = False +) -> Callable[[AnyCallableT], AnyCallableT]: ... + + +@overload +def validate_call(func: AnyCallableT, /) -> AnyCallableT: ... + + +def validate_call( + func: AnyCallableT | None = None, + /, + *, + config: ConfigDict | None = None, + validate_return: bool = False, +) -> AnyCallableT | Callable[[AnyCallableT], AnyCallableT]: + """!!! abstract "Usage Documentation" + [Validation Decorator](../concepts/validation_decorator.md) + + Returns a decorated wrapper around the function that validates the arguments and, optionally, the return value. + + Usage may be either as a plain decorator `@validate_call` or with arguments `@validate_call(...)`. + + Args: + func: The function to be decorated. + config: The configuration dictionary. + validate_return: Whether to validate the return value. + + Returns: + The decorated function. + """ + parent_namespace = _typing_extra.parent_frame_namespace() + + def validate(function: AnyCallableT) -> AnyCallableT: + _check_function_type(function) + validate_call_wrapper = _validate_call.ValidateCallWrapper( + cast(_generate_schema.ValidateCallSupportedTypes, function), config, validate_return, parent_namespace + ) + return _validate_call.update_wrapper_attributes(function, validate_call_wrapper.__call__) # type: ignore + + if func is not None: + return validate(func) + else: + return validate diff --git a/lib/python3.12/site-packages/pydantic/validators.py b/lib/python3.12/site-packages/pydantic/validators.py new file mode 100644 index 0000000000000000000000000000000000000000..7921b04f05d5ef0ad20365e7372a48d86724dc51 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/validators.py @@ -0,0 +1,5 @@ +"""The `validators` module is a backport module from V1.""" + +from ._migration import getattr_migration + +__getattr__ = getattr_migration(__name__) diff --git a/lib/python3.12/site-packages/pydantic/version.py b/lib/python3.12/site-packages/pydantic/version.py new file mode 100644 index 0000000000000000000000000000000000000000..22d3c9a7b54bb2c0af7af4b9305e0e775fbc0d37 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/version.py @@ -0,0 +1,84 @@ +"""The `version` module holds the version information for Pydantic.""" + +from __future__ import annotations as _annotations + +from pydantic_core import __version__ as __pydantic_core_version__ + +__all__ = 'VERSION', 'version_info' + +VERSION = '2.11.10' +"""The version of Pydantic.""" + + +def version_short() -> str: + """Return the `major.minor` part of Pydantic version. + + It returns '2.1' if Pydantic version is '2.1.1'. + """ + return '.'.join(VERSION.split('.')[:2]) + + +def version_info() -> str: + """Return complete version information for Pydantic and its dependencies.""" + import importlib.metadata + import platform + import sys + from pathlib import Path + + import pydantic_core._pydantic_core as pdc + + from ._internal import _git as git + + # get data about packages that are closely related to pydantic, use pydantic or often conflict with pydantic + package_names = { + 'email-validator', + 'fastapi', + 'mypy', + 'pydantic-extra-types', + 'pydantic-settings', + 'pyright', + 'typing_extensions', + } + related_packages = [] + + for dist in importlib.metadata.distributions(): + name = dist.metadata['Name'] + if name in package_names: + related_packages.append(f'{name}-{dist.version}') + + pydantic_dir = Path(__file__).parents[1].resolve() + most_recent_commit = ( + git.git_revision(pydantic_dir) if git.is_git_repo(pydantic_dir) and git.have_git() else 'unknown' + ) + + info = { + 'pydantic version': VERSION, + 'pydantic-core version': pdc.__version__, + 'pydantic-core build': getattr(pdc, 'build_info', None) or pdc.build_profile, + 'python version': sys.version, + 'platform': platform.platform(), + 'related packages': ' '.join(related_packages), + 'commit': most_recent_commit, + } + return '\n'.join('{:>30} {}'.format(k + ':', str(v).replace('\n', ' ')) for k, v in info.items()) + + +def check_pydantic_core_version() -> bool: + """Check that the installed `pydantic-core` dependency is compatible.""" + # Keep this in sync with the version constraint in the `pyproject.toml` dependencies: + return __pydantic_core_version__ == '2.33.2' + + +def parse_mypy_version(version: str) -> tuple[int, int, int]: + """Parse `mypy` string version to a 3-tuple of ints. + + It parses normal version like `1.11.0` and extra info followed by a `+` sign + like `1.11.0+dev.d6d9d8cd4f27c52edac1f537e236ec48a01e54cb.dirty`. + + Args: + version: The mypy version string. + + Returns: + A triple of ints, e.g. `(1, 11, 0)`. + """ + return tuple(map(int, version.partition('+')[0].split('.'))) # pyright: ignore[reportReturnType] diff --git a/lib/python3.12/site-packages/pydantic/warnings.py b/lib/python3.12/site-packages/pydantic/warnings.py new file mode 100644 index 0000000000000000000000000000000000000000..6e874dd12c61df0c69b12de40e5122fce2738495 --- /dev/null +++ b/lib/python3.12/site-packages/pydantic/warnings.py @@ -0,0 +1,96 @@ +"""Pydantic-specific warnings.""" + +from __future__ import annotations as _annotations + +from .version import version_short + +__all__ = ( + 'PydanticDeprecatedSince20', + 'PydanticDeprecatedSince26', + 'PydanticDeprecatedSince29', + 'PydanticDeprecatedSince210', + 'PydanticDeprecatedSince211', + 'PydanticDeprecationWarning', + 'PydanticExperimentalWarning', +) + + +class PydanticDeprecationWarning(DeprecationWarning): + """A Pydantic specific deprecation warning. + + This warning is raised when using deprecated functionality in Pydantic. It provides information on when the + deprecation was introduced and the expected version in which the corresponding functionality will be removed. + + Attributes: + message: Description of the warning. + since: Pydantic version in what the deprecation was introduced. + expected_removal: Pydantic version in what the corresponding functionality expected to be removed. + """ + + message: str + since: tuple[int, int] + expected_removal: tuple[int, int] + + def __init__( + self, message: str, *args: object, since: tuple[int, int], expected_removal: tuple[int, int] | None = None + ) -> None: + super().__init__(message, *args) + self.message = message.rstrip('.') + self.since = since + self.expected_removal = expected_removal if expected_removal is not None else (since[0] + 1, 0) + + def __str__(self) -> str: + message = ( + f'{self.message}. Deprecated in Pydantic V{self.since[0]}.{self.since[1]}' + f' to be removed in V{self.expected_removal[0]}.{self.expected_removal[1]}.' + ) + if self.since == (2, 0): + message += f' See Pydantic V2 Migration Guide at https://errors.pydantic.dev/{version_short()}/migration/' + return message + + +class PydanticDeprecatedSince20(PydanticDeprecationWarning): + """A specific `PydanticDeprecationWarning` subclass defining functionality deprecated since Pydantic 2.0.""" + + def __init__(self, message: str, *args: object) -> None: + super().__init__(message, *args, since=(2, 0), expected_removal=(3, 0)) + + +class PydanticDeprecatedSince26(PydanticDeprecationWarning): + """A specific `PydanticDeprecationWarning` subclass defining functionality deprecated since Pydantic 2.6.""" + + def __init__(self, message: str, *args: object) -> None: + super().__init__(message, *args, since=(2, 6), expected_removal=(3, 0)) + + +class PydanticDeprecatedSince29(PydanticDeprecationWarning): + """A specific `PydanticDeprecationWarning` subclass defining functionality deprecated since Pydantic 2.9.""" + + def __init__(self, message: str, *args: object) -> None: + super().__init__(message, *args, since=(2, 9), expected_removal=(3, 0)) + + +class PydanticDeprecatedSince210(PydanticDeprecationWarning): + """A specific `PydanticDeprecationWarning` subclass defining functionality deprecated since Pydantic 2.10.""" + + def __init__(self, message: str, *args: object) -> None: + super().__init__(message, *args, since=(2, 10), expected_removal=(3, 0)) + + +class PydanticDeprecatedSince211(PydanticDeprecationWarning): + """A specific `PydanticDeprecationWarning` subclass defining functionality deprecated since Pydantic 2.11.""" + + def __init__(self, message: str, *args: object) -> None: + super().__init__(message, *args, since=(2, 11), expected_removal=(3, 0)) + + +class GenericBeforeBaseModelWarning(Warning): + pass + + +class PydanticExperimentalWarning(Warning): + """A Pydantic specific experimental functionality warning. + + This warning is raised when using experimental functionality in Pydantic. + It is raised to warn users that the functionality may change or be removed in future versions of Pydantic. + """ diff --git a/lib/python3.12/site-packages/pytz/__init__.py b/lib/python3.12/site-packages/pytz/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8e44c3c404a9bc3b8472d17028f9aeb3c98ca9ef --- /dev/null +++ b/lib/python3.12/site-packages/pytz/__init__.py @@ -0,0 +1,1556 @@ +''' +datetime.tzinfo timezone definitions generated from the +Olson timezone database: + + ftp://elsie.nci.nih.gov/pub/tz*.tar.gz + +See the datetime section of the Python Library Reference for information +on how to use these modules. +''' + +import sys +import datetime +import os.path + +from pytz.exceptions import AmbiguousTimeError +from pytz.exceptions import InvalidTimeError +from pytz.exceptions import NonExistentTimeError +from pytz.exceptions import UnknownTimeZoneError +from pytz.lazy import LazyDict, LazyList, LazySet # noqa +from pytz.tzinfo import unpickler, BaseTzInfo +from pytz.tzfile import build_tzinfo + + +# The IANA (nee Olson) database is updated several times a year. +OLSON_VERSION = '2025b' +VERSION = '2025.2' # pip compatible version number. +__version__ = VERSION + +OLSEN_VERSION = OLSON_VERSION # Old releases had this misspelling + +__all__ = [ + 'timezone', 'utc', 'country_timezones', 'country_names', + 'AmbiguousTimeError', 'InvalidTimeError', + 'NonExistentTimeError', 'UnknownTimeZoneError', + 'all_timezones', 'all_timezones_set', + 'common_timezones', 'common_timezones_set', + 'BaseTzInfo', 'FixedOffset', +] + + +if sys.version_info[0] > 2: # Python 3.x + + # Python 3.x doesn't have unicode(), making writing code + # for Python 2.3 and Python 3.x a pain. + unicode = str + + def ascii(s): + r""" + >>> ascii('Hello') + 'Hello' + >>> ascii('\N{TRADE MARK SIGN}') #doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + UnicodeEncodeError: ... + """ + if type(s) == bytes: + s = s.decode('ASCII') + else: + s.encode('ASCII') # Raise an exception if not ASCII + return s # But the string - not a byte string. + +else: # Python 2.x + + def ascii(s): + r""" + >>> ascii('Hello') + 'Hello' + >>> ascii(u'Hello') + 'Hello' + >>> ascii(u'\N{TRADE MARK SIGN}') #doctest: +IGNORE_EXCEPTION_DETAIL + Traceback (most recent call last): + ... + UnicodeEncodeError: ... + """ + return s.encode('ASCII') + + +def open_resource(name): + """Open a resource from the zoneinfo subdir for reading. + + Uses the pkg_resources module if available and no standard file + found at the calculated location. + + It is possible to specify different location for zoneinfo + subdir by using the PYTZ_TZDATADIR environment variable. + """ + name_parts = name.lstrip('/').split('/') + for part in name_parts: + if part == os.path.pardir or os.sep in part: + raise ValueError('Bad path segment: %r' % part) + zoneinfo_dir = os.environ.get('PYTZ_TZDATADIR', None) + if zoneinfo_dir is not None: + filename = os.path.join(zoneinfo_dir, *name_parts) + else: + filename = os.path.join(os.path.dirname(__file__), + 'zoneinfo', *name_parts) + if not os.path.exists(filename): + # http://bugs.launchpad.net/bugs/383171 - we avoid using this + # unless absolutely necessary to help when a broken version of + # pkg_resources is installed. + try: + from pkg_resources import resource_stream + except ImportError: + resource_stream = None + + if resource_stream is not None: + return resource_stream(__name__, 'zoneinfo/' + name) + return open(filename, 'rb') + + +def resource_exists(name): + """Return true if the given resource exists""" + try: + if os.environ.get('PYTZ_SKIPEXISTSCHECK', ''): + # In "standard" distributions, we can assume that + # all the listed timezones are present. As an + # import-speed optimization, you can set the + # PYTZ_SKIPEXISTSCHECK flag to skip checking + # for the presence of the resource file on disk. + return True + open_resource(name).close() + return True + except IOError: + return False + + +_tzinfo_cache = {} + + +def timezone(zone): + r''' Return a datetime.tzinfo implementation for the given timezone + + >>> from datetime import datetime, timedelta + >>> utc = timezone('UTC') + >>> eastern = timezone('US/Eastern') + >>> eastern.zone + 'US/Eastern' + >>> timezone(unicode('US/Eastern')) is eastern + True + >>> utc_dt = datetime(2002, 10, 27, 6, 0, 0, tzinfo=utc) + >>> loc_dt = utc_dt.astimezone(eastern) + >>> fmt = '%Y-%m-%d %H:%M:%S %Z (%z)' + >>> loc_dt.strftime(fmt) + '2002-10-27 01:00:00 EST (-0500)' + >>> (loc_dt - timedelta(minutes=10)).strftime(fmt) + '2002-10-27 00:50:00 EST (-0500)' + >>> eastern.normalize(loc_dt - timedelta(minutes=10)).strftime(fmt) + '2002-10-27 01:50:00 EDT (-0400)' + >>> (loc_dt + timedelta(minutes=10)).strftime(fmt) + '2002-10-27 01:10:00 EST (-0500)' + + Raises UnknownTimeZoneError if passed an unknown zone. + + >>> try: + ... timezone('Asia/Shangri-La') + ... except UnknownTimeZoneError: + ... print('Unknown') + Unknown + + >>> try: + ... timezone(unicode('\N{TRADE MARK SIGN}')) + ... except UnknownTimeZoneError: + ... print('Unknown') + Unknown + + ''' + if zone is None: + raise UnknownTimeZoneError(None) + + if zone.upper() == 'UTC': + return utc + + try: + zone = ascii(zone) + except UnicodeEncodeError: + # All valid timezones are ASCII + raise UnknownTimeZoneError(zone) + + zone = _case_insensitive_zone_lookup(_unmunge_zone(zone)) + if zone not in _tzinfo_cache: + if zone in all_timezones_set: # noqa + fp = open_resource(zone) + try: + _tzinfo_cache[zone] = build_tzinfo(zone, fp) + finally: + fp.close() + else: + raise UnknownTimeZoneError(zone) + + return _tzinfo_cache[zone] + + +def _unmunge_zone(zone): + """Undo the time zone name munging done by older versions of pytz.""" + return zone.replace('_plus_', '+').replace('_minus_', '-') + + +_all_timezones_lower_to_standard = None + + +def _case_insensitive_zone_lookup(zone): + """case-insensitively matching timezone, else return zone unchanged""" + global _all_timezones_lower_to_standard + if _all_timezones_lower_to_standard is None: + _all_timezones_lower_to_standard = dict((tz.lower(), tz) for tz in _all_timezones_unchecked) # noqa + return _all_timezones_lower_to_standard.get(zone.lower()) or zone # noqa + + +ZERO = datetime.timedelta(0) +HOUR = datetime.timedelta(hours=1) + + +class UTC(BaseTzInfo): + """UTC + + Optimized UTC implementation. It unpickles using the single module global + instance defined beneath this class declaration. + """ + zone = "UTC" + + _utcoffset = ZERO + _dst = ZERO + _tzname = zone + + def fromutc(self, dt): + if dt.tzinfo is None: + return self.localize(dt) + return super(utc.__class__, self).fromutc(dt) + + def utcoffset(self, dt): + return ZERO + + def tzname(self, dt): + return "UTC" + + def dst(self, dt): + return ZERO + + def __reduce__(self): + return _UTC, () + + def localize(self, dt, is_dst=False): + '''Convert naive time to local time''' + if dt.tzinfo is not None: + raise ValueError('Not naive datetime (tzinfo is already set)') + return dt.replace(tzinfo=self) + + def normalize(self, dt, is_dst=False): + '''Correct the timezone information on the given datetime''' + if dt.tzinfo is self: + return dt + if dt.tzinfo is None: + raise ValueError('Naive time - no tzinfo set') + return dt.astimezone(self) + + def __repr__(self): + return "" + + def __str__(self): + return "UTC" + + +UTC = utc = UTC() # UTC is a singleton + + +def _UTC(): + """Factory function for utc unpickling. + + Makes sure that unpickling a utc instance always returns the same + module global. + + These examples belong in the UTC class above, but it is obscured; or in + the README.rst, but we are not depending on Python 2.4 so integrating + the README.rst examples with the unit tests is not trivial. + + >>> import datetime, pickle + >>> dt = datetime.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 + >>> utc is UTC is timezone('UTC') + True + >>> utc is timezone('GMT') + False + """ + return utc + + +_UTC.__safe_for_unpickling__ = True + + +def _p(*args): + """Factory function for unpickling pytz tzinfo instances. + + Just a wrapper around tzinfo.unpickler to save a few bytes in each pickle + by shortening the path. + """ + return unpickler(*args) + + +_p.__safe_for_unpickling__ = True + + +class _CountryTimezoneDict(LazyDict): + """Map ISO 3166 country code to a list of timezone names commonly used + in that country. + + iso3166_code is the two letter code used to identify the country. + + >>> def print_list(list_of_strings): + ... 'We use a helper so doctests work under Python 2.3 -> 3.x' + ... for s in list_of_strings: + ... print(s) + + >>> print_list(country_timezones['nz']) + Pacific/Auckland + Pacific/Chatham + >>> print_list(country_timezones['ch']) + Europe/Zurich + >>> print_list(country_timezones['CH']) + Europe/Zurich + >>> print_list(country_timezones[unicode('ch')]) + Europe/Zurich + >>> print_list(country_timezones['XXX']) + Traceback (most recent call last): + ... + KeyError: 'XXX' + + Previously, this information was exposed as a function rather than a + dictionary. This is still supported:: + + >>> print_list(country_timezones('nz')) + Pacific/Auckland + Pacific/Chatham + """ + def __call__(self, iso3166_code): + """Backwards compatibility.""" + return self[iso3166_code] + + def _fill(self): + data = {} + zone_tab = open_resource('zone.tab') + try: + for line in zone_tab: + line = line.decode('UTF-8') + if line.startswith('#'): + continue + code, coordinates, zone = line.split(None, 4)[:3] + if zone not in all_timezones_set: # noqa + continue + try: + data[code].append(zone) + except KeyError: + data[code] = [zone] + self.data = data + finally: + zone_tab.close() + + +country_timezones = _CountryTimezoneDict() + + +class _CountryNameDict(LazyDict): + '''Dictionary proving ISO3166 code -> English name. + + >>> print(country_names['au']) + Australia + ''' + def _fill(self): + data = {} + zone_tab = open_resource('iso3166.tab') + try: + for line in zone_tab.readlines(): + line = line.decode('UTF-8') + if line.startswith('#'): + continue + code, name = line.split(None, 1) + data[code] = name.strip() + self.data = data + finally: + zone_tab.close() + + +country_names = _CountryNameDict() + + +# Time-zone info based solely on fixed offsets + +class _FixedOffset(datetime.tzinfo): + + zone = None # to match the standard pytz API + + def __init__(self, minutes): + if abs(minutes) >= 1440: + raise ValueError("absolute offset is too large", minutes) + self._minutes = minutes + self._offset = datetime.timedelta(minutes=minutes) + + def utcoffset(self, dt): + return self._offset + + def __reduce__(self): + return FixedOffset, (self._minutes, ) + + def dst(self, dt): + return ZERO + + def tzname(self, dt): + return None + + def __repr__(self): + return 'pytz.FixedOffset(%d)' % self._minutes + + def localize(self, dt, is_dst=False): + '''Convert naive time to local time''' + if dt.tzinfo is not None: + raise ValueError('Not naive datetime (tzinfo is already set)') + return dt.replace(tzinfo=self) + + def normalize(self, dt, is_dst=False): + '''Correct the timezone information on the given datetime''' + if dt.tzinfo is self: + return dt + if dt.tzinfo is None: + raise ValueError('Naive time - no tzinfo set') + return dt.astimezone(self) + + +def FixedOffset(offset, _tzinfos={}): + """return a fixed-offset timezone based off a number of minutes. + + >>> one = FixedOffset(-330) + >>> one + pytz.FixedOffset(-330) + >>> str(one.utcoffset(datetime.datetime.now())) + '-1 day, 18:30:00' + >>> str(one.dst(datetime.datetime.now())) + '0:00:00' + + >>> two = FixedOffset(1380) + >>> two + pytz.FixedOffset(1380) + >>> str(two.utcoffset(datetime.datetime.now())) + '23:00:00' + >>> str(two.dst(datetime.datetime.now())) + '0:00:00' + + The datetime.timedelta must be between the range of -1 and 1 day, + non-inclusive. + + >>> FixedOffset(1440) + Traceback (most recent call last): + ... + ValueError: ('absolute offset is too large', 1440) + + >>> FixedOffset(-1440) + Traceback (most recent call last): + ... + ValueError: ('absolute offset is too large', -1440) + + An offset of 0 is special-cased to return UTC. + + >>> FixedOffset(0) is UTC + True + + There should always be only one instance of a FixedOffset per timedelta. + This should be true for multiple creation calls. + + >>> FixedOffset(-330) is one + True + >>> FixedOffset(1380) is two + True + + It should also be true for pickling. + + >>> import pickle + >>> pickle.loads(pickle.dumps(one)) is one + True + >>> pickle.loads(pickle.dumps(two)) is two + True + """ + if offset == 0: + return UTC + + info = _tzinfos.get(offset) + if info is None: + # We haven't seen this one before. we need to save it. + + # Use setdefault to avoid a race condition and make sure we have + # only one + info = _tzinfos.setdefault(offset, _FixedOffset(offset)) + + return info + + +FixedOffset.__safe_for_unpickling__ = True + + +def _test(): + import doctest + sys.path.insert(0, os.pardir) + import pytz + return doctest.testmod(pytz) + + +if __name__ == '__main__': + _test() +_all_timezones_unchecked = \ +['Africa/Abidjan', + 'Africa/Accra', + 'Africa/Addis_Ababa', + 'Africa/Algiers', + 'Africa/Asmara', + 'Africa/Asmera', + 'Africa/Bamako', + 'Africa/Bangui', + 'Africa/Banjul', + 'Africa/Bissau', + 'Africa/Blantyre', + 'Africa/Brazzaville', + 'Africa/Bujumbura', + 'Africa/Cairo', + 'Africa/Casablanca', + 'Africa/Ceuta', + 'Africa/Conakry', + 'Africa/Dakar', + 'Africa/Dar_es_Salaam', + 'Africa/Djibouti', + 'Africa/Douala', + 'Africa/El_Aaiun', + 'Africa/Freetown', + 'Africa/Gaborone', + 'Africa/Harare', + 'Africa/Johannesburg', + 'Africa/Juba', + 'Africa/Kampala', + 'Africa/Khartoum', + 'Africa/Kigali', + 'Africa/Kinshasa', + 'Africa/Lagos', + 'Africa/Libreville', + 'Africa/Lome', + 'Africa/Luanda', + 'Africa/Lubumbashi', + 'Africa/Lusaka', + 'Africa/Malabo', + 'Africa/Maputo', + 'Africa/Maseru', + 'Africa/Mbabane', + 'Africa/Mogadishu', + 'Africa/Monrovia', + 'Africa/Nairobi', + 'Africa/Ndjamena', + 'Africa/Niamey', + 'Africa/Nouakchott', + 'Africa/Ouagadougou', + 'Africa/Porto-Novo', + 'Africa/Sao_Tome', + 'Africa/Timbuktu', + 'Africa/Tripoli', + 'Africa/Tunis', + 'Africa/Windhoek', + 'America/Adak', + 'America/Anchorage', + 'America/Anguilla', + 'America/Antigua', + 'America/Araguaina', + 'America/Argentina/Buenos_Aires', + 'America/Argentina/Catamarca', + 'America/Argentina/ComodRivadavia', + 'America/Argentina/Cordoba', + 'America/Argentina/Jujuy', + 'America/Argentina/La_Rioja', + 'America/Argentina/Mendoza', + 'America/Argentina/Rio_Gallegos', + 'America/Argentina/Salta', + 'America/Argentina/San_Juan', + 'America/Argentina/San_Luis', + 'America/Argentina/Tucuman', + 'America/Argentina/Ushuaia', + 'America/Aruba', + 'America/Asuncion', + 'America/Atikokan', + 'America/Atka', + 'America/Bahia', + 'America/Bahia_Banderas', + 'America/Barbados', + 'America/Belem', + 'America/Belize', + 'America/Blanc-Sablon', + 'America/Boa_Vista', + 'America/Bogota', + 'America/Boise', + 'America/Buenos_Aires', + 'America/Cambridge_Bay', + 'America/Campo_Grande', + 'America/Cancun', + 'America/Caracas', + 'America/Catamarca', + 'America/Cayenne', + 'America/Cayman', + 'America/Chicago', + 'America/Chihuahua', + 'America/Ciudad_Juarez', + 'America/Coral_Harbour', + 'America/Cordoba', + 'America/Costa_Rica', + 'America/Coyhaique', + 'America/Creston', + 'America/Cuiaba', + 'America/Curacao', + 'America/Danmarkshavn', + 'America/Dawson', + 'America/Dawson_Creek', + 'America/Denver', + 'America/Detroit', + 'America/Dominica', + 'America/Edmonton', + 'America/Eirunepe', + 'America/El_Salvador', + 'America/Ensenada', + 'America/Fort_Nelson', + 'America/Fort_Wayne', + 'America/Fortaleza', + 'America/Glace_Bay', + 'America/Godthab', + 'America/Goose_Bay', + 'America/Grand_Turk', + 'America/Grenada', + 'America/Guadeloupe', + 'America/Guatemala', + 'America/Guayaquil', + 'America/Guyana', + 'America/Halifax', + 'America/Havana', + 'America/Hermosillo', + 'America/Indiana/Indianapolis', + 'America/Indiana/Knox', + 'America/Indiana/Marengo', + 'America/Indiana/Petersburg', + 'America/Indiana/Tell_City', + 'America/Indiana/Vevay', + 'America/Indiana/Vincennes', + 'America/Indiana/Winamac', + 'America/Indianapolis', + 'America/Inuvik', + 'America/Iqaluit', + 'America/Jamaica', + 'America/Jujuy', + 'America/Juneau', + 'America/Kentucky/Louisville', + 'America/Kentucky/Monticello', + 'America/Knox_IN', + 'America/Kralendijk', + 'America/La_Paz', + 'America/Lima', + 'America/Los_Angeles', + 'America/Louisville', + 'America/Lower_Princes', + 'America/Maceio', + 'America/Managua', + 'America/Manaus', + 'America/Marigot', + 'America/Martinique', + 'America/Matamoros', + 'America/Mazatlan', + 'America/Mendoza', + 'America/Menominee', + 'America/Merida', + 'America/Metlakatla', + 'America/Mexico_City', + 'America/Miquelon', + 'America/Moncton', + 'America/Monterrey', + 'America/Montevideo', + 'America/Montreal', + 'America/Montserrat', + 'America/Nassau', + 'America/New_York', + 'America/Nipigon', + 'America/Nome', + 'America/Noronha', + 'America/North_Dakota/Beulah', + 'America/North_Dakota/Center', + 'America/North_Dakota/New_Salem', + 'America/Nuuk', + 'America/Ojinaga', + 'America/Panama', + 'America/Pangnirtung', + 'America/Paramaribo', + 'America/Phoenix', + 'America/Port-au-Prince', + 'America/Port_of_Spain', + 'America/Porto_Acre', + 'America/Porto_Velho', + 'America/Puerto_Rico', + 'America/Punta_Arenas', + 'America/Rainy_River', + 'America/Rankin_Inlet', + 'America/Recife', + 'America/Regina', + 'America/Resolute', + 'America/Rio_Branco', + 'America/Rosario', + 'America/Santa_Isabel', + 'America/Santarem', + 'America/Santiago', + 'America/Santo_Domingo', + 'America/Sao_Paulo', + 'America/Scoresbysund', + 'America/Shiprock', + 'America/Sitka', + 'America/St_Barthelemy', + 'America/St_Johns', + 'America/St_Kitts', + 'America/St_Lucia', + 'America/St_Thomas', + 'America/St_Vincent', + 'America/Swift_Current', + 'America/Tegucigalpa', + 'America/Thule', + 'America/Thunder_Bay', + 'America/Tijuana', + 'America/Toronto', + 'America/Tortola', + 'America/Vancouver', + 'America/Virgin', + 'America/Whitehorse', + 'America/Winnipeg', + 'America/Yakutat', + 'America/Yellowknife', + 'Antarctica/Casey', + 'Antarctica/Davis', + 'Antarctica/DumontDUrville', + 'Antarctica/Macquarie', + 'Antarctica/Mawson', + 'Antarctica/McMurdo', + 'Antarctica/Palmer', + 'Antarctica/Rothera', + 'Antarctica/South_Pole', + 'Antarctica/Syowa', + 'Antarctica/Troll', + 'Antarctica/Vostok', + 'Arctic/Longyearbyen', + 'Asia/Aden', + 'Asia/Almaty', + 'Asia/Amman', + 'Asia/Anadyr', + 'Asia/Aqtau', + 'Asia/Aqtobe', + 'Asia/Ashgabat', + 'Asia/Ashkhabad', + 'Asia/Atyrau', + 'Asia/Baghdad', + 'Asia/Bahrain', + 'Asia/Baku', + 'Asia/Bangkok', + 'Asia/Barnaul', + 'Asia/Beirut', + 'Asia/Bishkek', + 'Asia/Brunei', + 'Asia/Calcutta', + 'Asia/Chita', + 'Asia/Choibalsan', + 'Asia/Chongqing', + 'Asia/Chungking', + 'Asia/Colombo', + 'Asia/Dacca', + 'Asia/Damascus', + 'Asia/Dhaka', + 'Asia/Dili', + 'Asia/Dubai', + 'Asia/Dushanbe', + 'Asia/Famagusta', + 'Asia/Gaza', + 'Asia/Harbin', + 'Asia/Hebron', + 'Asia/Ho_Chi_Minh', + 'Asia/Hong_Kong', + 'Asia/Hovd', + 'Asia/Irkutsk', + 'Asia/Istanbul', + 'Asia/Jakarta', + 'Asia/Jayapura', + 'Asia/Jerusalem', + 'Asia/Kabul', + 'Asia/Kamchatka', + 'Asia/Karachi', + 'Asia/Kashgar', + 'Asia/Kathmandu', + 'Asia/Katmandu', + 'Asia/Khandyga', + 'Asia/Kolkata', + 'Asia/Krasnoyarsk', + 'Asia/Kuala_Lumpur', + 'Asia/Kuching', + 'Asia/Kuwait', + 'Asia/Macao', + 'Asia/Macau', + 'Asia/Magadan', + 'Asia/Makassar', + 'Asia/Manila', + 'Asia/Muscat', + 'Asia/Nicosia', + 'Asia/Novokuznetsk', + 'Asia/Novosibirsk', + 'Asia/Omsk', + 'Asia/Oral', + 'Asia/Phnom_Penh', + 'Asia/Pontianak', + 'Asia/Pyongyang', + 'Asia/Qatar', + 'Asia/Qostanay', + 'Asia/Qyzylorda', + 'Asia/Rangoon', + 'Asia/Riyadh', + 'Asia/Saigon', + 'Asia/Sakhalin', + 'Asia/Samarkand', + 'Asia/Seoul', + 'Asia/Shanghai', + 'Asia/Singapore', + 'Asia/Srednekolymsk', + 'Asia/Taipei', + 'Asia/Tashkent', + 'Asia/Tbilisi', + 'Asia/Tehran', + 'Asia/Tel_Aviv', + 'Asia/Thimbu', + 'Asia/Thimphu', + 'Asia/Tokyo', + 'Asia/Tomsk', + 'Asia/Ujung_Pandang', + 'Asia/Ulaanbaatar', + 'Asia/Ulan_Bator', + 'Asia/Urumqi', + 'Asia/Ust-Nera', + 'Asia/Vientiane', + 'Asia/Vladivostok', + 'Asia/Yakutsk', + 'Asia/Yangon', + 'Asia/Yekaterinburg', + 'Asia/Yerevan', + 'Atlantic/Azores', + 'Atlantic/Bermuda', + 'Atlantic/Canary', + 'Atlantic/Cape_Verde', + 'Atlantic/Faeroe', + 'Atlantic/Faroe', + 'Atlantic/Jan_Mayen', + 'Atlantic/Madeira', + 'Atlantic/Reykjavik', + 'Atlantic/South_Georgia', + 'Atlantic/St_Helena', + 'Atlantic/Stanley', + 'Australia/ACT', + 'Australia/Adelaide', + 'Australia/Brisbane', + 'Australia/Broken_Hill', + 'Australia/Canberra', + 'Australia/Currie', + 'Australia/Darwin', + 'Australia/Eucla', + 'Australia/Hobart', + 'Australia/LHI', + 'Australia/Lindeman', + 'Australia/Lord_Howe', + 'Australia/Melbourne', + 'Australia/NSW', + 'Australia/North', + 'Australia/Perth', + 'Australia/Queensland', + 'Australia/South', + 'Australia/Sydney', + 'Australia/Tasmania', + 'Australia/Victoria', + 'Australia/West', + 'Australia/Yancowinna', + 'Brazil/Acre', + 'Brazil/DeNoronha', + 'Brazil/East', + 'Brazil/West', + 'CET', + 'CST6CDT', + 'Canada/Atlantic', + 'Canada/Central', + 'Canada/Eastern', + 'Canada/Mountain', + 'Canada/Newfoundland', + 'Canada/Pacific', + 'Canada/Saskatchewan', + 'Canada/Yukon', + 'Chile/Continental', + 'Chile/EasterIsland', + 'Cuba', + 'EET', + 'EST', + 'EST5EDT', + 'Egypt', + 'Eire', + 'Etc/GMT', + 'Etc/GMT+0', + 'Etc/GMT+1', + 'Etc/GMT+10', + 'Etc/GMT+11', + 'Etc/GMT+12', + 'Etc/GMT+2', + 'Etc/GMT+3', + 'Etc/GMT+4', + 'Etc/GMT+5', + 'Etc/GMT+6', + 'Etc/GMT+7', + 'Etc/GMT+8', + 'Etc/GMT+9', + 'Etc/GMT-0', + 'Etc/GMT-1', + 'Etc/GMT-10', + 'Etc/GMT-11', + 'Etc/GMT-12', + 'Etc/GMT-13', + 'Etc/GMT-14', + 'Etc/GMT-2', + 'Etc/GMT-3', + 'Etc/GMT-4', + 'Etc/GMT-5', + 'Etc/GMT-6', + 'Etc/GMT-7', + 'Etc/GMT-8', + 'Etc/GMT-9', + 'Etc/GMT0', + 'Etc/Greenwich', + 'Etc/UCT', + 'Etc/UTC', + 'Etc/Universal', + 'Etc/Zulu', + 'Europe/Amsterdam', + 'Europe/Andorra', + 'Europe/Astrakhan', + 'Europe/Athens', + 'Europe/Belfast', + 'Europe/Belgrade', + 'Europe/Berlin', + 'Europe/Bratislava', + 'Europe/Brussels', + 'Europe/Bucharest', + 'Europe/Budapest', + 'Europe/Busingen', + 'Europe/Chisinau', + 'Europe/Copenhagen', + 'Europe/Dublin', + 'Europe/Gibraltar', + 'Europe/Guernsey', + 'Europe/Helsinki', + 'Europe/Isle_of_Man', + 'Europe/Istanbul', + 'Europe/Jersey', + 'Europe/Kaliningrad', + 'Europe/Kiev', + 'Europe/Kirov', + 'Europe/Kyiv', + 'Europe/Lisbon', + 'Europe/Ljubljana', + 'Europe/London', + 'Europe/Luxembourg', + 'Europe/Madrid', + 'Europe/Malta', + 'Europe/Mariehamn', + 'Europe/Minsk', + 'Europe/Monaco', + 'Europe/Moscow', + 'Europe/Nicosia', + 'Europe/Oslo', + 'Europe/Paris', + 'Europe/Podgorica', + 'Europe/Prague', + 'Europe/Riga', + 'Europe/Rome', + 'Europe/Samara', + 'Europe/San_Marino', + 'Europe/Sarajevo', + 'Europe/Saratov', + 'Europe/Simferopol', + 'Europe/Skopje', + 'Europe/Sofia', + 'Europe/Stockholm', + 'Europe/Tallinn', + 'Europe/Tirane', + 'Europe/Tiraspol', + 'Europe/Ulyanovsk', + 'Europe/Uzhgorod', + 'Europe/Vaduz', + 'Europe/Vatican', + 'Europe/Vienna', + 'Europe/Vilnius', + 'Europe/Volgograd', + 'Europe/Warsaw', + 'Europe/Zagreb', + 'Europe/Zaporozhye', + 'Europe/Zurich', + 'GB', + 'GB-Eire', + 'GMT', + 'GMT+0', + 'GMT-0', + 'GMT0', + 'Greenwich', + 'HST', + 'Hongkong', + 'Iceland', + 'Indian/Antananarivo', + 'Indian/Chagos', + 'Indian/Christmas', + 'Indian/Cocos', + 'Indian/Comoro', + 'Indian/Kerguelen', + 'Indian/Mahe', + 'Indian/Maldives', + 'Indian/Mauritius', + 'Indian/Mayotte', + 'Indian/Reunion', + 'Iran', + 'Israel', + 'Jamaica', + 'Japan', + 'Kwajalein', + 'Libya', + 'MET', + 'MST', + 'MST7MDT', + 'Mexico/BajaNorte', + 'Mexico/BajaSur', + 'Mexico/General', + 'NZ', + 'NZ-CHAT', + 'Navajo', + 'PRC', + 'PST8PDT', + 'Pacific/Apia', + 'Pacific/Auckland', + 'Pacific/Bougainville', + 'Pacific/Chatham', + 'Pacific/Chuuk', + 'Pacific/Easter', + 'Pacific/Efate', + 'Pacific/Enderbury', + 'Pacific/Fakaofo', + 'Pacific/Fiji', + 'Pacific/Funafuti', + 'Pacific/Galapagos', + 'Pacific/Gambier', + 'Pacific/Guadalcanal', + 'Pacific/Guam', + 'Pacific/Honolulu', + 'Pacific/Johnston', + 'Pacific/Kanton', + 'Pacific/Kiritimati', + 'Pacific/Kosrae', + 'Pacific/Kwajalein', + 'Pacific/Majuro', + 'Pacific/Marquesas', + 'Pacific/Midway', + 'Pacific/Nauru', + 'Pacific/Niue', + 'Pacific/Norfolk', + 'Pacific/Noumea', + 'Pacific/Pago_Pago', + 'Pacific/Palau', + 'Pacific/Pitcairn', + 'Pacific/Pohnpei', + 'Pacific/Ponape', + 'Pacific/Port_Moresby', + 'Pacific/Rarotonga', + 'Pacific/Saipan', + 'Pacific/Samoa', + 'Pacific/Tahiti', + 'Pacific/Tarawa', + 'Pacific/Tongatapu', + 'Pacific/Truk', + 'Pacific/Wake', + 'Pacific/Wallis', + 'Pacific/Yap', + 'Poland', + 'Portugal', + 'ROC', + 'ROK', + 'Singapore', + 'Turkey', + 'UCT', + 'US/Alaska', + 'US/Aleutian', + 'US/Arizona', + 'US/Central', + 'US/East-Indiana', + 'US/Eastern', + 'US/Hawaii', + 'US/Indiana-Starke', + 'US/Michigan', + 'US/Mountain', + 'US/Pacific', + 'US/Samoa', + 'UTC', + 'Universal', + 'W-SU', + 'WET', + 'Zulu'] +all_timezones = LazyList( + tz for tz in _all_timezones_unchecked if resource_exists(tz)) + +all_timezones_set = LazySet(all_timezones) +common_timezones = \ +['Africa/Abidjan', + 'Africa/Accra', + 'Africa/Addis_Ababa', + 'Africa/Algiers', + 'Africa/Asmara', + 'Africa/Bamako', + 'Africa/Bangui', + 'Africa/Banjul', + 'Africa/Bissau', + 'Africa/Blantyre', + 'Africa/Brazzaville', + 'Africa/Bujumbura', + 'Africa/Cairo', + 'Africa/Casablanca', + 'Africa/Ceuta', + 'Africa/Conakry', + 'Africa/Dakar', + 'Africa/Dar_es_Salaam', + 'Africa/Djibouti', + 'Africa/Douala', + 'Africa/El_Aaiun', + 'Africa/Freetown', + 'Africa/Gaborone', + 'Africa/Harare', + 'Africa/Johannesburg', + 'Africa/Juba', + 'Africa/Kampala', + 'Africa/Khartoum', + 'Africa/Kigali', + 'Africa/Kinshasa', + 'Africa/Lagos', + 'Africa/Libreville', + 'Africa/Lome', + 'Africa/Luanda', + 'Africa/Lubumbashi', + 'Africa/Lusaka', + 'Africa/Malabo', + 'Africa/Maputo', + 'Africa/Maseru', + 'Africa/Mbabane', + 'Africa/Mogadishu', + 'Africa/Monrovia', + 'Africa/Nairobi', + 'Africa/Ndjamena', + 'Africa/Niamey', + 'Africa/Nouakchott', + 'Africa/Ouagadougou', + 'Africa/Porto-Novo', + 'Africa/Sao_Tome', + 'Africa/Tripoli', + 'Africa/Tunis', + 'Africa/Windhoek', + 'America/Adak', + 'America/Anchorage', + 'America/Anguilla', + 'America/Antigua', + 'America/Araguaina', + 'America/Argentina/Buenos_Aires', + 'America/Argentina/Catamarca', + 'America/Argentina/Cordoba', + 'America/Argentina/Jujuy', + 'America/Argentina/La_Rioja', + 'America/Argentina/Mendoza', + 'America/Argentina/Rio_Gallegos', + 'America/Argentina/Salta', + 'America/Argentina/San_Juan', + 'America/Argentina/San_Luis', + 'America/Argentina/Tucuman', + 'America/Argentina/Ushuaia', + 'America/Aruba', + 'America/Asuncion', + 'America/Atikokan', + 'America/Bahia', + 'America/Bahia_Banderas', + 'America/Barbados', + 'America/Belem', + 'America/Belize', + 'America/Blanc-Sablon', + 'America/Boa_Vista', + 'America/Bogota', + 'America/Boise', + 'America/Cambridge_Bay', + 'America/Campo_Grande', + 'America/Cancun', + 'America/Caracas', + 'America/Cayenne', + 'America/Cayman', + 'America/Chicago', + 'America/Chihuahua', + 'America/Ciudad_Juarez', + 'America/Costa_Rica', + 'America/Coyhaique', + 'America/Creston', + 'America/Cuiaba', + 'America/Curacao', + 'America/Danmarkshavn', + 'America/Dawson', + 'America/Dawson_Creek', + 'America/Denver', + 'America/Detroit', + 'America/Dominica', + 'America/Edmonton', + 'America/Eirunepe', + 'America/El_Salvador', + 'America/Fort_Nelson', + 'America/Fortaleza', + 'America/Glace_Bay', + 'America/Goose_Bay', + 'America/Grand_Turk', + 'America/Grenada', + 'America/Guadeloupe', + 'America/Guatemala', + 'America/Guayaquil', + 'America/Guyana', + 'America/Halifax', + 'America/Havana', + 'America/Hermosillo', + 'America/Indiana/Indianapolis', + 'America/Indiana/Knox', + 'America/Indiana/Marengo', + 'America/Indiana/Petersburg', + 'America/Indiana/Tell_City', + 'America/Indiana/Vevay', + 'America/Indiana/Vincennes', + 'America/Indiana/Winamac', + 'America/Inuvik', + 'America/Iqaluit', + 'America/Jamaica', + 'America/Juneau', + 'America/Kentucky/Louisville', + 'America/Kentucky/Monticello', + 'America/Kralendijk', + 'America/La_Paz', + 'America/Lima', + 'America/Los_Angeles', + 'America/Lower_Princes', + 'America/Maceio', + 'America/Managua', + 'America/Manaus', + 'America/Marigot', + 'America/Martinique', + 'America/Matamoros', + 'America/Mazatlan', + 'America/Menominee', + 'America/Merida', + 'America/Metlakatla', + 'America/Mexico_City', + 'America/Miquelon', + 'America/Moncton', + 'America/Monterrey', + 'America/Montevideo', + 'America/Montserrat', + 'America/Nassau', + 'America/New_York', + 'America/Nome', + 'America/Noronha', + 'America/North_Dakota/Beulah', + 'America/North_Dakota/Center', + 'America/North_Dakota/New_Salem', + 'America/Nuuk', + 'America/Ojinaga', + 'America/Panama', + 'America/Paramaribo', + 'America/Phoenix', + 'America/Port-au-Prince', + 'America/Port_of_Spain', + 'America/Porto_Velho', + 'America/Puerto_Rico', + 'America/Punta_Arenas', + 'America/Rankin_Inlet', + 'America/Recife', + 'America/Regina', + 'America/Resolute', + 'America/Rio_Branco', + 'America/Santarem', + 'America/Santiago', + 'America/Santo_Domingo', + 'America/Sao_Paulo', + 'America/Scoresbysund', + 'America/Sitka', + 'America/St_Barthelemy', + 'America/St_Johns', + 'America/St_Kitts', + 'America/St_Lucia', + 'America/St_Thomas', + 'America/St_Vincent', + 'America/Swift_Current', + 'America/Tegucigalpa', + 'America/Thule', + 'America/Tijuana', + 'America/Toronto', + 'America/Tortola', + 'America/Vancouver', + 'America/Whitehorse', + 'America/Winnipeg', + 'America/Yakutat', + 'Antarctica/Casey', + 'Antarctica/Davis', + 'Antarctica/DumontDUrville', + 'Antarctica/Macquarie', + 'Antarctica/Mawson', + 'Antarctica/McMurdo', + 'Antarctica/Palmer', + 'Antarctica/Rothera', + 'Antarctica/Syowa', + 'Antarctica/Troll', + 'Antarctica/Vostok', + 'Arctic/Longyearbyen', + 'Asia/Aden', + 'Asia/Almaty', + 'Asia/Amman', + 'Asia/Anadyr', + 'Asia/Aqtau', + 'Asia/Aqtobe', + 'Asia/Ashgabat', + 'Asia/Atyrau', + 'Asia/Baghdad', + 'Asia/Bahrain', + 'Asia/Baku', + 'Asia/Bangkok', + 'Asia/Barnaul', + 'Asia/Beirut', + 'Asia/Bishkek', + 'Asia/Brunei', + 'Asia/Chita', + 'Asia/Colombo', + 'Asia/Damascus', + 'Asia/Dhaka', + 'Asia/Dili', + 'Asia/Dubai', + 'Asia/Dushanbe', + 'Asia/Famagusta', + 'Asia/Gaza', + 'Asia/Hebron', + 'Asia/Ho_Chi_Minh', + 'Asia/Hong_Kong', + 'Asia/Hovd', + 'Asia/Irkutsk', + 'Asia/Jakarta', + 'Asia/Jayapura', + 'Asia/Jerusalem', + 'Asia/Kabul', + 'Asia/Kamchatka', + 'Asia/Karachi', + 'Asia/Kathmandu', + 'Asia/Khandyga', + 'Asia/Kolkata', + 'Asia/Krasnoyarsk', + 'Asia/Kuala_Lumpur', + 'Asia/Kuching', + 'Asia/Kuwait', + 'Asia/Macau', + 'Asia/Magadan', + 'Asia/Makassar', + 'Asia/Manila', + 'Asia/Muscat', + 'Asia/Nicosia', + 'Asia/Novokuznetsk', + 'Asia/Novosibirsk', + 'Asia/Omsk', + 'Asia/Oral', + 'Asia/Phnom_Penh', + 'Asia/Pontianak', + 'Asia/Pyongyang', + 'Asia/Qatar', + 'Asia/Qostanay', + 'Asia/Qyzylorda', + 'Asia/Riyadh', + 'Asia/Sakhalin', + 'Asia/Samarkand', + 'Asia/Seoul', + 'Asia/Shanghai', + 'Asia/Singapore', + 'Asia/Srednekolymsk', + 'Asia/Taipei', + 'Asia/Tashkent', + 'Asia/Tbilisi', + 'Asia/Tehran', + 'Asia/Thimphu', + 'Asia/Tokyo', + 'Asia/Tomsk', + 'Asia/Ulaanbaatar', + 'Asia/Urumqi', + 'Asia/Ust-Nera', + 'Asia/Vientiane', + 'Asia/Vladivostok', + 'Asia/Yakutsk', + 'Asia/Yangon', + 'Asia/Yekaterinburg', + 'Asia/Yerevan', + 'Atlantic/Azores', + 'Atlantic/Bermuda', + 'Atlantic/Canary', + 'Atlantic/Cape_Verde', + 'Atlantic/Faroe', + 'Atlantic/Madeira', + 'Atlantic/Reykjavik', + 'Atlantic/South_Georgia', + 'Atlantic/St_Helena', + 'Atlantic/Stanley', + 'Australia/Adelaide', + 'Australia/Brisbane', + 'Australia/Broken_Hill', + 'Australia/Darwin', + 'Australia/Eucla', + 'Australia/Hobart', + 'Australia/Lindeman', + 'Australia/Lord_Howe', + 'Australia/Melbourne', + 'Australia/Perth', + 'Australia/Sydney', + 'Canada/Atlantic', + 'Canada/Central', + 'Canada/Eastern', + 'Canada/Mountain', + 'Canada/Newfoundland', + 'Canada/Pacific', + 'Europe/Amsterdam', + 'Europe/Andorra', + 'Europe/Astrakhan', + 'Europe/Athens', + 'Europe/Belgrade', + 'Europe/Berlin', + 'Europe/Bratislava', + 'Europe/Brussels', + 'Europe/Bucharest', + 'Europe/Budapest', + 'Europe/Busingen', + 'Europe/Chisinau', + 'Europe/Copenhagen', + 'Europe/Dublin', + 'Europe/Gibraltar', + 'Europe/Guernsey', + 'Europe/Helsinki', + 'Europe/Isle_of_Man', + 'Europe/Istanbul', + 'Europe/Jersey', + 'Europe/Kaliningrad', + 'Europe/Kirov', + 'Europe/Kyiv', + 'Europe/Lisbon', + 'Europe/Ljubljana', + 'Europe/London', + 'Europe/Luxembourg', + 'Europe/Madrid', + 'Europe/Malta', + 'Europe/Mariehamn', + 'Europe/Minsk', + 'Europe/Monaco', + 'Europe/Moscow', + 'Europe/Oslo', + 'Europe/Paris', + 'Europe/Podgorica', + 'Europe/Prague', + 'Europe/Riga', + 'Europe/Rome', + 'Europe/Samara', + 'Europe/San_Marino', + 'Europe/Sarajevo', + 'Europe/Saratov', + 'Europe/Simferopol', + 'Europe/Skopje', + 'Europe/Sofia', + 'Europe/Stockholm', + 'Europe/Tallinn', + 'Europe/Tirane', + 'Europe/Ulyanovsk', + 'Europe/Vaduz', + 'Europe/Vatican', + 'Europe/Vienna', + 'Europe/Vilnius', + 'Europe/Volgograd', + 'Europe/Warsaw', + 'Europe/Zagreb', + 'Europe/Zurich', + 'GMT', + 'Indian/Antananarivo', + 'Indian/Chagos', + 'Indian/Christmas', + 'Indian/Cocos', + 'Indian/Comoro', + 'Indian/Kerguelen', + 'Indian/Mahe', + 'Indian/Maldives', + 'Indian/Mauritius', + 'Indian/Mayotte', + 'Indian/Reunion', + 'Pacific/Apia', + 'Pacific/Auckland', + 'Pacific/Bougainville', + 'Pacific/Chatham', + 'Pacific/Chuuk', + 'Pacific/Easter', + 'Pacific/Efate', + 'Pacific/Fakaofo', + 'Pacific/Fiji', + 'Pacific/Funafuti', + 'Pacific/Galapagos', + 'Pacific/Gambier', + 'Pacific/Guadalcanal', + 'Pacific/Guam', + 'Pacific/Honolulu', + 'Pacific/Kanton', + 'Pacific/Kiritimati', + 'Pacific/Kosrae', + 'Pacific/Kwajalein', + 'Pacific/Majuro', + 'Pacific/Marquesas', + 'Pacific/Midway', + 'Pacific/Nauru', + 'Pacific/Niue', + 'Pacific/Norfolk', + 'Pacific/Noumea', + 'Pacific/Pago_Pago', + 'Pacific/Palau', + 'Pacific/Pitcairn', + 'Pacific/Pohnpei', + 'Pacific/Port_Moresby', + 'Pacific/Rarotonga', + 'Pacific/Saipan', + 'Pacific/Tahiti', + 'Pacific/Tarawa', + 'Pacific/Tongatapu', + 'Pacific/Wake', + 'Pacific/Wallis', + 'US/Alaska', + 'US/Arizona', + 'US/Central', + 'US/Eastern', + 'US/Hawaii', + 'US/Mountain', + 'US/Pacific', + 'UTC'] +common_timezones = LazyList( + tz for tz in common_timezones if tz in all_timezones) + +common_timezones_set = LazySet(common_timezones) diff --git a/lib/python3.12/site-packages/pytz/exceptions.py b/lib/python3.12/site-packages/pytz/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..4b20bde9ff9240ce8cc578e480f4d9aa8555bab4 --- /dev/null +++ b/lib/python3.12/site-packages/pytz/exceptions.py @@ -0,0 +1,59 @@ +''' +Custom exceptions raised by pytz. +''' + +__all__ = [ + 'UnknownTimeZoneError', 'InvalidTimeError', 'AmbiguousTimeError', + 'NonExistentTimeError', +] + + +class Error(Exception): + '''Base class for all exceptions raised by the pytz library''' + + +class UnknownTimeZoneError(KeyError, Error): + '''Exception raised when pytz is passed an unknown timezone. + + >>> isinstance(UnknownTimeZoneError(), LookupError) + True + + This class is actually a subclass of KeyError to provide backwards + compatibility with code relying on the undocumented behavior of earlier + pytz releases. + + >>> isinstance(UnknownTimeZoneError(), KeyError) + True + + And also a subclass of pytz.exceptions.Error, as are other pytz + exceptions. + + >>> isinstance(UnknownTimeZoneError(), Error) + True + + ''' + pass + + +class InvalidTimeError(Error): + '''Base class for invalid time exceptions.''' + + +class AmbiguousTimeError(InvalidTimeError): + '''Exception raised when attempting to create an ambiguous wallclock time. + + At the end of a DST transition period, a particular wallclock time will + occur twice (once before the clocks are set back, once after). Both + possibilities may be correct, unless further information is supplied. + + See DstTzInfo.normalize() for more info + ''' + + +class NonExistentTimeError(InvalidTimeError): + '''Exception raised when attempting to create a wallclock time that + cannot exist. + + At the start of a DST transition period, the wallclock time jumps forward. + The instants jumped over never occur. + ''' diff --git a/lib/python3.12/site-packages/pytz/lazy.py b/lib/python3.12/site-packages/pytz/lazy.py new file mode 100644 index 0000000000000000000000000000000000000000..39344fc1f8c77d5ec43539d0c8e655f4b5d7d6f6 --- /dev/null +++ b/lib/python3.12/site-packages/pytz/lazy.py @@ -0,0 +1,172 @@ +from threading import RLock +try: + from collections.abc import Mapping as DictMixin +except ImportError: # Python < 3.3 + try: + from UserDict import DictMixin # Python 2 + except ImportError: # Python 3.0-3.3 + from collections import Mapping as DictMixin + + +# With lazy loading, we might end up with multiple threads triggering +# it at the same time. We need a lock. +_fill_lock = RLock() + + +class LazyDict(DictMixin): + """Dictionary populated on first use.""" + data = None + + def __getitem__(self, key): + if self.data is None: + _fill_lock.acquire() + try: + if self.data is None: + self._fill() + finally: + _fill_lock.release() + return self.data[key.upper()] + + def __contains__(self, key): + if self.data is None: + _fill_lock.acquire() + try: + if self.data is None: + self._fill() + finally: + _fill_lock.release() + return key in self.data + + def __iter__(self): + if self.data is None: + _fill_lock.acquire() + try: + if self.data is None: + self._fill() + finally: + _fill_lock.release() + return iter(self.data) + + def __len__(self): + if self.data is None: + _fill_lock.acquire() + try: + if self.data is None: + self._fill() + finally: + _fill_lock.release() + return len(self.data) + + def keys(self): + if self.data is None: + _fill_lock.acquire() + try: + if self.data is None: + self._fill() + finally: + _fill_lock.release() + return self.data.keys() + + +class LazyList(list): + """List populated on first use.""" + + _props = [ + '__str__', '__repr__', '__unicode__', + '__hash__', '__sizeof__', '__cmp__', + '__lt__', '__le__', '__eq__', '__ne__', '__gt__', '__ge__', + 'append', 'count', 'index', 'extend', 'insert', 'pop', 'remove', + 'reverse', 'sort', '__add__', '__radd__', '__iadd__', '__mul__', + '__rmul__', '__imul__', '__contains__', '__len__', '__nonzero__', + '__getitem__', '__setitem__', '__delitem__', '__iter__', + '__reversed__', '__getslice__', '__setslice__', '__delslice__'] + + def __new__(cls, fill_iter=None): + + if fill_iter is None: + return list() + + # We need a new class as we will be dynamically messing with its + # methods. + class LazyList(list): + pass + + fill_iter = [fill_iter] + + def lazy(name): + def _lazy(self, *args, **kw): + _fill_lock.acquire() + try: + if len(fill_iter) > 0: + list.extend(self, fill_iter.pop()) + for method_name in cls._props: + delattr(LazyList, method_name) + finally: + _fill_lock.release() + return getattr(list, name)(self, *args, **kw) + return _lazy + + for name in cls._props: + setattr(LazyList, name, lazy(name)) + + new_list = LazyList() + return new_list + +# Not all versions of Python declare the same magic methods. +# Filter out properties that don't exist in this version of Python +# from the list. +LazyList._props = [prop for prop in LazyList._props if hasattr(list, prop)] + + +class LazySet(set): + """Set populated on first use.""" + + _props = ( + '__str__', '__repr__', '__unicode__', + '__hash__', '__sizeof__', '__cmp__', + '__lt__', '__le__', '__eq__', '__ne__', '__gt__', '__ge__', + '__contains__', '__len__', '__nonzero__', + '__getitem__', '__setitem__', '__delitem__', '__iter__', + '__sub__', '__and__', '__xor__', '__or__', + '__rsub__', '__rand__', '__rxor__', '__ror__', + '__isub__', '__iand__', '__ixor__', '__ior__', + 'add', 'clear', 'copy', 'difference', 'difference_update', + 'discard', 'intersection', 'intersection_update', 'isdisjoint', + 'issubset', 'issuperset', 'pop', 'remove', + 'symmetric_difference', 'symmetric_difference_update', + 'union', 'update') + + def __new__(cls, fill_iter=None): + + if fill_iter is None: + return set() + + class LazySet(set): + pass + + fill_iter = [fill_iter] + + def lazy(name): + def _lazy(self, *args, **kw): + _fill_lock.acquire() + try: + if len(fill_iter) > 0: + for i in fill_iter.pop(): + set.add(self, i) + for method_name in cls._props: + delattr(LazySet, method_name) + finally: + _fill_lock.release() + return getattr(set, name)(self, *args, **kw) + return _lazy + + for name in cls._props: + setattr(LazySet, name, lazy(name)) + + new_set = LazySet() + return new_set + +# Not all versions of Python declare the same magic methods. +# Filter out properties that don't exist in this version of Python +# from the list. +LazySet._props = [prop for prop in LazySet._props if hasattr(set, prop)] diff --git a/lib/python3.12/site-packages/pytz/reference.py b/lib/python3.12/site-packages/pytz/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..f765ca0af0b24e66dc3b7d51b9bf97e71b2b67aa --- /dev/null +++ b/lib/python3.12/site-packages/pytz/reference.py @@ -0,0 +1,140 @@ +''' +Reference tzinfo implementations from the Python docs. +Used for testing against as they are only correct for the years +1987 to 2006. Do not use these for real code. +''' + +from datetime import tzinfo, timedelta, datetime +from pytz import HOUR, ZERO, UTC + +__all__ = [ + 'FixedOffset', + 'LocalTimezone', + 'USTimeZone', + 'Eastern', + 'Central', + 'Mountain', + 'Pacific', + 'UTC' +] + + +# A class building tzinfo objects for fixed-offset time zones. +# Note that FixedOffset(0, "UTC") is a different way to build a +# UTC tzinfo object. +class FixedOffset(tzinfo): + """Fixed offset in minutes east from UTC.""" + + def __init__(self, offset, name): + self.__offset = timedelta(minutes=offset) + self.__name = name + + def utcoffset(self, dt): + return self.__offset + + def tzname(self, dt): + return self.__name + + def dst(self, dt): + return ZERO + + +import time as _time + +STDOFFSET = timedelta(seconds=-_time.timezone) +if _time.daylight: + DSTOFFSET = timedelta(seconds=-_time.altzone) +else: + DSTOFFSET = STDOFFSET + +DSTDIFF = DSTOFFSET - STDOFFSET + + +# A class capturing the platform's idea of local time. +class LocalTimezone(tzinfo): + + def utcoffset(self, dt): + if self._isdst(dt): + return DSTOFFSET + else: + return STDOFFSET + + def dst(self, dt): + if self._isdst(dt): + return DSTDIFF + else: + return ZERO + + def tzname(self, dt): + return _time.tzname[self._isdst(dt)] + + def _isdst(self, dt): + tt = (dt.year, dt.month, dt.day, + dt.hour, dt.minute, dt.second, + dt.weekday(), 0, -1) + stamp = _time.mktime(tt) + tt = _time.localtime(stamp) + return tt.tm_isdst > 0 + +Local = LocalTimezone() + + +def first_sunday_on_or_after(dt): + days_to_go = 6 - dt.weekday() + if days_to_go: + dt += timedelta(days_to_go) + return dt + + +# In the US, DST starts at 2am (standard time) on the first Sunday in April. +DSTSTART = datetime(1, 4, 1, 2) +# and ends at 2am (DST time; 1am standard time) on the last Sunday of Oct. +# which is the first Sunday on or after Oct 25. +DSTEND = datetime(1, 10, 25, 1) + + +# A complete implementation of current DST rules for major US time zones. +class USTimeZone(tzinfo): + + def __init__(self, hours, reprname, stdname, dstname): + self.stdoffset = timedelta(hours=hours) + self.reprname = reprname + self.stdname = stdname + self.dstname = dstname + + def __repr__(self): + return self.reprname + + def tzname(self, dt): + if self.dst(dt): + return self.dstname + else: + return self.stdname + + def utcoffset(self, dt): + return self.stdoffset + self.dst(dt) + + def dst(self, dt): + if dt is None or dt.tzinfo is None: + # An exception may be sensible here, in one or both cases. + # It depends on how you want to treat them. The default + # fromutc() implementation (called by the default astimezone() + # implementation) passes a datetime with dt.tzinfo is self. + return ZERO + assert dt.tzinfo is self + + # Find first Sunday in April & the last in October. + start = first_sunday_on_or_after(DSTSTART.replace(year=dt.year)) + end = first_sunday_on_or_after(DSTEND.replace(year=dt.year)) + + # Can't compare naive to aware objects, so strip the timezone from + # dt first. + if start <= dt.replace(tzinfo=None) < end: + return HOUR + else: + return ZERO + +Eastern = USTimeZone(-5, "Eastern", "EST", "EDT") +Central = USTimeZone(-6, "Central", "CST", "CDT") +Mountain = USTimeZone(-7, "Mountain", "MST", "MDT") +Pacific = USTimeZone(-8, "Pacific", "PST", "PDT") diff --git a/lib/python3.12/site-packages/pytz/tzfile.py b/lib/python3.12/site-packages/pytz/tzfile.py new file mode 100644 index 0000000000000000000000000000000000000000..99e74489b859e21fcaa68e93089035c3d81a73c8 --- /dev/null +++ b/lib/python3.12/site-packages/pytz/tzfile.py @@ -0,0 +1,133 @@ +''' +$Id: tzfile.py,v 1.8 2004/06/03 00:15:24 zenzen Exp $ +''' + +from datetime import datetime +from struct import unpack, calcsize + +from pytz.tzinfo import StaticTzInfo, DstTzInfo, memorized_ttinfo +from pytz.tzinfo import memorized_datetime, memorized_timedelta + + +def _byte_string(s): + """Cast a string or byte string to an ASCII byte string.""" + return s.encode('ASCII') + +_NULL = _byte_string('\0') + + +def _std_string(s): + """Cast a string or byte string to an ASCII string.""" + return str(s.decode('ASCII')) + + +def build_tzinfo(zone, fp): + head_fmt = '>4s c 15x 6l' + head_size = calcsize(head_fmt) + (magic, format, ttisgmtcnt, ttisstdcnt, leapcnt, timecnt, + typecnt, charcnt) = unpack(head_fmt, fp.read(head_size)) + + # Make sure it is a tzfile(5) file + assert magic == _byte_string('TZif'), 'Got magic %s' % repr(magic) + + # Read out the transition times, localtime indices and ttinfo structures. + data_fmt = '>%(timecnt)dl %(timecnt)dB %(ttinfo)s %(charcnt)ds' % dict( + timecnt=timecnt, ttinfo='lBB' * typecnt, charcnt=charcnt) + data_size = calcsize(data_fmt) + data = unpack(data_fmt, fp.read(data_size)) + + # make sure we unpacked the right number of values + assert len(data) == 2 * timecnt + 3 * typecnt + 1 + transitions = [memorized_datetime(trans) + for trans in data[:timecnt]] + lindexes = list(data[timecnt:2 * timecnt]) + ttinfo_raw = data[2 * timecnt:-1] + tznames_raw = data[-1] + del data + + # Process ttinfo into separate structs + ttinfo = [] + tznames = {} + i = 0 + while i < len(ttinfo_raw): + # have we looked up this timezone name yet? + tzname_offset = ttinfo_raw[i + 2] + if tzname_offset not in tznames: + nul = tznames_raw.find(_NULL, tzname_offset) + if nul < 0: + nul = len(tznames_raw) + tznames[tzname_offset] = _std_string( + tznames_raw[tzname_offset:nul]) + ttinfo.append((ttinfo_raw[i], + bool(ttinfo_raw[i + 1]), + tznames[tzname_offset])) + i += 3 + + # Now build the timezone object + if len(ttinfo) == 1 or len(transitions) == 0: + ttinfo[0][0], ttinfo[0][2] + cls = type(zone, (StaticTzInfo,), dict( + zone=zone, + _utcoffset=memorized_timedelta(ttinfo[0][0]), + _tzname=ttinfo[0][2])) + else: + # Early dates use the first standard time ttinfo + i = 0 + while ttinfo[i][1]: + i += 1 + if ttinfo[i] == ttinfo[lindexes[0]]: + transitions[0] = datetime.min + else: + transitions.insert(0, datetime.min) + lindexes.insert(0, i) + + # calculate transition info + transition_info = [] + for i in range(len(transitions)): + inf = ttinfo[lindexes[i]] + utcoffset = inf[0] + if not inf[1]: + dst = 0 + else: + for j in range(i - 1, -1, -1): + prev_inf = ttinfo[lindexes[j]] + if not prev_inf[1]: + break + dst = inf[0] - prev_inf[0] # dst offset + + # Bad dst? Look further. DST > 24 hours happens when + # a timzone has moved across the international dateline. + if dst <= 0 or dst > 3600 * 3: + for j in range(i + 1, len(transitions)): + stdinf = ttinfo[lindexes[j]] + if not stdinf[1]: + dst = inf[0] - stdinf[0] + if dst > 0: + break # Found a useful std time. + + tzname = inf[2] + + # Round utcoffset and dst to the nearest minute or the + # datetime library will complain. Conversions to these timezones + # might be up to plus or minus 30 seconds out, but it is + # the best we can do. + utcoffset = int((utcoffset + 30) // 60) * 60 + dst = int((dst + 30) // 60) * 60 + transition_info.append(memorized_ttinfo(utcoffset, dst, tzname)) + + cls = type(zone, (DstTzInfo,), dict( + zone=zone, + _utc_transition_times=transitions, + _transition_info=transition_info)) + + return cls() + +if __name__ == '__main__': + import os.path + from pprint import pprint + base = os.path.join(os.path.dirname(__file__), 'zoneinfo') + tz = build_tzinfo('Australia/Melbourne', + open(os.path.join(base, 'Australia', 'Melbourne'), 'rb')) + tz = build_tzinfo('US/Eastern', + open(os.path.join(base, 'US', 'Eastern'), 'rb')) + pprint(tz._utc_transition_times) diff --git a/lib/python3.12/site-packages/pytz/tzinfo.py b/lib/python3.12/site-packages/pytz/tzinfo.py new file mode 100644 index 0000000000000000000000000000000000000000..49b5c3febdbce74624c0d2a7aea5b0eb839212cc --- /dev/null +++ b/lib/python3.12/site-packages/pytz/tzinfo.py @@ -0,0 +1,580 @@ +'''Base classes and helpers for building zone specific tzinfo classes''' + +from datetime import datetime, timedelta, tzinfo +from bisect import bisect_right +try: + set +except NameError: + from sets import Set as set + +import pytz +from pytz.exceptions import AmbiguousTimeError, NonExistentTimeError + +__all__ = [] + +_timedelta_cache = {} + + +def memorized_timedelta(seconds): + '''Create only one instance of each distinct timedelta''' + try: + return _timedelta_cache[seconds] + except KeyError: + delta = timedelta(seconds=seconds) + _timedelta_cache[seconds] = delta + return delta + + +_epoch = datetime(1970, 1, 1, 0, 0) # datetime.utcfromtimestamp(0) +_datetime_cache = {0: _epoch} + + +def memorized_datetime(seconds): + '''Create only one instance of each distinct datetime''' + try: + return _datetime_cache[seconds] + except KeyError: + # NB. We can't just do datetime.fromtimestamp(seconds, tz=timezone.utc).replace(tzinfo=None) + # as this fails with negative values under Windows (Bug #90096) + dt = _epoch + timedelta(seconds=seconds) + _datetime_cache[seconds] = dt + return dt + + +_ttinfo_cache = {} + + +def memorized_ttinfo(*args): + '''Create only one instance of each distinct tuple''' + try: + return _ttinfo_cache[args] + except KeyError: + ttinfo = ( + memorized_timedelta(args[0]), + memorized_timedelta(args[1]), + args[2] + ) + _ttinfo_cache[args] = ttinfo + return ttinfo + + +_notime = memorized_timedelta(0) + + +def _to_seconds(td): + '''Convert a timedelta to seconds''' + return td.seconds + td.days * 24 * 60 * 60 + + +class BaseTzInfo(tzinfo): + # Overridden in subclass + _utcoffset = None + _tzname = None + zone = None + + def __str__(self): + return self.zone + + +class StaticTzInfo(BaseTzInfo): + '''A timezone that has a constant offset from UTC + + These timezones are rare, as most locations have changed their + offset at some point in their history + ''' + def fromutc(self, dt): + '''See datetime.tzinfo.fromutc''' + if dt.tzinfo is not None and dt.tzinfo is not self: + raise ValueError('fromutc: dt.tzinfo is not self') + return (dt + self._utcoffset).replace(tzinfo=self) + + def utcoffset(self, dt, is_dst=None): + '''See datetime.tzinfo.utcoffset + + is_dst is ignored for StaticTzInfo, and exists only to + retain compatibility with DstTzInfo. + ''' + return self._utcoffset + + def dst(self, dt, is_dst=None): + '''See datetime.tzinfo.dst + + is_dst is ignored for StaticTzInfo, and exists only to + retain compatibility with DstTzInfo. + ''' + return _notime + + def tzname(self, dt, is_dst=None): + '''See datetime.tzinfo.tzname + + is_dst is ignored for StaticTzInfo, and exists only to + retain compatibility with DstTzInfo. + ''' + return self._tzname + + def localize(self, dt, is_dst=False): + '''Convert naive time to local time''' + if dt.tzinfo is not None: + raise ValueError('Not naive datetime (tzinfo is already set)') + return dt.replace(tzinfo=self) + + def normalize(self, dt, is_dst=False): + '''Correct the timezone information on the given datetime. + + This is normally a no-op, as StaticTzInfo timezones never have + ambiguous cases to correct: + + >>> from pytz import timezone + >>> gmt = timezone('GMT') + >>> isinstance(gmt, StaticTzInfo) + True + >>> dt = datetime(2011, 5, 8, 1, 2, 3, tzinfo=gmt) + >>> gmt.normalize(dt) is dt + True + + The supported method of converting between timezones is to use + datetime.astimezone(). Currently normalize() also works: + + >>> la = timezone('America/Los_Angeles') + >>> dt = la.localize(datetime(2011, 5, 7, 1, 2, 3)) + >>> fmt = '%Y-%m-%d %H:%M:%S %Z (%z)' + >>> gmt.normalize(dt).strftime(fmt) + '2011-05-07 08:02:03 GMT (+0000)' + ''' + if dt.tzinfo is self: + return dt + if dt.tzinfo is None: + raise ValueError('Naive time - no tzinfo set') + return dt.astimezone(self) + + def __repr__(self): + return '' % (self.zone,) + + def __reduce__(self): + # Special pickle to zone remains a singleton and to cope with + # database changes. + return pytz._p, (self.zone,) + + +class DstTzInfo(BaseTzInfo): + '''A timezone that has a variable offset from UTC + + The offset might change if daylight saving time comes into effect, + or at a point in history when the region decides to change their + timezone definition. + ''' + # Overridden in subclass + + # Sorted list of DST transition times, UTC + _utc_transition_times = None + + # [(utcoffset, dstoffset, tzname)] corresponding to + # _utc_transition_times entries + _transition_info = None + + zone = None + + # Set in __init__ + + _tzinfos = None + _dst = None # DST offset + + def __init__(self, _inf=None, _tzinfos=None): + if _inf: + self._tzinfos = _tzinfos + self._utcoffset, self._dst, self._tzname = _inf + else: + _tzinfos = {} + self._tzinfos = _tzinfos + self._utcoffset, self._dst, self._tzname = ( + self._transition_info[0]) + _tzinfos[self._transition_info[0]] = self + for inf in self._transition_info[1:]: + if inf not in _tzinfos: + _tzinfos[inf] = self.__class__(inf, _tzinfos) + + def fromutc(self, dt): + '''See datetime.tzinfo.fromutc''' + if (dt.tzinfo is not None and + getattr(dt.tzinfo, '_tzinfos', None) is not self._tzinfos): + raise ValueError('fromutc: dt.tzinfo is not self') + dt = dt.replace(tzinfo=None) + idx = max(0, bisect_right(self._utc_transition_times, dt) - 1) + inf = self._transition_info[idx] + return (dt + inf[0]).replace(tzinfo=self._tzinfos[inf]) + + def normalize(self, dt): + '''Correct the timezone information on the given datetime + + If date arithmetic crosses DST boundaries, the tzinfo + is not magically adjusted. This method normalizes the + tzinfo to the correct one. + + To test, first we need to do some setup + + >>> from pytz import timezone + >>> utc = timezone('UTC') + >>> eastern = timezone('US/Eastern') + >>> fmt = '%Y-%m-%d %H:%M:%S %Z (%z)' + + We next create a datetime right on an end-of-DST transition point, + the instant when the wallclocks are wound back one hour. + + >>> 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)' + + Now, if we subtract a few minutes from it, note that the timezone + information has not changed. + + >>> before = loc_dt - timedelta(minutes=10) + >>> before.strftime(fmt) + '2002-10-27 00:50:00 EST (-0500)' + + But we can fix that by calling the normalize method + + >>> before = eastern.normalize(before) + >>> before.strftime(fmt) + '2002-10-27 01:50:00 EDT (-0400)' + + The supported method of converting between timezones is to use + datetime.astimezone(). Currently, normalize() also works: + + >>> th = timezone('Asia/Bangkok') + >>> am = timezone('Europe/Amsterdam') + >>> dt = th.localize(datetime(2011, 5, 7, 1, 2, 3)) + >>> fmt = '%Y-%m-%d %H:%M:%S %Z (%z)' + >>> am.normalize(dt).strftime(fmt) + '2011-05-06 20:02:03 CEST (+0200)' + ''' + if dt.tzinfo is None: + raise ValueError('Naive time - no tzinfo set') + + # Convert dt in localtime to UTC + offset = dt.tzinfo._utcoffset + dt = dt.replace(tzinfo=None) + dt = dt - offset + # convert it back, and return it + return self.fromutc(dt) + + def localize(self, dt, is_dst=False): + '''Convert naive time to local time. + + This method should be used to construct localtimes, rather + than passing a tzinfo argument to a datetime constructor. + + is_dst is used to determine the correct timezone in the ambigous + period at the end of daylight saving time. + + >>> from pytz import timezone + >>> fmt = '%Y-%m-%d %H:%M:%S %Z (%z)' + >>> amdam = timezone('Europe/Amsterdam') + >>> dt = datetime(2004, 10, 31, 2, 0, 0) + >>> loc_dt1 = amdam.localize(dt, is_dst=True) + >>> loc_dt2 = amdam.localize(dt, is_dst=False) + >>> loc_dt1.strftime(fmt) + '2004-10-31 02:00:00 CEST (+0200)' + >>> loc_dt2.strftime(fmt) + '2004-10-31 02:00:00 CET (+0100)' + >>> str(loc_dt2 - loc_dt1) + '1:00:00' + + Use is_dst=None to raise an AmbiguousTimeError for ambiguous + times at the end of daylight saving time + + >>> try: + ... loc_dt1 = amdam.localize(dt, is_dst=None) + ... except AmbiguousTimeError: + ... print('Ambiguous') + Ambiguous + + is_dst defaults to False + + >>> amdam.localize(dt) == amdam.localize(dt, False) + True + + is_dst is also used to determine the correct timezone in the + wallclock times jumped over at the start of daylight saving time. + + >>> pacific = timezone('US/Pacific') + >>> dt = datetime(2008, 3, 9, 2, 0, 0) + >>> ploc_dt1 = pacific.localize(dt, is_dst=True) + >>> ploc_dt2 = pacific.localize(dt, is_dst=False) + >>> ploc_dt1.strftime(fmt) + '2008-03-09 02:00:00 PDT (-0700)' + >>> ploc_dt2.strftime(fmt) + '2008-03-09 02:00:00 PST (-0800)' + >>> str(ploc_dt2 - ploc_dt1) + '1:00:00' + + Use is_dst=None to raise a NonExistentTimeError for these skipped + times. + + >>> try: + ... loc_dt1 = pacific.localize(dt, is_dst=None) + ... except NonExistentTimeError: + ... print('Non-existent') + Non-existent + ''' + if dt.tzinfo is not None: + raise ValueError('Not naive datetime (tzinfo is already set)') + + # Find the two best possibilities. + possible_loc_dt = set() + for delta in [timedelta(days=-1), timedelta(days=1)]: + loc_dt = dt + delta + idx = max(0, bisect_right( + self._utc_transition_times, loc_dt) - 1) + inf = self._transition_info[idx] + tzinfo = self._tzinfos[inf] + loc_dt = tzinfo.normalize(dt.replace(tzinfo=tzinfo)) + if loc_dt.replace(tzinfo=None) == dt: + possible_loc_dt.add(loc_dt) + + if len(possible_loc_dt) == 1: + return possible_loc_dt.pop() + + # If there are no possibly correct timezones, we are attempting + # to convert a time that never happened - the time period jumped + # during the start-of-DST transition period. + if len(possible_loc_dt) == 0: + # If we refuse to guess, raise an exception. + if is_dst is None: + raise NonExistentTimeError(dt) + + # If we are forcing the pre-DST side of the DST transition, we + # obtain the correct timezone by winding the clock forward a few + # hours. + elif is_dst: + return self.localize( + dt + timedelta(hours=6), is_dst=True) - timedelta(hours=6) + + # If we are forcing the post-DST side of the DST transition, we + # obtain the correct timezone by winding the clock back. + else: + return self.localize( + dt - timedelta(hours=6), + is_dst=False) + timedelta(hours=6) + + # If we get this far, we have multiple possible timezones - this + # is an ambiguous case occurring during the end-of-DST transition. + + # If told to be strict, raise an exception since we have an + # ambiguous case + if is_dst is None: + raise AmbiguousTimeError(dt) + + # Filter out the possiblilities that don't match the requested + # is_dst + filtered_possible_loc_dt = [ + p for p in possible_loc_dt if bool(p.tzinfo._dst) == is_dst + ] + + # Hopefully we only have one possibility left. Return it. + if len(filtered_possible_loc_dt) == 1: + return filtered_possible_loc_dt[0] + + if len(filtered_possible_loc_dt) == 0: + filtered_possible_loc_dt = list(possible_loc_dt) + + # If we get this far, we have in a wierd timezone transition + # where the clocks have been wound back but is_dst is the same + # in both (eg. Europe/Warsaw 1915 when they switched to CET). + # At this point, we just have to guess unless we allow more + # hints to be passed in (such as the UTC offset or abbreviation), + # but that is just getting silly. + # + # Choose the earliest (by UTC) applicable timezone if is_dst=True + # Choose the latest (by UTC) applicable timezone if is_dst=False + # i.e., behave like end-of-DST transition + dates = {} # utc -> local + for local_dt in filtered_possible_loc_dt: + utc_time = ( + local_dt.replace(tzinfo=None) - local_dt.tzinfo._utcoffset) + assert utc_time not in dates + dates[utc_time] = local_dt + return dates[[min, max][not is_dst](dates)] + + def utcoffset(self, dt, is_dst=None): + '''See datetime.tzinfo.utcoffset + + The is_dst parameter may be used to remove ambiguity during DST + transitions. + + >>> from pytz import timezone + >>> tz = timezone('America/St_Johns') + >>> ambiguous = datetime(2009, 10, 31, 23, 30) + + >>> str(tz.utcoffset(ambiguous, is_dst=False)) + '-1 day, 20:30:00' + + >>> str(tz.utcoffset(ambiguous, is_dst=True)) + '-1 day, 21:30:00' + + >>> try: + ... tz.utcoffset(ambiguous) + ... except AmbiguousTimeError: + ... print('Ambiguous') + Ambiguous + + ''' + if dt is None: + return None + elif dt.tzinfo is not self: + dt = self.localize(dt, is_dst) + return dt.tzinfo._utcoffset + else: + return self._utcoffset + + def dst(self, dt, is_dst=None): + '''See datetime.tzinfo.dst + + The is_dst parameter may be used to remove ambiguity during DST + transitions. + + >>> from pytz import timezone + >>> tz = timezone('America/St_Johns') + + >>> normal = datetime(2009, 9, 1) + + >>> str(tz.dst(normal)) + '1:00:00' + >>> str(tz.dst(normal, is_dst=False)) + '1:00:00' + >>> str(tz.dst(normal, is_dst=True)) + '1:00:00' + + >>> ambiguous = datetime(2009, 10, 31, 23, 30) + + >>> str(tz.dst(ambiguous, is_dst=False)) + '0:00:00' + >>> str(tz.dst(ambiguous, is_dst=True)) + '1:00:00' + >>> try: + ... tz.dst(ambiguous) + ... except AmbiguousTimeError: + ... print('Ambiguous') + Ambiguous + + ''' + if dt is None: + return None + elif dt.tzinfo is not self: + dt = self.localize(dt, is_dst) + return dt.tzinfo._dst + else: + return self._dst + + def tzname(self, dt, is_dst=None): + '''See datetime.tzinfo.tzname + + The is_dst parameter may be used to remove ambiguity during DST + transitions. + + >>> from pytz import timezone + >>> tz = timezone('America/St_Johns') + + >>> normal = datetime(2009, 9, 1) + + >>> tz.tzname(normal) + 'NDT' + >>> tz.tzname(normal, is_dst=False) + 'NDT' + >>> tz.tzname(normal, is_dst=True) + 'NDT' + + >>> ambiguous = datetime(2009, 10, 31, 23, 30) + + >>> tz.tzname(ambiguous, is_dst=False) + 'NST' + >>> tz.tzname(ambiguous, is_dst=True) + 'NDT' + >>> try: + ... tz.tzname(ambiguous) + ... except AmbiguousTimeError: + ... print('Ambiguous') + Ambiguous + ''' + if dt is None: + return self.zone + elif dt.tzinfo is not self: + dt = self.localize(dt, is_dst) + return dt.tzinfo._tzname + else: + return self._tzname + + def __repr__(self): + if self._dst: + dst = 'DST' + else: + dst = 'STD' + if self._utcoffset > _notime: + return '' % ( + self.zone, self._tzname, self._utcoffset, dst + ) + else: + return '' % ( + self.zone, self._tzname, self._utcoffset, dst + ) + + def __reduce__(self): + # Special pickle to zone remains a singleton and to cope with + # database changes. + return pytz._p, ( + self.zone, + _to_seconds(self._utcoffset), + _to_seconds(self._dst), + self._tzname + ) + + +def unpickler(zone, utcoffset=None, dstoffset=None, tzname=None): + """Factory function for unpickling pytz tzinfo instances. + + This is shared for both StaticTzInfo and DstTzInfo instances, because + database changes could cause a zones implementation to switch between + these two base classes and we can't break pickles on a pytz version + upgrade. + """ + # Raises a KeyError if zone no longer exists, which should never happen + # and would be a bug. + tz = pytz.timezone(zone) + + # A StaticTzInfo - just return it + if utcoffset is None: + return tz + + # This pickle was created from a DstTzInfo. We need to + # determine which of the list of tzinfo instances for this zone + # to use in order to restore the state of any datetime instances using + # it correctly. + utcoffset = memorized_timedelta(utcoffset) + dstoffset = memorized_timedelta(dstoffset) + try: + return tz._tzinfos[(utcoffset, dstoffset, tzname)] + except KeyError: + # The particular state requested in this timezone no longer exists. + # This indicates a corrupt pickle, or the timezone database has been + # corrected violently enough to make this particular + # (utcoffset,dstoffset) no longer exist in the zone, or the + # abbreviation has been changed. + pass + + # See if we can find an entry differing only by tzname. Abbreviations + # get changed from the initial guess by the database maintainers to + # match reality when this information is discovered. + for localized_tz in tz._tzinfos.values(): + if (localized_tz._utcoffset == utcoffset and + localized_tz._dst == dstoffset): + return localized_tz + + # This (utcoffset, dstoffset) information has been removed from the + # zone. Add it back. This might occur when the database maintainers have + # corrected incorrect information. datetime instances using this + # incorrect information will continue to do so, exactly as they were + # before being pickled. This is purely an overly paranoid safety net - I + # doubt this will ever been needed in real life. + inf = (utcoffset, dstoffset, tzname) + tz._tzinfos[inf] = tz.__class__(inf, tz._tzinfos) + return tz._tzinfos[inf] diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aden b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aden new file mode 100644 index 0000000000000000000000000000000000000000..8c8062471dce91a5be827d6908795ee7391a4afc Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aden differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Almaty b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Almaty new file mode 100644 index 0000000000000000000000000000000000000000..855abbd6e3ee338c3d46bee7b592f3d645d2c273 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Almaty differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Amman b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Amman new file mode 100644 index 0000000000000000000000000000000000000000..0a8e350a334bbee29b155fb9944c6831460d973e Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Amman differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Anadyr b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Anadyr new file mode 100644 index 0000000000000000000000000000000000000000..35c531c0709fabace97e52cacef74b4cfbbcd820 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Anadyr differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aqtau b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aqtau new file mode 100644 index 0000000000000000000000000000000000000000..0e1c16d32eca7652aa7e1392e606cf5b54e26bc5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aqtau differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aqtobe b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aqtobe new file mode 100644 index 0000000000000000000000000000000000000000..3b5d6eb41883a7b32819f387b4c03c93e9d489cf Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Aqtobe differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ashgabat b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ashgabat new file mode 100644 index 0000000000000000000000000000000000000000..2bd1cb3da0f5a11024f1609b09ae68645e6ae75c Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ashgabat differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ashkhabad b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ashkhabad new file mode 100644 index 0000000000000000000000000000000000000000..2bd1cb3da0f5a11024f1609b09ae68645e6ae75c Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ashkhabad differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Atyrau b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Atyrau new file mode 100644 index 0000000000000000000000000000000000000000..e7ea9c545a56ae0b0fabd02e47748004a6d87c3b Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Atyrau differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Baghdad b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Baghdad new file mode 100644 index 0000000000000000000000000000000000000000..c0e607234a07cd2650aabc1fd022bd5d4f99b175 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Baghdad differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bahrain b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bahrain new file mode 100644 index 0000000000000000000000000000000000000000..098997e7dd972ca4acedc535c4a645ee966baa33 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bahrain differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Baku b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Baku new file mode 100644 index 0000000000000000000000000000000000000000..ae0ce4e7c3d273b537fa3653784e73be68ebb441 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Baku differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bangkok b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bangkok new file mode 100644 index 0000000000000000000000000000000000000000..fa799db39e7625dd74bd9caa5c29b4819a7cbd3f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bangkok differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Barnaul b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Barnaul new file mode 100644 index 0000000000000000000000000000000000000000..2f6b8101d474642410af84260fb2b88c493945b1 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Barnaul differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Beirut b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Beirut new file mode 100644 index 0000000000000000000000000000000000000000..fb266ede2279b6aff913538d9d5aae3935e53aeb Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Beirut differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bishkek b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bishkek new file mode 100644 index 0000000000000000000000000000000000000000..547fd5e1bd16152073f7237ad265b5c643e7cbfe Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Bishkek differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Brunei b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Brunei new file mode 100644 index 0000000000000000000000000000000000000000..098c6a0b0afc2c3ef3a6937104d270074ff2a9b6 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Brunei differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Calcutta b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Calcutta new file mode 100644 index 0000000000000000000000000000000000000000..0014046d29a38e9b8006f746fea794d7f71eb479 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Calcutta differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chita b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chita new file mode 100644 index 0000000000000000000000000000000000000000..75b3d7b3a6b3bf2b7459532fce1142912b897efd Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chita differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Choibalsan b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Choibalsan new file mode 100644 index 0000000000000000000000000000000000000000..2aa5cc4b84d369b03684c30cdb66b7abf087f467 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Choibalsan differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chongqing b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chongqing new file mode 100644 index 0000000000000000000000000000000000000000..91f6f8bc2e234bafd484146986bdb289082c3588 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chongqing differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chungking b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chungking new file mode 100644 index 0000000000000000000000000000000000000000..91f6f8bc2e234bafd484146986bdb289082c3588 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Chungking differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Colombo b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Colombo new file mode 100644 index 0000000000000000000000000000000000000000..353fe2aa3564c79809e05e12711f90669e5e2388 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Colombo differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dacca b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dacca new file mode 100644 index 0000000000000000000000000000000000000000..3cf597d83f90403132573d78247e6d02e66d0583 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dacca differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Damascus b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Damascus new file mode 100644 index 0000000000000000000000000000000000000000..afd956c878856e0e239887963b32ddb8327ccf38 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Damascus differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dhaka b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dhaka new file mode 100644 index 0000000000000000000000000000000000000000..3cf597d83f90403132573d78247e6d02e66d0583 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dhaka differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dili b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dili new file mode 100644 index 0000000000000000000000000000000000000000..4614e4fc4353a7094ba5173a0bd09158ef49d9a9 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dili differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dubai b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dubai new file mode 100644 index 0000000000000000000000000000000000000000..b3ac791aef4e73d6d644c40c614f37f15d462cdd Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dubai differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dushanbe b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dushanbe new file mode 100644 index 0000000000000000000000000000000000000000..89e875beaef780b2ceae08ee4703cd813f33d6a4 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Dushanbe differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Famagusta b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Famagusta new file mode 100644 index 0000000000000000000000000000000000000000..653b146a60e5e5641a07bfc82f9590dad1ed69f5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Famagusta differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Gaza b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Gaza new file mode 100644 index 0000000000000000000000000000000000000000..dd5781e8ae26ed66aa2650c580c561fae50df4b5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Gaza differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Harbin b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Harbin new file mode 100644 index 0000000000000000000000000000000000000000..91f6f8bc2e234bafd484146986bdb289082c3588 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Harbin differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hebron b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hebron new file mode 100644 index 0000000000000000000000000000000000000000..a64fc9e7b2599e8f60d228ef3cf1743b5296cd4f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hebron differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ho_Chi_Minh b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ho_Chi_Minh new file mode 100644 index 0000000000000000000000000000000000000000..9c45ed991a230919d129c63c660aa8d1ad74d0ea Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ho_Chi_Minh differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hong_Kong b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hong_Kong new file mode 100644 index 0000000000000000000000000000000000000000..f9f7b134dd5c35e4718c6fa3697024cb95442c53 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hong_Kong differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hovd b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hovd new file mode 100644 index 0000000000000000000000000000000000000000..8b9abca344daf199b554888fc4fff562a7a18e4a Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Hovd differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Irkutsk b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Irkutsk new file mode 100644 index 0000000000000000000000000000000000000000..e74a4d3f6b3a799b0fd12272c3aae0f084fb2783 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Irkutsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Istanbul b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Istanbul new file mode 100644 index 0000000000000000000000000000000000000000..7c2336dd80c3c9cbf71cb53d2b2c1f89a65a8ba5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Istanbul differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jakarta b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jakarta new file mode 100644 index 0000000000000000000000000000000000000000..ec4bd5747a8c9c528dfd22c8e3171851784ade59 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jakarta differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jayapura b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jayapura new file mode 100644 index 0000000000000000000000000000000000000000..3002c82022f7e9a4ff1ad545ea617106a13b24e0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jayapura differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jerusalem b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jerusalem new file mode 100644 index 0000000000000000000000000000000000000000..1ebd0664aa29c0abd722661f761031ec0304631c Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Jerusalem differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kabul b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kabul new file mode 100644 index 0000000000000000000000000000000000000000..661efc83294edbe5f615daa87a1f4a5590aae167 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kabul differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kamchatka b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kamchatka new file mode 100644 index 0000000000000000000000000000000000000000..99776f515fd53a86a6886311d91ff33b401396c3 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kamchatka differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Karachi b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Karachi new file mode 100644 index 0000000000000000000000000000000000000000..ba65c0e8d31cb1d59935878af29baac722f45936 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Karachi differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kashgar b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kashgar new file mode 100644 index 0000000000000000000000000000000000000000..62bdcac14db3f464ff561e32db2c6b55c0cb1866 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kashgar differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kathmandu b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kathmandu new file mode 100644 index 0000000000000000000000000000000000000000..751cf4a8939e898f8abe4358a98b2a5f59f65e21 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kathmandu differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Katmandu b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Katmandu new file mode 100644 index 0000000000000000000000000000000000000000..751cf4a8939e898f8abe4358a98b2a5f59f65e21 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Katmandu differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Khandyga b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Khandyga new file mode 100644 index 0000000000000000000000000000000000000000..7cdc99a9803f4e1ce3c5f7f9e6cd761f178536ae Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Khandyga differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kolkata b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kolkata new file mode 100644 index 0000000000000000000000000000000000000000..0014046d29a38e9b8006f746fea794d7f71eb479 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kolkata differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Krasnoyarsk b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Krasnoyarsk new file mode 100644 index 0000000000000000000000000000000000000000..4c27b2decd1159cf47285f0b18b5ed5178641e82 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Krasnoyarsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuala_Lumpur b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuala_Lumpur new file mode 100644 index 0000000000000000000000000000000000000000..3d9f191e3acad56fb4f73ab44a74c50cdf91c12e Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuala_Lumpur differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuching b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuching new file mode 100644 index 0000000000000000000000000000000000000000..098c6a0b0afc2c3ef3a6937104d270074ff2a9b6 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuching differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuwait b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuwait new file mode 100644 index 0000000000000000000000000000000000000000..8c8062471dce91a5be827d6908795ee7391a4afc Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Kuwait differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Macao b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Macao new file mode 100644 index 0000000000000000000000000000000000000000..cac65063d0dbf48e37c547fba3b67f34110d5a90 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Macao differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Macau b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Macau new file mode 100644 index 0000000000000000000000000000000000000000..cac65063d0dbf48e37c547fba3b67f34110d5a90 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Macau differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Magadan b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Magadan new file mode 100644 index 0000000000000000000000000000000000000000..70c198baf743457c58a8fe4d1f28186d8541a505 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Magadan differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Makassar b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Makassar new file mode 100644 index 0000000000000000000000000000000000000000..556ba866933d37f3cfcf8042045d64e209bae30f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Makassar differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Manila b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Manila new file mode 100644 index 0000000000000000000000000000000000000000..168f72096d96e9f8a6dc4c6073086c19dbc20d0d Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Manila differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Muscat b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Muscat new file mode 100644 index 0000000000000000000000000000000000000000..b3ac791aef4e73d6d644c40c614f37f15d462cdd Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Muscat differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Nicosia b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Nicosia new file mode 100644 index 0000000000000000000000000000000000000000..f7f10ab7665e94ca44fd8cd98a362cd4b304eff1 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Nicosia differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Novokuznetsk b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Novokuznetsk new file mode 100644 index 0000000000000000000000000000000000000000..a5e1b79600d21019a3b93f3e2d796129d9b1d655 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Novokuznetsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Novosibirsk b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Novosibirsk new file mode 100644 index 0000000000000000000000000000000000000000..4ac7582ad5ee895bf9b354436142fd7ece1996ee Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Novosibirsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Omsk b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Omsk new file mode 100644 index 0000000000000000000000000000000000000000..16c5f3cfed75151d50ffc2dc483ac6279816a01f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Omsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Oral b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Oral new file mode 100644 index 0000000000000000000000000000000000000000..3b9ecacf6eef72366df1cc847c0d4666efc8a48c Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Oral differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Phnom_Penh b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Phnom_Penh new file mode 100644 index 0000000000000000000000000000000000000000..fa799db39e7625dd74bd9caa5c29b4819a7cbd3f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Phnom_Penh differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Pontianak b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Pontianak new file mode 100644 index 0000000000000000000000000000000000000000..12ce24cbeae404efe6921081d21289be452ff88d Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Pontianak differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Pyongyang b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Pyongyang new file mode 100644 index 0000000000000000000000000000000000000000..7ad7e0b2cf8fa4fc844fe8cc9c58e4f3018cead1 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Pyongyang differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qatar b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qatar new file mode 100644 index 0000000000000000000000000000000000000000..098997e7dd972ca4acedc535c4a645ee966baa33 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qatar differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qostanay b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qostanay new file mode 100644 index 0000000000000000000000000000000000000000..2ee9ef7e9859bc95da2450e3b282fd4275cbbb39 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qostanay differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qyzylorda b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qyzylorda new file mode 100644 index 0000000000000000000000000000000000000000..27b522a7d5e24eafdf29dd541ebbca69ce4db7b8 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Qyzylorda differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Rangoon b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Rangoon new file mode 100644 index 0000000000000000000000000000000000000000..eef37b42e8a0e7179f8113bea01f4a71d668e8ef Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Rangoon differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Riyadh b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Riyadh new file mode 100644 index 0000000000000000000000000000000000000000..8c8062471dce91a5be827d6908795ee7391a4afc Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Riyadh differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Saigon b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Saigon new file mode 100644 index 0000000000000000000000000000000000000000..9c45ed991a230919d129c63c660aa8d1ad74d0ea Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Saigon differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Sakhalin b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Sakhalin new file mode 100644 index 0000000000000000000000000000000000000000..beb77b449654decdfec1f6bfc16fafc8bc8ceeaa Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Sakhalin differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Samarkand b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Samarkand new file mode 100644 index 0000000000000000000000000000000000000000..8a93767bfef2cecbeee15db413a0a7c277e761a1 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Samarkand differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Seoul b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Seoul new file mode 100644 index 0000000000000000000000000000000000000000..96199e73e73aafacd89e48cb2855a96d7a134e1d Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Seoul differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Shanghai b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Shanghai new file mode 100644 index 0000000000000000000000000000000000000000..91f6f8bc2e234bafd484146986bdb289082c3588 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Shanghai differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Singapore b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Singapore new file mode 100644 index 0000000000000000000000000000000000000000..3d9f191e3acad56fb4f73ab44a74c50cdf91c12e Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Singapore differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Srednekolymsk b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Srednekolymsk new file mode 100644 index 0000000000000000000000000000000000000000..d21e7eeed5408affb3277bde903d6f5f2090ea61 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Srednekolymsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Taipei b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Taipei new file mode 100644 index 0000000000000000000000000000000000000000..24c43444b6751343d2915843d03e55753d4d7359 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Taipei differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tashkent b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tashkent new file mode 100644 index 0000000000000000000000000000000000000000..a9f6cd93c849c8b1677bc54f8bb836c91b40d63d Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tashkent differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tbilisi b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tbilisi new file mode 100644 index 0000000000000000000000000000000000000000..3b131bb10e100bc2dd619d57367a4e43c86e5316 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tbilisi differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tehran b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tehran new file mode 100644 index 0000000000000000000000000000000000000000..78f28cc402027bd337cf67b937bf8c31a75ee8b3 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tehran differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tel_Aviv b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tel_Aviv new file mode 100644 index 0000000000000000000000000000000000000000..1ebd0664aa29c0abd722661f761031ec0304631c Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tel_Aviv differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Thimbu b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Thimbu new file mode 100644 index 0000000000000000000000000000000000000000..95a9de9657f5b8b15aacb6e05c085c7e11bb8f55 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Thimbu differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Thimphu b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Thimphu new file mode 100644 index 0000000000000000000000000000000000000000..95a9de9657f5b8b15aacb6e05c085c7e11bb8f55 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Thimphu differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tokyo b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tokyo new file mode 100644 index 0000000000000000000000000000000000000000..26f4d34d67b46513491f26c2e661c6e653cc130d Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tokyo differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tomsk b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tomsk new file mode 100644 index 0000000000000000000000000000000000000000..a6e494a78cef4baaa30d06cea4186391788eb30b Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Tomsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ujung_Pandang b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ujung_Pandang new file mode 100644 index 0000000000000000000000000000000000000000..556ba866933d37f3cfcf8042045d64e209bae30f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ujung_Pandang differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ulaanbaatar b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ulaanbaatar new file mode 100644 index 0000000000000000000000000000000000000000..2aa5cc4b84d369b03684c30cdb66b7abf087f467 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ulaanbaatar differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ulan_Bator b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ulan_Bator new file mode 100644 index 0000000000000000000000000000000000000000..2aa5cc4b84d369b03684c30cdb66b7abf087f467 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ulan_Bator differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Urumqi b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Urumqi new file mode 100644 index 0000000000000000000000000000000000000000..62bdcac14db3f464ff561e32db2c6b55c0cb1866 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Urumqi differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ust-Nera b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ust-Nera new file mode 100644 index 0000000000000000000000000000000000000000..d05726aba9fd67bf230290b3e2d74b75e46ee214 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Ust-Nera differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Vientiane b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Vientiane new file mode 100644 index 0000000000000000000000000000000000000000..fa799db39e7625dd74bd9caa5c29b4819a7cbd3f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Vientiane differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Vladivostok b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Vladivostok new file mode 100644 index 0000000000000000000000000000000000000000..274a10b43d8089ed522daf9ded5b6aa3653075c9 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Vladivostok differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yakutsk b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yakutsk new file mode 100644 index 0000000000000000000000000000000000000000..ae65a5f9b94d196b29b68ddec7862e098d829b59 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yakutsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yangon b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yangon new file mode 100644 index 0000000000000000000000000000000000000000..eef37b42e8a0e7179f8113bea01f4a71d668e8ef Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yangon differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yekaterinburg b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yekaterinburg new file mode 100644 index 0000000000000000000000000000000000000000..d4d19ccf1e91121589a7cabff444bc916f67f9ce Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yekaterinburg differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yerevan b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yerevan new file mode 100644 index 0000000000000000000000000000000000000000..0d5f6853ac79c41fcb445a23062a27aca9e471f3 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Asia/Yerevan differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/CET b/lib/python3.12/site-packages/pytz/zoneinfo/CET new file mode 100644 index 0000000000000000000000000000000000000000..40d7124e5346af056c75e2f7012a51d94e8154b7 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/CET differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/CST6CDT b/lib/python3.12/site-packages/pytz/zoneinfo/CST6CDT new file mode 100644 index 0000000000000000000000000000000000000000..c6981a06b1d9c26f447518efe265a6454726eae7 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/CST6CDT differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Chile/Continental b/lib/python3.12/site-packages/pytz/zoneinfo/Chile/Continental new file mode 100644 index 0000000000000000000000000000000000000000..010c6bd04cae79078540da560ce38400bfe0ade6 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Chile/Continental differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Chile/EasterIsland b/lib/python3.12/site-packages/pytz/zoneinfo/Chile/EasterIsland new file mode 100644 index 0000000000000000000000000000000000000000..184cb6a83b3392d0492c42297531c85e7e38c4f5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Chile/EasterIsland differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Cuba b/lib/python3.12/site-packages/pytz/zoneinfo/Cuba new file mode 100644 index 0000000000000000000000000000000000000000..b69ac4510784f23ee794cd6d11e62315a7318e5e Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Cuba differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/EET b/lib/python3.12/site-packages/pytz/zoneinfo/EET new file mode 100644 index 0000000000000000000000000000000000000000..9f3a0678d766881389e129c93def7fffd74f14f1 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/EET differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/EST b/lib/python3.12/site-packages/pytz/zoneinfo/EST new file mode 100644 index 0000000000000000000000000000000000000000..9964b9a33452f4b636f43703b7cdec4891cbda5f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/EST differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/EST5EDT b/lib/python3.12/site-packages/pytz/zoneinfo/EST5EDT new file mode 100644 index 0000000000000000000000000000000000000000..a8b9ab1992257d721ad627b14f535c3d4b020888 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/EST5EDT differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Egypt b/lib/python3.12/site-packages/pytz/zoneinfo/Egypt new file mode 100644 index 0000000000000000000000000000000000000000..dd538c65db6ed0a0e47feb7b6001640516958e19 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Egypt differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Eire b/lib/python3.12/site-packages/pytz/zoneinfo/Eire new file mode 100644 index 0000000000000000000000000000000000000000..c729def42fc0822e6b24f7bf503c67cc7e6965f0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Eire differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Dublin b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Dublin new file mode 100644 index 0000000000000000000000000000000000000000..c729def42fc0822e6b24f7bf503c67cc7e6965f0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Dublin differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Helsinki b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Helsinki new file mode 100644 index 0000000000000000000000000000000000000000..b4f8f9cbb57450549933f83ac90dd56a2ca75344 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Helsinki differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Kaliningrad b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Kaliningrad new file mode 100644 index 0000000000000000000000000000000000000000..cc99beabe4ffc5107c4719d1201d6583b9ead03a Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Kaliningrad differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Kyiv b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Kyiv new file mode 100644 index 0000000000000000000000000000000000000000..52efea88065b220e44fd876de3bf3090fe62cc79 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Kyiv differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/London b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/London new file mode 100644 index 0000000000000000000000000000000000000000..5ad74220e83e6c30a0aeefda4a00271b8ebdfcad Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/London differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Luxembourg b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Luxembourg new file mode 100644 index 0000000000000000000000000000000000000000..40d7124e5346af056c75e2f7012a51d94e8154b7 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Luxembourg differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Madrid b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Madrid new file mode 100644 index 0000000000000000000000000000000000000000..53f4cd101c18058a1ee4d4b9d1184ada11d74232 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Madrid differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Mariehamn b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Mariehamn new file mode 100644 index 0000000000000000000000000000000000000000..b4f8f9cbb57450549933f83ac90dd56a2ca75344 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Mariehamn differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Minsk b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Minsk new file mode 100644 index 0000000000000000000000000000000000000000..3731e40d8634c4d58ddb3378d828c92948f06d67 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Minsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Moscow b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Moscow new file mode 100644 index 0000000000000000000000000000000000000000..ddb3f4e99a1030f33b56fad986c8d9c16e59eb32 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Moscow differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Nicosia b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Nicosia new file mode 100644 index 0000000000000000000000000000000000000000..f7f10ab7665e94ca44fd8cd98a362cd4b304eff1 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Nicosia differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Prague b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Prague new file mode 100644 index 0000000000000000000000000000000000000000..ce8f433ece44f0b96b18d3b5780730e7f9cad9f5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Prague differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/San_Marino b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/San_Marino new file mode 100644 index 0000000000000000000000000000000000000000..32b2899a306dde401fa2e3952d06f5f4d9952bed Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/San_Marino differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Ulyanovsk b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Ulyanovsk new file mode 100644 index 0000000000000000000000000000000000000000..d668233b37f268a745e995177195b06ffab1e69b Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Ulyanovsk differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vaduz b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vaduz new file mode 100644 index 0000000000000000000000000000000000000000..ad6cf59281a1046d9dcd045fda521585e3e33e06 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vaduz differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vienna b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vienna new file mode 100644 index 0000000000000000000000000000000000000000..3582bb15cd7322088839b0134987ad10e717b6b5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vienna differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vilnius b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vilnius new file mode 100644 index 0000000000000000000000000000000000000000..7abd63fa608e0186b9f154d9fcc32472c28f6759 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Vilnius differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Volgograd b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Volgograd new file mode 100644 index 0000000000000000000000000000000000000000..9d51a38c0ea719c5042837d47e0f427057b845f0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Europe/Volgograd differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Factory b/lib/python3.12/site-packages/pytz/zoneinfo/Factory new file mode 100644 index 0000000000000000000000000000000000000000..60aa2a0d695ba577ff87624d479f1eb25c8f1caf Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Factory differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/GB b/lib/python3.12/site-packages/pytz/zoneinfo/GB new file mode 100644 index 0000000000000000000000000000000000000000..5ad74220e83e6c30a0aeefda4a00271b8ebdfcad Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/GB differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/GB-Eire b/lib/python3.12/site-packages/pytz/zoneinfo/GB-Eire new file mode 100644 index 0000000000000000000000000000000000000000..5ad74220e83e6c30a0aeefda4a00271b8ebdfcad Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/GB-Eire differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/GMT b/lib/python3.12/site-packages/pytz/zoneinfo/GMT new file mode 100644 index 0000000000000000000000000000000000000000..c63474664a289aa3c3c0d8b2ce06d484679754c0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/GMT differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/GMT+0 b/lib/python3.12/site-packages/pytz/zoneinfo/GMT+0 new file mode 100644 index 0000000000000000000000000000000000000000..c63474664a289aa3c3c0d8b2ce06d484679754c0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/GMT+0 differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/GMT-0 b/lib/python3.12/site-packages/pytz/zoneinfo/GMT-0 new file mode 100644 index 0000000000000000000000000000000000000000..c63474664a289aa3c3c0d8b2ce06d484679754c0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/GMT-0 differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/GMT0 b/lib/python3.12/site-packages/pytz/zoneinfo/GMT0 new file mode 100644 index 0000000000000000000000000000000000000000..c63474664a289aa3c3c0d8b2ce06d484679754c0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/GMT0 differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Greenwich b/lib/python3.12/site-packages/pytz/zoneinfo/Greenwich new file mode 100644 index 0000000000000000000000000000000000000000..c63474664a289aa3c3c0d8b2ce06d484679754c0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Greenwich differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/HST b/lib/python3.12/site-packages/pytz/zoneinfo/HST new file mode 100644 index 0000000000000000000000000000000000000000..c7cd060159bd22fc5e6f10ac5a2089afb2c19c6a Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/HST differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Hongkong b/lib/python3.12/site-packages/pytz/zoneinfo/Hongkong new file mode 100644 index 0000000000000000000000000000000000000000..f9f7b134dd5c35e4718c6fa3697024cb95442c53 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Hongkong differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Iceland b/lib/python3.12/site-packages/pytz/zoneinfo/Iceland new file mode 100644 index 0000000000000000000000000000000000000000..28b32ab2e0b9053f39a91d9f28b6072e41423954 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Iceland differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Iran b/lib/python3.12/site-packages/pytz/zoneinfo/Iran new file mode 100644 index 0000000000000000000000000000000000000000..78f28cc402027bd337cf67b937bf8c31a75ee8b3 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Iran differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Israel b/lib/python3.12/site-packages/pytz/zoneinfo/Israel new file mode 100644 index 0000000000000000000000000000000000000000..1ebd0664aa29c0abd722661f761031ec0304631c Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Israel differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Jamaica b/lib/python3.12/site-packages/pytz/zoneinfo/Jamaica new file mode 100644 index 0000000000000000000000000000000000000000..2a9b7fd52d37a1ffe9fc589daa04d88c6c71a6e0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Jamaica differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Japan b/lib/python3.12/site-packages/pytz/zoneinfo/Japan new file mode 100644 index 0000000000000000000000000000000000000000..26f4d34d67b46513491f26c2e661c6e653cc130d Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Japan differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Kwajalein b/lib/python3.12/site-packages/pytz/zoneinfo/Kwajalein new file mode 100644 index 0000000000000000000000000000000000000000..1887a607422edd499fdf24afe91a04294f1caf6f Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Kwajalein differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Libya b/lib/python3.12/site-packages/pytz/zoneinfo/Libya new file mode 100644 index 0000000000000000000000000000000000000000..07b393bb7db14cef1e906ebe63cfbbe8cddc79d5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Libya differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/MET b/lib/python3.12/site-packages/pytz/zoneinfo/MET new file mode 100644 index 0000000000000000000000000000000000000000..40d7124e5346af056c75e2f7012a51d94e8154b7 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/MET differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/MST b/lib/python3.12/site-packages/pytz/zoneinfo/MST new file mode 100644 index 0000000000000000000000000000000000000000..ab37e845566aa95659b7b85be0051d0c67a7e53a Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/MST differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/MST7MDT b/lib/python3.12/site-packages/pytz/zoneinfo/MST7MDT new file mode 100644 index 0000000000000000000000000000000000000000..abb2b974a47eb3e5c8b4f5d4370baf4898b239ab Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/MST7MDT differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/NZ b/lib/python3.12/site-packages/pytz/zoneinfo/NZ new file mode 100644 index 0000000000000000000000000000000000000000..6575fdce31183d8238b18f2f30ab5b9227c7071c Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/NZ differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/NZ-CHAT b/lib/python3.12/site-packages/pytz/zoneinfo/NZ-CHAT new file mode 100644 index 0000000000000000000000000000000000000000..bde46cf7e4b7909714b93cee39a5d953387d756c Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/NZ-CHAT differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Navajo b/lib/python3.12/site-packages/pytz/zoneinfo/Navajo new file mode 100644 index 0000000000000000000000000000000000000000..abb2b974a47eb3e5c8b4f5d4370baf4898b239ab Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Navajo differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/PRC b/lib/python3.12/site-packages/pytz/zoneinfo/PRC new file mode 100644 index 0000000000000000000000000000000000000000..91f6f8bc2e234bafd484146986bdb289082c3588 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/PRC differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/PST8PDT b/lib/python3.12/site-packages/pytz/zoneinfo/PST8PDT new file mode 100644 index 0000000000000000000000000000000000000000..610e7af5fc13d9784de30d272c7c39d7938873a0 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/PST8PDT differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Poland b/lib/python3.12/site-packages/pytz/zoneinfo/Poland new file mode 100644 index 0000000000000000000000000000000000000000..e33cf67171da78aa9e6eb02e50f9b9603da4c3f4 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Poland differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Portugal b/lib/python3.12/site-packages/pytz/zoneinfo/Portugal new file mode 100644 index 0000000000000000000000000000000000000000..616de167b8a1c534c1c526e7bbff82f7d5da94f2 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Portugal differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/ROC b/lib/python3.12/site-packages/pytz/zoneinfo/ROC new file mode 100644 index 0000000000000000000000000000000000000000..24c43444b6751343d2915843d03e55753d4d7359 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/ROC differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/ROK b/lib/python3.12/site-packages/pytz/zoneinfo/ROK new file mode 100644 index 0000000000000000000000000000000000000000..96199e73e73aafacd89e48cb2855a96d7a134e1d Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/ROK differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Singapore b/lib/python3.12/site-packages/pytz/zoneinfo/Singapore new file mode 100644 index 0000000000000000000000000000000000000000..3d9f191e3acad56fb4f73ab44a74c50cdf91c12e Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Singapore differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Turkey b/lib/python3.12/site-packages/pytz/zoneinfo/Turkey new file mode 100644 index 0000000000000000000000000000000000000000..7c2336dd80c3c9cbf71cb53d2b2c1f89a65a8ba5 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Turkey differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/UCT b/lib/python3.12/site-packages/pytz/zoneinfo/UCT new file mode 100644 index 0000000000000000000000000000000000000000..91558be0c2bf903b2364215ba26d5227d6126508 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/UCT differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/UTC b/lib/python3.12/site-packages/pytz/zoneinfo/UTC new file mode 100644 index 0000000000000000000000000000000000000000..91558be0c2bf903b2364215ba26d5227d6126508 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/UTC differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Universal b/lib/python3.12/site-packages/pytz/zoneinfo/Universal new file mode 100644 index 0000000000000000000000000000000000000000..91558be0c2bf903b2364215ba26d5227d6126508 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Universal differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/W-SU b/lib/python3.12/site-packages/pytz/zoneinfo/W-SU new file mode 100644 index 0000000000000000000000000000000000000000..ddb3f4e99a1030f33b56fad986c8d9c16e59eb32 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/W-SU differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/WET b/lib/python3.12/site-packages/pytz/zoneinfo/WET new file mode 100644 index 0000000000000000000000000000000000000000..616de167b8a1c534c1c526e7bbff82f7d5da94f2 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/WET differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/Zulu b/lib/python3.12/site-packages/pytz/zoneinfo/Zulu new file mode 100644 index 0000000000000000000000000000000000000000..91558be0c2bf903b2364215ba26d5227d6126508 Binary files /dev/null and b/lib/python3.12/site-packages/pytz/zoneinfo/Zulu differ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/iso3166.tab b/lib/python3.12/site-packages/pytz/zoneinfo/iso3166.tab new file mode 100644 index 0000000000000000000000000000000000000000..402c015ec6b1071499155f7d28739db68be7763f --- /dev/null +++ b/lib/python3.12/site-packages/pytz/zoneinfo/iso3166.tab @@ -0,0 +1,279 @@ +# ISO 3166 alpha-2 country codes +# +# This file is in the public domain, so clarified as of +# 2009-05-17 by Arthur David Olson. +# +# From Paul Eggert (2023-09-06): +# This file contains a table of two-letter country codes. Columns are +# separated by a single tab. Lines beginning with '#' are comments. +# All text uses UTF-8 encoding. The columns of the table are as follows: +# +# 1. ISO 3166-1 alpha-2 country code, current as of +# ISO/TC 46 N1108 (2023-04-05). See: ISO/TC 46 Documents +# https://www.iso.org/committee/48750.html?view=documents +# 2. The usual English name for the coded region. This sometimes +# departs from ISO-listed names, sometimes so that sorted subsets +# of names are useful (e.g., "Samoa (American)" and "Samoa +# (western)" rather than "American Samoa" and "Samoa"), +# sometimes to avoid confusion among non-experts (e.g., +# "Czech Republic" and "Turkey" rather than "Czechia" and "Türkiye"), +# and sometimes to omit needless detail or churn (e.g., "Netherlands" +# rather than "Netherlands (the)" or "Netherlands (Kingdom of the)"). +# +# The table is sorted by country code. +# +# This table is intended as an aid for users, to help them select time +# zone data appropriate for their practical needs. It is not intended +# to take or endorse any position on legal or territorial claims. +# +#country- +#code name of country, territory, area, or subdivision +AD Andorra +AE United Arab Emirates +AF Afghanistan +AG Antigua & Barbuda +AI Anguilla +AL Albania +AM Armenia +AO Angola +AQ Antarctica +AR Argentina +AS Samoa (American) +AT Austria +AU Australia +AW Aruba +AX Åland Islands +AZ Azerbaijan +BA Bosnia & Herzegovina +BB Barbados +BD Bangladesh +BE Belgium +BF Burkina Faso +BG Bulgaria +BH Bahrain +BI Burundi +BJ Benin +BL St Barthelemy +BM Bermuda +BN Brunei +BO Bolivia +BQ Caribbean NL +BR Brazil +BS Bahamas +BT Bhutan +BV Bouvet Island +BW Botswana +BY Belarus +BZ Belize +CA Canada +CC Cocos (Keeling) Islands +CD Congo (Dem. Rep.) +CF Central African Rep. +CG Congo (Rep.) +CH Switzerland +CI Côte d'Ivoire +CK Cook Islands +CL Chile +CM Cameroon +CN China +CO Colombia +CR Costa Rica +CU Cuba +CV Cape Verde +CW Curaçao +CX Christmas Island +CY Cyprus +CZ Czech Republic +DE Germany +DJ Djibouti +DK Denmark +DM Dominica +DO Dominican Republic +DZ Algeria +EC Ecuador +EE Estonia +EG Egypt +EH Western Sahara +ER Eritrea +ES Spain +ET Ethiopia +FI Finland +FJ Fiji +FK Falkland Islands +FM Micronesia +FO Faroe Islands +FR France +GA Gabon +GB Britain (UK) +GD Grenada +GE Georgia +GF French Guiana +GG Guernsey +GH Ghana +GI Gibraltar +GL Greenland +GM Gambia +GN Guinea +GP Guadeloupe +GQ Equatorial Guinea +GR Greece +GS South Georgia & the South Sandwich Islands +GT Guatemala +GU Guam +GW Guinea-Bissau +GY Guyana +HK Hong Kong +HM Heard Island & McDonald Islands +HN Honduras +HR Croatia +HT Haiti +HU Hungary +ID Indonesia +IE Ireland +IL Israel +IM Isle of Man +IN India +IO British Indian Ocean Territory +IQ Iraq +IR Iran +IS Iceland +IT Italy +JE Jersey +JM Jamaica +JO Jordan +JP Japan +KE Kenya +KG Kyrgyzstan +KH Cambodia +KI Kiribati +KM Comoros +KN St Kitts & Nevis +KP Korea (North) +KR Korea (South) +KW Kuwait +KY Cayman Islands +KZ Kazakhstan +LA Laos +LB Lebanon +LC St Lucia +LI Liechtenstein +LK Sri Lanka +LR Liberia +LS Lesotho +LT Lithuania +LU Luxembourg +LV Latvia +LY Libya +MA Morocco +MC Monaco +MD Moldova +ME Montenegro +MF St Martin (French) +MG Madagascar +MH Marshall Islands +MK North Macedonia +ML Mali +MM Myanmar (Burma) +MN Mongolia +MO Macau +MP Northern Mariana Islands +MQ Martinique +MR Mauritania +MS Montserrat +MT Malta +MU Mauritius +MV Maldives +MW Malawi +MX Mexico +MY Malaysia +MZ Mozambique +NA Namibia +NC New Caledonia +NE Niger +NF Norfolk Island +NG Nigeria +NI Nicaragua +NL Netherlands +NO Norway +NP Nepal +NR Nauru +NU Niue +NZ New Zealand +OM Oman +PA Panama +PE Peru +PF French Polynesia +PG Papua New Guinea +PH Philippines +PK Pakistan +PL Poland +PM St Pierre & Miquelon +PN Pitcairn +PR Puerto Rico +PS Palestine +PT Portugal +PW Palau +PY Paraguay +QA Qatar +RE Réunion +RO Romania +RS Serbia +RU Russia +RW Rwanda +SA Saudi Arabia +SB Solomon Islands +SC Seychelles +SD Sudan +SE Sweden +SG Singapore +SH St Helena +SI Slovenia +SJ Svalbard & Jan Mayen +SK Slovakia +SL Sierra Leone +SM San Marino +SN Senegal +SO Somalia +SR Suriname +SS South Sudan +ST Sao Tome & Principe +SV El Salvador +SX St Maarten (Dutch) +SY Syria +SZ Eswatini (Swaziland) +TC Turks & Caicos Is +TD Chad +TF French S. Terr. +TG Togo +TH Thailand +TJ Tajikistan +TK Tokelau +TL East Timor +TM Turkmenistan +TN Tunisia +TO Tonga +TR Turkey +TT Trinidad & Tobago +TV Tuvalu +TW Taiwan +TZ Tanzania +UA Ukraine +UG Uganda +UM US minor outlying islands +US United States +UY Uruguay +UZ Uzbekistan +VA Vatican City +VC St Vincent +VE Venezuela +VG Virgin Islands (UK) +VI Virgin Islands (US) +VN Vietnam +VU Vanuatu +WF Wallis & Futuna +WS Samoa (western) +YE Yemen +YT Mayotte +ZA South Africa +ZM Zambia +ZW Zimbabwe diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/leapseconds b/lib/python3.12/site-packages/pytz/zoneinfo/leapseconds new file mode 100644 index 0000000000000000000000000000000000000000..76f771427f25b91e6944ffd2fee6031bfd73979a --- /dev/null +++ b/lib/python3.12/site-packages/pytz/zoneinfo/leapseconds @@ -0,0 +1,79 @@ +# Allowance for leap seconds added to each time zone file. + +# This file is in the public domain. + +# This file is generated automatically from the data in the public-domain +# NIST/IERS format leap-seconds.list file, which can be copied from +# +# or, in a variant with different comments, from +# . +# For more about leap-seconds.list, please see +# The NTP Timescale and Leap Seconds +# . + +# The rules for leap seconds are specified in Annex 1 (Time scales) of: +# Standard-frequency and time-signal emissions. +# International Telecommunication Union - Radiocommunication Sector +# (ITU-R) Recommendation TF.460-6 (02/2002) +# . +# The International Earth Rotation and Reference Systems Service (IERS) +# periodically uses leap seconds to keep UTC to within 0.9 s of UT1 +# (a proxy for Earth's angle in space as measured by astronomers) +# and publishes leap second data in a copyrighted file +# . +# See: Levine J. Coordinated Universal Time and the leap second. +# URSI Radio Sci Bull. 2016;89(4):30-6. doi:10.23919/URSIRSB.2016.7909995 +# . + +# There were no leap seconds before 1972, as no official mechanism +# accounted for the discrepancy between atomic time (TAI) and the earth's +# rotation. The first ("1 Jan 1972") data line in leap-seconds.list +# does not denote a leap second; it denotes the start of the current definition +# of UTC. + +# All leap-seconds are Stationary (S) at the given UTC time. +# The correction (+ or -) is made at the given time, so in the unlikely +# event of a negative leap second, a line would look like this: +# Leap YEAR MON DAY 23:59:59 - S +# Typical lines look like this: +# Leap YEAR MON DAY 23:59:60 + S +Leap 1972 Jun 30 23:59:60 + S +Leap 1972 Dec 31 23:59:60 + S +Leap 1973 Dec 31 23:59:60 + S +Leap 1974 Dec 31 23:59:60 + S +Leap 1975 Dec 31 23:59:60 + S +Leap 1976 Dec 31 23:59:60 + S +Leap 1977 Dec 31 23:59:60 + S +Leap 1978 Dec 31 23:59:60 + S +Leap 1979 Dec 31 23:59:60 + S +Leap 1981 Jun 30 23:59:60 + S +Leap 1982 Jun 30 23:59:60 + S +Leap 1983 Jun 30 23:59:60 + S +Leap 1985 Jun 30 23:59:60 + S +Leap 1987 Dec 31 23:59:60 + S +Leap 1989 Dec 31 23:59:60 + S +Leap 1990 Dec 31 23:59:60 + S +Leap 1992 Jun 30 23:59:60 + S +Leap 1993 Jun 30 23:59:60 + S +Leap 1994 Jun 30 23:59:60 + S +Leap 1995 Dec 31 23:59:60 + S +Leap 1997 Jun 30 23:59:60 + S +Leap 1998 Dec 31 23:59:60 + S +Leap 2005 Dec 31 23:59:60 + S +Leap 2008 Dec 31 23:59:60 + S +Leap 2012 Jun 30 23:59:60 + S +Leap 2015 Jun 30 23:59:60 + S +Leap 2016 Dec 31 23:59:60 + S + +# UTC timestamp when this leap second list expires. +# Any additional leap seconds will come after this. +# This Expires line is commented out for now, +# so that pre-2020a zic implementations do not reject this file. +#Expires 2025 Dec 28 00:00:00 + +# POSIX timestamps for the data in this file: +#updated 1736208000 (2025-01-07 00:00:00 UTC) +#expires 1766880000 (2025-12-28 00:00:00 UTC) + +# Updated through IERS Bulletin C (https://hpiers.obspm.fr/iers/bul/bulc/bulletinc.dat) +# File expires on 28 December 2025 diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/tzdata.zi b/lib/python3.12/site-packages/pytz/zoneinfo/tzdata.zi new file mode 100644 index 0000000000000000000000000000000000000000..0bcae52e5efa123d39e17553fda5ebade3bcc394 --- /dev/null +++ b/lib/python3.12/site-packages/pytz/zoneinfo/tzdata.zi @@ -0,0 +1,4300 @@ +# version unknown +# This zic input file is in the public domain. +R d 1916 o - Jun 14 23s 1 S +R d 1916 1919 - O Su>=1 23s 0 - +R d 1917 o - Mar 24 23s 1 S +R d 1918 o - Mar 9 23s 1 S +R d 1919 o - Mar 1 23s 1 S +R d 1920 o - F 14 23s 1 S +R d 1920 o - O 23 23s 0 - +R d 1921 o - Mar 14 23s 1 S +R d 1921 o - Jun 21 23s 0 - +R d 1939 o - S 11 23s 1 S +R d 1939 o - N 19 1 0 - +R d 1944 1945 - Ap M>=1 2 1 S +R d 1944 o - O 8 2 0 - +R d 1945 o - S 16 1 0 - +R d 1971 o - Ap 25 23s 1 S +R d 1971 o - S 26 23s 0 - +R d 1977 o - May 6 0 1 S +R d 1977 o - O 21 0 0 - +R d 1978 o - Mar 24 1 1 S +R d 1978 o - S 22 3 0 - +R d 1980 o - Ap 25 0 1 S +R d 1980 o - O 31 2 0 - +R K 1940 o - Jul 15 0 1 S +R K 1940 o - O 1 0 0 - +R K 1941 o - Ap 15 0 1 S +R K 1941 o - S 16 0 0 - +R K 1942 1944 - Ap 1 0 1 S +R K 1942 o - O 27 0 0 - +R K 1943 1945 - N 1 0 0 - +R K 1945 o - Ap 16 0 1 S +R K 1957 o - May 10 0 1 S +R K 1957 1958 - O 1 0 0 - +R K 1958 o - May 1 0 1 S +R K 1959 1981 - May 1 1 1 S +R K 1959 1965 - S 30 3 0 - +R K 1966 1994 - O 1 3 0 - +R K 1982 o - Jul 25 1 1 S +R K 1983 o - Jul 12 1 1 S +R K 1984 1988 - May 1 1 1 S +R K 1989 o - May 6 1 1 S +R K 1990 1994 - May 1 1 1 S +R K 1995 2010 - Ap lastF 0s 1 S +R K 1995 2005 - S lastTh 24 0 - +R K 2006 o - S 21 24 0 - +R K 2007 o - S Th>=1 24 0 - +R K 2008 o - Au lastTh 24 0 - +R K 2009 o - Au 20 24 0 - +R K 2010 o - Au 10 24 0 - +R K 2010 o - S 9 24 1 S +R K 2010 o - S lastTh 24 0 - +R K 2014 o - May 15 24 1 S +R K 2014 o - Jun 26 24 0 - +R K 2014 o - Jul 31 24 1 S +R K 2014 o - S lastTh 24 0 - +R K 2023 ma - Ap lastF 0 1 S +R K 2023 ma - O lastTh 24 0 - +R L 1951 o - O 14 2 1 S +R L 1952 o - Ja 1 0 0 - +R L 1953 o - O 9 2 1 S +R L 1954 o - Ja 1 0 0 - +R L 1955 o - S 30 0 1 S +R L 1956 o - Ja 1 0 0 - +R L 1982 1984 - Ap 1 0 1 S +R L 1982 1985 - O 1 0 0 - +R L 1985 o - Ap 6 0 1 S +R L 1986 o - Ap 4 0 1 S +R L 1986 o - O 3 0 0 - +R L 1987 1989 - Ap 1 0 1 S +R L 1987 1989 - O 1 0 0 - +R L 1997 o - Ap 4 0 1 S +R L 1997 o - O 4 0 0 - +R L 2013 o - Mar lastF 1 1 S +R L 2013 o - O lastF 2 0 - +R MU 1982 o - O 10 0 1 - +R MU 1983 o - Mar 21 0 0 - +R MU 2008 o - O lastSu 2 1 - +R MU 2009 o - Mar lastSu 2 0 - +R M 1939 o - S 12 0 1 - +R M 1939 o - N 19 0 0 - +R M 1940 o - F 25 0 1 - +R M 1945 o - N 18 0 0 - +R M 1950 o - Jun 11 0 1 - +R M 1950 o - O 29 0 0 - +R M 1967 o - Jun 3 12 1 - +R M 1967 o - O 1 0 0 - +R M 1974 o - Jun 24 0 1 - +R M 1974 o - S 1 0 0 - +R M 1976 1977 - May 1 0 1 - +R M 1976 o - Au 1 0 0 - +R M 1977 o - S 28 0 0 - +R M 1978 o - Jun 1 0 1 - +R M 1978 o - Au 4 0 0 - +R M 2008 o - Jun 1 0 1 - +R M 2008 o - S 1 0 0 - +R M 2009 o - Jun 1 0 1 - +R M 2009 o - Au 21 0 0 - +R M 2010 o - May 2 0 1 - +R M 2010 o - Au 8 0 0 - +R M 2011 o - Ap 3 0 1 - +R M 2011 o - Jul 31 0 0 - +R M 2012 2013 - Ap lastSu 2 1 - +R M 2012 o - Jul 20 3 0 - +R M 2012 o - Au 20 2 1 - +R M 2012 o - S 30 3 0 - +R M 2013 o - Jul 7 3 0 - +R M 2013 o - Au 10 2 1 - +R M 2013 2018 - O lastSu 3 0 - +R M 2014 2018 - Mar lastSu 2 1 - +R M 2014 o - Jun 28 3 0 - +R M 2014 o - Au 2 2 1 - +R M 2015 o - Jun 14 3 0 - +R M 2015 o - Jul 19 2 1 - +R M 2016 o - Jun 5 3 0 - +R M 2016 o - Jul 10 2 1 - +R M 2017 o - May 21 3 0 - +R M 2017 o - Jul 2 2 1 - +R M 2018 o - May 13 3 0 - +R M 2018 o - Jun 17 2 1 - +R M 2019 o - May 5 3 -1 - +R M 2019 o - Jun 9 2 0 - +R M 2020 o - Ap 19 3 -1 - +R M 2020 o - May 31 2 0 - +R M 2021 o - Ap 11 3 -1 - +R M 2021 o - May 16 2 0 - +R M 2022 o - Mar 27 3 -1 - +R M 2022 o - May 8 2 0 - +R M 2023 o - Mar 19 3 -1 - +R M 2023 o - Ap 23 2 0 - +R M 2024 o - Mar 10 3 -1 - +R M 2024 o - Ap 14 2 0 - +R M 2025 o - F 23 3 -1 - +R M 2025 o - Ap 6 2 0 - +R M 2026 o - F 15 3 -1 - +R M 2026 o - Mar 22 2 0 - +R M 2027 o - F 7 3 -1 - +R M 2027 o - Mar 14 2 0 - +R M 2028 o - Ja 23 3 -1 - +R M 2028 o - Mar 5 2 0 - +R M 2029 o - Ja 14 3 -1 - +R M 2029 o - F 18 2 0 - +R M 2029 o - D 30 3 -1 - +R M 2030 o - F 10 2 0 - +R M 2030 o - D 22 3 -1 - +R M 2031 o - Ja 26 2 0 - +R M 2031 o - D 14 3 -1 - +R M 2032 o - Ja 18 2 0 - +R M 2032 o - N 28 3 -1 - +R M 2033 o - Ja 9 2 0 - +R M 2033 o - N 20 3 -1 - +R M 2033 o - D 25 2 0 - +R M 2034 o - N 5 3 -1 - +R M 2034 o - D 17 2 0 - +R M 2035 o - O 28 3 -1 - +R M 2035 o - D 9 2 0 - +R M 2036 o - O 19 3 -1 - +R M 2036 o - N 23 2 0 - +R M 2037 o - O 4 3 -1 - +R M 2037 o - N 15 2 0 - +R M 2038 o - S 26 3 -1 - +R M 2038 o - O 31 2 0 - +R M 2039 o - S 18 3 -1 - +R M 2039 o - O 23 2 0 - +R M 2040 o - S 2 3 -1 - +R M 2040 o - O 14 2 0 - +R M 2041 o - Au 25 3 -1 - +R M 2041 o - S 29 2 0 - +R M 2042 o - Au 10 3 -1 - +R M 2042 o - S 21 2 0 - +R M 2043 o - Au 2 3 -1 - +R M 2043 o - S 13 2 0 - +R M 2044 o - Jul 24 3 -1 - +R M 2044 o - Au 28 2 0 - +R M 2045 o - Jul 9 3 -1 - +R M 2045 o - Au 20 2 0 - +R M 2046 o - Jul 1 3 -1 - +R M 2046 o - Au 5 2 0 - +R M 2047 o - Jun 23 3 -1 - +R M 2047 o - Jul 28 2 0 - +R M 2048 o - Jun 7 3 -1 - +R M 2048 o - Jul 19 2 0 - +R M 2049 o - May 30 3 -1 - +R M 2049 o - Jul 4 2 0 - +R M 2050 o - May 15 3 -1 - +R M 2050 o - Jun 26 2 0 - +R M 2051 o - May 7 3 -1 - +R M 2051 o - Jun 18 2 0 - +R M 2052 o - Ap 28 3 -1 - +R M 2052 o - Jun 2 2 0 - +R M 2053 o - Ap 13 3 -1 - +R M 2053 o - May 25 2 0 - +R M 2054 o - Ap 5 3 -1 - +R M 2054 o - May 10 2 0 - +R M 2055 o - Mar 28 3 -1 - +R M 2055 o - May 2 2 0 - +R M 2056 o - Mar 12 3 -1 - +R M 2056 o - Ap 23 2 0 - +R M 2057 o - Mar 4 3 -1 - +R M 2057 o - Ap 8 2 0 - +R M 2058 o - F 17 3 -1 - +R M 2058 o - Mar 31 2 0 - +R M 2059 o - F 9 3 -1 - +R M 2059 o - Mar 23 2 0 - +R M 2060 o - F 1 3 -1 - +R M 2060 o - Mar 7 2 0 - +R M 2061 o - Ja 16 3 -1 - +R M 2061 o - F 27 2 0 - +R M 2062 o - Ja 8 3 -1 - +R M 2062 o - F 12 2 0 - +R M 2062 o - D 31 3 -1 - +R M 2063 o - F 4 2 0 - +R M 2063 o - D 16 3 -1 - +R M 2064 o - Ja 27 2 0 - +R M 2064 o - D 7 3 -1 - +R M 2065 o - Ja 11 2 0 - +R M 2065 o - N 22 3 -1 - +R M 2066 o - Ja 3 2 0 - +R M 2066 o - N 14 3 -1 - +R M 2066 o - D 26 2 0 - +R M 2067 o - N 6 3 -1 - +R M 2067 o - D 11 2 0 - +R M 2068 o - O 21 3 -1 - +R M 2068 o - D 2 2 0 - +R M 2069 o - O 13 3 -1 - +R M 2069 o - N 17 2 0 - +R M 2070 o - O 5 3 -1 - +R M 2070 o - N 9 2 0 - +R M 2071 o - S 20 3 -1 - +R M 2071 o - N 1 2 0 - +R M 2072 o - S 11 3 -1 - +R M 2072 o - O 16 2 0 - +R M 2073 o - Au 27 3 -1 - +R M 2073 o - O 8 2 0 - +R M 2074 o - Au 19 3 -1 - +R M 2074 o - S 30 2 0 - +R M 2075 o - Au 11 3 -1 - +R M 2075 o - S 15 2 0 - +R M 2076 o - Jul 26 3 -1 - +R M 2076 o - S 6 2 0 - +R M 2077 o - Jul 18 3 -1 - +R M 2077 o - Au 22 2 0 - +R M 2078 o - Jul 10 3 -1 - +R M 2078 o - Au 14 2 0 - +R M 2079 o - Jun 25 3 -1 - +R M 2079 o - Au 6 2 0 - +R M 2080 o - Jun 16 3 -1 - +R M 2080 o - Jul 21 2 0 - +R M 2081 o - Jun 1 3 -1 - +R M 2081 o - Jul 13 2 0 - +R M 2082 o - May 24 3 -1 - +R M 2082 o - Jun 28 2 0 - +R M 2083 o - May 16 3 -1 - +R M 2083 o - Jun 20 2 0 - +R M 2084 o - Ap 30 3 -1 - +R M 2084 o - Jun 11 2 0 - +R M 2085 o - Ap 22 3 -1 - +R M 2085 o - May 27 2 0 - +R M 2086 o - Ap 14 3 -1 - +R M 2086 o - May 19 2 0 - +R M 2087 o - Mar 30 3 -1 - +R M 2087 o - May 11 2 0 - +R NA 1994 o - Mar 21 0 -1 WAT +R NA 1994 2017 - S Su>=1 2 0 CAT +R NA 1995 2017 - Ap Su>=1 2 -1 WAT +R SA 1942 1943 - S Su>=15 2 1 - +R SA 1943 1944 - Mar Su>=15 2 0 - +R SD 1970 o - May 1 0 1 S +R SD 1970 1985 - O 15 0 0 - +R SD 1971 o - Ap 30 0 1 S +R SD 1972 1985 - Ap lastSu 0 1 S +R n 1939 o - Ap 15 23s 1 S +R n 1939 o - N 18 23s 0 - +R n 1940 o - F 25 23s 1 S +R n 1941 o - O 6 0 0 - +R n 1942 o - Mar 9 0 1 S +R n 1942 o - N 2 3 0 - +R n 1943 o - Mar 29 2 1 S +R n 1943 o - Ap 17 2 0 - +R n 1943 o - Ap 25 2 1 S +R n 1943 o - O 4 2 0 - +R n 1944 1945 - Ap M>=1 2 1 S +R n 1944 o - O 8 0 0 - +R n 1945 o - S 16 0 0 - +R n 1977 o - Ap 30 0s 1 S +R n 1977 o - S 24 0s 0 - +R n 1978 o - May 1 0s 1 S +R n 1978 o - O 1 0s 0 - +R n 1988 o - Jun 1 0s 1 S +R n 1988 1990 - S lastSu 0s 0 - +R n 1989 o - Mar 26 0s 1 S +R n 1990 o - May 1 0s 1 S +R n 2005 o - May 1 0s 1 S +R n 2005 o - S 30 1s 0 - +R n 2006 2008 - Mar lastSu 2s 1 S +R n 2006 2008 - O lastSu 2s 0 - +R Tr 2005 ma - Mar lastSu 1u 2 +02 +R Tr 2004 ma - O lastSu 1u 0 +00 +R AM 2011 o - Mar lastSu 2s 1 - +R AM 2011 o - O lastSu 2s 0 - +R AZ 1997 2015 - Mar lastSu 4 1 - +R AZ 1997 2015 - O lastSu 5 0 - +R BD 2009 o - Jun 19 23 1 - +R BD 2009 o - D 31 24 0 - +R Sh 1919 o - Ap 12 24 1 D +R Sh 1919 o - S 30 24 0 S +R Sh 1940 o - Jun 1 0 1 D +R Sh 1940 o - O 12 24 0 S +R Sh 1941 o - Mar 15 0 1 D +R Sh 1941 o - N 1 24 0 S +R Sh 1942 o - Ja 31 0 1 D +R Sh 1945 o - S 1 24 0 S +R Sh 1946 o - May 15 0 1 D +R Sh 1946 o - S 30 24 0 S +R Sh 1947 o - Ap 15 0 1 D +R Sh 1947 o - O 31 24 0 S +R Sh 1948 1949 - May 1 0 1 D +R Sh 1948 1949 - S 30 24 0 S +R CN 1986 o - May 4 2 1 D +R CN 1986 1991 - S Su>=11 2 0 S +R CN 1987 1991 - Ap Su>=11 2 1 D +R HK 1946 o - Ap 21 0 1 S +R HK 1946 o - D 1 3:30s 0 - +R HK 1947 o - Ap 13 3:30s 1 S +R HK 1947 o - N 30 3:30s 0 - +R HK 1948 o - May 2 3:30s 1 S +R HK 1948 1952 - O Su>=28 3:30s 0 - +R HK 1949 1953 - Ap Su>=1 3:30 1 S +R HK 1953 1964 - O Su>=31 3:30 0 - +R HK 1954 1964 - Mar Su>=18 3:30 1 S +R HK 1965 1976 - Ap Su>=16 3:30 1 S +R HK 1965 1976 - O Su>=16 3:30 0 - +R HK 1973 o - D 30 3:30 1 S +R HK 1979 o - May 13 3:30 1 S +R HK 1979 o - O 21 3:30 0 - +R f 1946 o - May 15 0 1 D +R f 1946 o - O 1 0 0 S +R f 1947 o - Ap 15 0 1 D +R f 1947 o - N 1 0 0 S +R f 1948 1951 - May 1 0 1 D +R f 1948 1951 - O 1 0 0 S +R f 1952 o - Mar 1 0 1 D +R f 1952 1954 - N 1 0 0 S +R f 1953 1959 - Ap 1 0 1 D +R f 1955 1961 - O 1 0 0 S +R f 1960 1961 - Jun 1 0 1 D +R f 1974 1975 - Ap 1 0 1 D +R f 1974 1975 - O 1 0 0 S +R f 1979 o - Jul 1 0 1 D +R f 1979 o - O 1 0 0 S +R _ 1942 1943 - Ap 30 23 1 - +R _ 1942 o - N 17 23 0 - +R _ 1943 o - S 30 23 0 S +R _ 1946 o - Ap 30 23s 1 D +R _ 1946 o - S 30 23s 0 S +R _ 1947 o - Ap 19 23s 1 D +R _ 1947 o - N 30 23s 0 S +R _ 1948 o - May 2 23s 1 D +R _ 1948 o - O 31 23s 0 S +R _ 1949 1950 - Ap Sa>=1 23s 1 D +R _ 1949 1950 - O lastSa 23s 0 S +R _ 1951 o - Mar 31 23s 1 D +R _ 1951 o - O 28 23s 0 S +R _ 1952 1953 - Ap Sa>=1 23s 1 D +R _ 1952 o - N 1 23s 0 S +R _ 1953 1954 - O lastSa 23s 0 S +R _ 1954 1956 - Mar Sa>=17 23s 1 D +R _ 1955 o - N 5 23s 0 S +R _ 1956 1964 - N Su>=1 3:30 0 S +R _ 1957 1964 - Mar Su>=18 3:30 1 D +R _ 1965 1973 - Ap Su>=16 3:30 1 D +R _ 1965 1966 - O Su>=16 2:30 0 S +R _ 1967 1976 - O Su>=16 3:30 0 S +R _ 1973 o - D 30 3:30 1 D +R _ 1975 1976 - Ap Su>=16 3:30 1 D +R _ 1979 o - May 13 3:30 1 D +R _ 1979 o - O Su>=16 3:30 0 S +R CY 1975 o - Ap 13 0 1 S +R CY 1975 o - O 12 0 0 - +R CY 1976 o - May 15 0 1 S +R CY 1976 o - O 11 0 0 - +R CY 1977 1980 - Ap Su>=1 0 1 S +R CY 1977 o - S 25 0 0 - +R CY 1978 o - O 2 0 0 - +R CY 1979 1997 - S lastSu 0 0 - +R CY 1981 1998 - Mar lastSu 0 1 S +R i 1910 o - Ja 1 0 0 - +R i 1977 o - Mar 21 23 1 - +R i 1977 o - O 20 24 0 - +R i 1978 o - Mar 24 24 1 - +R i 1978 o - Au 5 1 0 - +R i 1979 o - May 26 24 1 - +R i 1979 o - S 18 24 0 - +R i 1980 o - Mar 20 24 1 - +R i 1980 o - S 22 24 0 - +R i 1991 o - May 2 24 1 - +R i 1992 1995 - Mar 21 24 1 - +R i 1991 1995 - S 21 24 0 - +R i 1996 o - Mar 20 24 1 - +R i 1996 o - S 20 24 0 - +R i 1997 1999 - Mar 21 24 1 - +R i 1997 1999 - S 21 24 0 - +R i 2000 o - Mar 20 24 1 - +R i 2000 o - S 20 24 0 - +R i 2001 2003 - Mar 21 24 1 - +R i 2001 2003 - S 21 24 0 - +R i 2004 o - Mar 20 24 1 - +R i 2004 o - S 20 24 0 - +R i 2005 o - Mar 21 24 1 - +R i 2005 o - S 21 24 0 - +R i 2008 o - Mar 20 24 1 - +R i 2008 o - S 20 24 0 - +R i 2009 2011 - Mar 21 24 1 - +R i 2009 2011 - S 21 24 0 - +R i 2012 o - Mar 20 24 1 - +R i 2012 o - S 20 24 0 - +R i 2013 2015 - Mar 21 24 1 - +R i 2013 2015 - S 21 24 0 - +R i 2016 o - Mar 20 24 1 - +R i 2016 o - S 20 24 0 - +R i 2017 2019 - Mar 21 24 1 - +R i 2017 2019 - S 21 24 0 - +R i 2020 o - Mar 20 24 1 - +R i 2020 o - S 20 24 0 - +R i 2021 2022 - Mar 21 24 1 - +R i 2021 2022 - S 21 24 0 - +R IQ 1982 o - May 1 0 1 - +R IQ 1982 1984 - O 1 0 0 - +R IQ 1983 o - Mar 31 0 1 - +R IQ 1984 1985 - Ap 1 0 1 - +R IQ 1985 1990 - S lastSu 1s 0 - +R IQ 1986 1990 - Mar lastSu 1s 1 - +R IQ 1991 2007 - Ap 1 3s 1 - +R IQ 1991 2007 - O 1 3s 0 - +R Z 1940 o - May 31 24u 1 D +R Z 1940 o - S 30 24u 0 S +R Z 1940 o - N 16 24u 1 D +R Z 1942 1946 - O 31 24u 0 S +R Z 1943 1944 - Mar 31 24u 1 D +R Z 1945 1946 - Ap 15 24u 1 D +R Z 1948 o - May 22 24u 2 DD +R Z 1948 o - Au 31 24u 1 D +R Z 1948 1949 - O 31 24u 0 S +R Z 1949 o - Ap 30 24u 1 D +R Z 1950 o - Ap 15 24u 1 D +R Z 1950 o - S 14 24u 0 S +R Z 1951 o - Mar 31 24u 1 D +R Z 1951 o - N 10 24u 0 S +R Z 1952 o - Ap 19 24u 1 D +R Z 1952 o - O 18 24u 0 S +R Z 1953 o - Ap 11 24u 1 D +R Z 1953 o - S 12 24u 0 S +R Z 1954 o - Jun 12 24u 1 D +R Z 1954 o - S 11 24u 0 S +R Z 1955 o - Jun 11 24u 1 D +R Z 1955 o - S 10 24u 0 S +R Z 1956 o - Jun 2 24u 1 D +R Z 1956 o - S 29 24u 0 S +R Z 1957 o - Ap 27 24u 1 D +R Z 1957 o - S 21 24u 0 S +R Z 1974 o - Jul 6 24 1 D +R Z 1974 o - O 12 24 0 S +R Z 1975 o - Ap 19 24 1 D +R Z 1975 o - Au 30 24 0 S +R Z 1980 o - Au 2 24s 1 D +R Z 1980 o - S 13 24s 0 S +R Z 1984 o - May 5 24s 1 D +R Z 1984 o - Au 25 24s 0 S +R Z 1985 o - Ap 13 24 1 D +R Z 1985 o - Au 31 24 0 S +R Z 1986 o - May 17 24 1 D +R Z 1986 o - S 6 24 0 S +R Z 1987 o - Ap 14 24 1 D +R Z 1987 o - S 12 24 0 S +R Z 1988 o - Ap 9 24 1 D +R Z 1988 o - S 3 24 0 S +R Z 1989 o - Ap 29 24 1 D +R Z 1989 o - S 2 24 0 S +R Z 1990 o - Mar 24 24 1 D +R Z 1990 o - Au 25 24 0 S +R Z 1991 o - Mar 23 24 1 D +R Z 1991 o - Au 31 24 0 S +R Z 1992 o - Mar 28 24 1 D +R Z 1992 o - S 5 24 0 S +R Z 1993 o - Ap 2 0 1 D +R Z 1993 o - S 5 0 0 S +R Z 1994 o - Ap 1 0 1 D +R Z 1994 o - Au 28 0 0 S +R Z 1995 o - Mar 31 0 1 D +R Z 1995 o - S 3 0 0 S +R Z 1996 o - Mar 14 24 1 D +R Z 1996 o - S 15 24 0 S +R Z 1997 o - Mar 20 24 1 D +R Z 1997 o - S 13 24 0 S +R Z 1998 o - Mar 20 0 1 D +R Z 1998 o - S 6 0 0 S +R Z 1999 o - Ap 2 2 1 D +R Z 1999 o - S 3 2 0 S +R Z 2000 o - Ap 14 2 1 D +R Z 2000 o - O 6 1 0 S +R Z 2001 o - Ap 9 1 1 D +R Z 2001 o - S 24 1 0 S +R Z 2002 o - Mar 29 1 1 D +R Z 2002 o - O 7 1 0 S +R Z 2003 o - Mar 28 1 1 D +R Z 2003 o - O 3 1 0 S +R Z 2004 o - Ap 7 1 1 D +R Z 2004 o - S 22 1 0 S +R Z 2005 2012 - Ap F<=1 2 1 D +R Z 2005 o - O 9 2 0 S +R Z 2006 o - O 1 2 0 S +R Z 2007 o - S 16 2 0 S +R Z 2008 o - O 5 2 0 S +R Z 2009 o - S 27 2 0 S +R Z 2010 o - S 12 2 0 S +R Z 2011 o - O 2 2 0 S +R Z 2012 o - S 23 2 0 S +R Z 2013 ma - Mar F>=23 2 1 D +R Z 2013 ma - O lastSu 2 0 S +R JP 1948 o - May Sa>=1 24 1 D +R JP 1948 1951 - S Sa>=8 25 0 S +R JP 1949 o - Ap Sa>=1 24 1 D +R JP 1950 1951 - May Sa>=1 24 1 D +R J 1973 o - Jun 6 0 1 S +R J 1973 1975 - O 1 0 0 - +R J 1974 1977 - May 1 0 1 S +R J 1976 o - N 1 0 0 - +R J 1977 o - O 1 0 0 - +R J 1978 o - Ap 30 0 1 S +R J 1978 o - S 30 0 0 - +R J 1985 o - Ap 1 0 1 S +R J 1985 o - O 1 0 0 - +R J 1986 1988 - Ap F>=1 0 1 S +R J 1986 1990 - O F>=1 0 0 - +R J 1989 o - May 8 0 1 S +R J 1990 o - Ap 27 0 1 S +R J 1991 o - Ap 17 0 1 S +R J 1991 o - S 27 0 0 - +R J 1992 o - Ap 10 0 1 S +R J 1992 1993 - O F>=1 0 0 - +R J 1993 1998 - Ap F>=1 0 1 S +R J 1994 o - S F>=15 0 0 - +R J 1995 1998 - S F>=15 0s 0 - +R J 1999 o - Jul 1 0s 1 S +R J 1999 2002 - S lastF 0s 0 - +R J 2000 2001 - Mar lastTh 0s 1 S +R J 2002 2012 - Mar lastTh 24 1 S +R J 2003 o - O 24 0s 0 - +R J 2004 o - O 15 0s 0 - +R J 2005 o - S lastF 0s 0 - +R J 2006 2011 - O lastF 0s 0 - +R J 2013 o - D 20 0 0 - +R J 2014 2021 - Mar lastTh 24 1 S +R J 2014 2022 - O lastF 0s 0 - +R J 2022 o - F lastTh 24 1 S +R KG 1992 1996 - Ap Su>=7 0s 1 - +R KG 1992 1996 - S lastSu 0 0 - +R KG 1997 2005 - Mar lastSu 2:30 1 - +R KG 1997 2004 - O lastSu 2:30 0 - +R KR 1948 o - Jun 1 0 1 D +R KR 1948 o - S 12 24 0 S +R KR 1949 o - Ap 3 0 1 D +R KR 1949 1951 - S Sa>=7 24 0 S +R KR 1950 o - Ap 1 0 1 D +R KR 1951 o - May 6 0 1 D +R KR 1955 o - May 5 0 1 D +R KR 1955 o - S 8 24 0 S +R KR 1956 o - May 20 0 1 D +R KR 1956 o - S 29 24 0 S +R KR 1957 1960 - May Su>=1 0 1 D +R KR 1957 1960 - S Sa>=17 24 0 S +R KR 1987 1988 - May Su>=8 2 1 D +R KR 1987 1988 - O Su>=8 3 0 S +R l 1920 o - Mar 28 0 1 S +R l 1920 o - O 25 0 0 - +R l 1921 o - Ap 3 0 1 S +R l 1921 o - O 3 0 0 - +R l 1922 o - Mar 26 0 1 S +R l 1922 o - O 8 0 0 - +R l 1923 o - Ap 22 0 1 S +R l 1923 o - S 16 0 0 - +R l 1957 1961 - May 1 0 1 S +R l 1957 1961 - O 1 0 0 - +R l 1972 o - Jun 22 0 1 S +R l 1972 1977 - O 1 0 0 - +R l 1973 1977 - May 1 0 1 S +R l 1978 o - Ap 30 0 1 S +R l 1978 o - S 30 0 0 - +R l 1984 1987 - May 1 0 1 S +R l 1984 1991 - O 16 0 0 - +R l 1988 o - Jun 1 0 1 S +R l 1989 o - May 10 0 1 S +R l 1990 1992 - May 1 0 1 S +R l 1992 o - O 4 0 0 - +R l 1993 ma - Mar lastSu 0 1 S +R l 1993 1998 - S lastSu 0 0 - +R l 1999 ma - O lastSu 0 0 - +R NB 1935 1941 - S 14 0 0:20 - +R NB 1935 1941 - D 14 0 0 - +R X 1983 1984 - Ap 1 0 1 - +R X 1983 o - O 1 0 0 - +R X 1985 1998 - Mar lastSu 0 1 - +R X 1984 1998 - S lastSu 0 0 - +R X 2001 o - Ap lastSa 2 1 - +R X 2001 2006 - S lastSa 2 0 - +R X 2002 2006 - Mar lastSa 2 1 - +R X 2015 2016 - Mar lastSa 2 1 - +R X 2015 2016 - S lastSa 0 0 - +R PK 2002 o - Ap Su>=2 0 1 S +R PK 2002 o - O Su>=2 0 0 - +R PK 2008 o - Jun 1 0 1 S +R PK 2008 2009 - N 1 0 0 - +R PK 2009 o - Ap 15 0 1 S +R P 1999 2005 - Ap F>=15 0 1 S +R P 1999 2003 - O F>=15 0 0 - +R P 2004 o - O 1 1 0 - +R P 2005 o - O 4 2 0 - +R P 2006 2007 - Ap 1 0 1 S +R P 2006 o - S 22 0 0 - +R P 2007 o - S 13 2 0 - +R P 2008 2009 - Mar lastF 0 1 S +R P 2008 o - S 1 0 0 - +R P 2009 o - S 4 1 0 - +R P 2010 o - Mar 26 0 1 S +R P 2010 o - Au 11 0 0 - +R P 2011 o - Ap 1 0:1 1 S +R P 2011 o - Au 1 0 0 - +R P 2011 o - Au 30 0 1 S +R P 2011 o - S 30 0 0 - +R P 2012 2014 - Mar lastTh 24 1 S +R P 2012 o - S 21 1 0 - +R P 2013 o - S 27 0 0 - +R P 2014 o - O 24 0 0 - +R P 2015 o - Mar 28 0 1 S +R P 2015 o - O 23 1 0 - +R P 2016 2018 - Mar Sa<=30 1 1 S +R P 2016 2018 - O Sa<=30 1 0 - +R P 2019 o - Mar 29 0 1 S +R P 2019 o - O Sa<=30 0 0 - +R P 2020 2021 - Mar Sa<=30 0 1 S +R P 2020 o - O 24 1 0 - +R P 2021 o - O 29 1 0 - +R P 2022 o - Mar 27 0 1 S +R P 2022 2035 - O Sa<=30 2 0 - +R P 2023 o - Ap 29 2 1 S +R P 2024 o - Ap 20 2 1 S +R P 2025 o - Ap 12 2 1 S +R P 2026 2054 - Mar Sa<=30 2 1 S +R P 2036 o - O 18 2 0 - +R P 2037 o - O 10 2 0 - +R P 2038 o - S 25 2 0 - +R P 2039 o - S 17 2 0 - +R P 2040 o - S 1 2 0 - +R P 2040 o - O 20 2 1 S +R P 2040 2067 - O Sa<=30 2 0 - +R P 2041 o - Au 24 2 0 - +R P 2041 o - O 5 2 1 S +R P 2042 o - Au 16 2 0 - +R P 2042 o - S 27 2 1 S +R P 2043 o - Au 1 2 0 - +R P 2043 o - S 19 2 1 S +R P 2044 o - Jul 23 2 0 - +R P 2044 o - S 3 2 1 S +R P 2045 o - Jul 15 2 0 - +R P 2045 o - Au 26 2 1 S +R P 2046 o - Jun 30 2 0 - +R P 2046 o - Au 18 2 1 S +R P 2047 o - Jun 22 2 0 - +R P 2047 o - Au 3 2 1 S +R P 2048 o - Jun 6 2 0 - +R P 2048 o - Jul 25 2 1 S +R P 2049 o - May 29 2 0 - +R P 2049 o - Jul 10 2 1 S +R P 2050 o - May 21 2 0 - +R P 2050 o - Jul 2 2 1 S +R P 2051 o - May 6 2 0 - +R P 2051 o - Jun 24 2 1 S +R P 2052 o - Ap 27 2 0 - +R P 2052 o - Jun 8 2 1 S +R P 2053 o - Ap 12 2 0 - +R P 2053 o - May 31 2 1 S +R P 2054 o - Ap 4 2 0 - +R P 2054 o - May 23 2 1 S +R P 2055 o - May 8 2 1 S +R P 2056 o - Ap 29 2 1 S +R P 2057 o - Ap 14 2 1 S +R P 2058 o - Ap 6 2 1 S +R P 2059 ma - Mar Sa<=30 2 1 S +R P 2068 o - O 20 2 0 - +R P 2069 o - O 12 2 0 - +R P 2070 o - O 4 2 0 - +R P 2071 o - S 19 2 0 - +R P 2072 o - S 10 2 0 - +R P 2072 o - O 22 2 1 S +R P 2072 ma - O Sa<=30 2 0 - +R P 2073 o - S 2 2 0 - +R P 2073 o - O 14 2 1 S +R P 2074 o - Au 18 2 0 - +R P 2074 o - O 6 2 1 S +R P 2075 o - Au 10 2 0 - +R P 2075 o - S 21 2 1 S +R P 2076 o - Jul 25 2 0 - +R P 2076 o - S 12 2 1 S +R P 2077 o - Jul 17 2 0 - +R P 2077 o - S 4 2 1 S +R P 2078 o - Jul 9 2 0 - +R P 2078 o - Au 20 2 1 S +R P 2079 o - Jun 24 2 0 - +R P 2079 o - Au 12 2 1 S +R P 2080 o - Jun 15 2 0 - +R P 2080 o - Jul 27 2 1 S +R P 2081 o - Jun 7 2 0 - +R P 2081 o - Jul 19 2 1 S +R P 2082 o - May 23 2 0 - +R P 2082 o - Jul 11 2 1 S +R P 2083 o - May 15 2 0 - +R P 2083 o - Jun 26 2 1 S +R P 2084 o - Ap 29 2 0 - +R P 2084 o - Jun 17 2 1 S +R P 2085 o - Ap 21 2 0 - +R P 2085 o - Jun 9 2 1 S +R P 2086 o - Ap 13 2 0 - +R P 2086 o - May 25 2 1 S +R PH 1936 o - O 31 24 1 D +R PH 1937 o - Ja 15 24 0 S +R PH 1941 o - D 15 24 1 D +R PH 1945 o - N 30 24 0 S +R PH 1954 o - Ap 11 24 1 D +R PH 1954 o - Jun 4 24 0 S +R PH 1977 o - Mar 27 24 1 D +R PH 1977 o - S 21 24 0 S +R PH 1990 o - May 21 0 1 D +R PH 1990 o - Jul 28 24 0 S +R S 1920 1923 - Ap Su>=15 2 1 S +R S 1920 1923 - O Su>=1 2 0 - +R S 1962 o - Ap 29 2 1 S +R S 1962 o - O 1 2 0 - +R S 1963 1965 - May 1 2 1 S +R S 1963 o - S 30 2 0 - +R S 1964 o - O 1 2 0 - +R S 1965 o - S 30 2 0 - +R S 1966 o - Ap 24 2 1 S +R S 1966 1976 - O 1 2 0 - +R S 1967 1978 - May 1 2 1 S +R S 1977 1978 - S 1 2 0 - +R S 1983 1984 - Ap 9 2 1 S +R S 1983 1984 - O 1 2 0 - +R S 1986 o - F 16 2 1 S +R S 1986 o - O 9 2 0 - +R S 1987 o - Mar 1 2 1 S +R S 1987 1988 - O 31 2 0 - +R S 1988 o - Mar 15 2 1 S +R S 1989 o - Mar 31 2 1 S +R S 1989 o - O 1 2 0 - +R S 1990 o - Ap 1 2 1 S +R S 1990 o - S 30 2 0 - +R S 1991 o - Ap 1 0 1 S +R S 1991 1992 - O 1 0 0 - +R S 1992 o - Ap 8 0 1 S +R S 1993 o - Mar 26 0 1 S +R S 1993 o - S 25 0 0 - +R S 1994 1996 - Ap 1 0 1 S +R S 1994 2005 - O 1 0 0 - +R S 1997 1998 - Mar lastM 0 1 S +R S 1999 2006 - Ap 1 0 1 S +R S 2006 o - S 22 0 0 - +R S 2007 o - Mar lastF 0 1 S +R S 2007 o - N F>=1 0 0 - +R S 2008 o - Ap F>=1 0 1 S +R S 2008 o - N 1 0 0 - +R S 2009 o - Mar lastF 0 1 S +R S 2010 2011 - Ap F>=1 0 1 S +R S 2012 2022 - Mar lastF 0 1 S +R S 2009 2022 - O lastF 0 0 - +R AU 1917 o - Ja 1 2s 1 D +R AU 1917 o - Mar lastSu 2s 0 S +R AU 1942 o - Ja 1 2s 1 D +R AU 1942 o - Mar lastSu 2s 0 S +R AU 1942 o - S 27 2s 1 D +R AU 1943 1944 - Mar lastSu 2s 0 S +R AU 1943 o - O 3 2s 1 D +R AW 1974 o - O lastSu 2s 1 D +R AW 1975 o - Mar Su>=1 2s 0 S +R AW 1983 o - O lastSu 2s 1 D +R AW 1984 o - Mar Su>=1 2s 0 S +R AW 1991 o - N 17 2s 1 D +R AW 1992 o - Mar Su>=1 2s 0 S +R AW 2006 o - D 3 2s 1 D +R AW 2007 2009 - Mar lastSu 2s 0 S +R AW 2007 2008 - O lastSu 2s 1 D +R AQ 1971 o - O lastSu 2s 1 D +R AQ 1972 o - F lastSu 2s 0 S +R AQ 1989 1991 - O lastSu 2s 1 D +R AQ 1990 1992 - Mar Su>=1 2s 0 S +R Ho 1992 1993 - O lastSu 2s 1 D +R Ho 1993 1994 - Mar Su>=1 2s 0 S +R AS 1971 1985 - O lastSu 2s 1 D +R AS 1986 o - O 19 2s 1 D +R AS 1987 2007 - O lastSu 2s 1 D +R AS 1972 o - F 27 2s 0 S +R AS 1973 1985 - Mar Su>=1 2s 0 S +R AS 1986 1990 - Mar Su>=15 2s 0 S +R AS 1991 o - Mar 3 2s 0 S +R AS 1992 o - Mar 22 2s 0 S +R AS 1993 o - Mar 7 2s 0 S +R AS 1994 o - Mar 20 2s 0 S +R AS 1995 2005 - Mar lastSu 2s 0 S +R AS 2006 o - Ap 2 2s 0 S +R AS 2007 o - Mar lastSu 2s 0 S +R AS 2008 ma - Ap Su>=1 2s 0 S +R AS 2008 ma - O Su>=1 2s 1 D +R AT 1916 o - O Su>=1 2s 1 D +R AT 1917 o - Mar lastSu 2s 0 S +R AT 1917 1918 - O Su>=22 2s 1 D +R AT 1918 1919 - Mar Su>=1 2s 0 S +R AT 1967 o - O Su>=1 2s 1 D +R AT 1968 o - Mar Su>=29 2s 0 S +R AT 1968 1985 - O lastSu 2s 1 D +R AT 1969 1971 - Mar Su>=8 2s 0 S +R AT 1972 o - F lastSu 2s 0 S +R AT 1973 1981 - Mar Su>=1 2s 0 S +R AT 1982 1983 - Mar lastSu 2s 0 S +R AT 1984 1986 - Mar Su>=1 2s 0 S +R AT 1986 o - O Su>=15 2s 1 D +R AT 1987 1990 - Mar Su>=15 2s 0 S +R AT 1987 o - O Su>=22 2s 1 D +R AT 1988 1990 - O lastSu 2s 1 D +R AT 1991 1999 - O Su>=1 2s 1 D +R AT 1991 2005 - Mar lastSu 2s 0 S +R AT 2000 o - Au lastSu 2s 1 D +R AT 2001 ma - O Su>=1 2s 1 D +R AT 2006 o - Ap Su>=1 2s 0 S +R AT 2007 o - Mar lastSu 2s 0 S +R AT 2008 ma - Ap Su>=1 2s 0 S +R AV 1971 1985 - O lastSu 2s 1 D +R AV 1972 o - F lastSu 2s 0 S +R AV 1973 1985 - Mar Su>=1 2s 0 S +R AV 1986 1990 - Mar Su>=15 2s 0 S +R AV 1986 1987 - O Su>=15 2s 1 D +R AV 1988 1999 - O lastSu 2s 1 D +R AV 1991 1994 - Mar Su>=1 2s 0 S +R AV 1995 2005 - Mar lastSu 2s 0 S +R AV 2000 o - Au lastSu 2s 1 D +R AV 2001 2007 - O lastSu 2s 1 D +R AV 2006 o - Ap Su>=1 2s 0 S +R AV 2007 o - Mar lastSu 2s 0 S +R AV 2008 ma - Ap Su>=1 2s 0 S +R AV 2008 ma - O Su>=1 2s 1 D +R AN 1971 1985 - O lastSu 2s 1 D +R AN 1972 o - F 27 2s 0 S +R AN 1973 1981 - Mar Su>=1 2s 0 S +R AN 1982 o - Ap Su>=1 2s 0 S +R AN 1983 1985 - Mar Su>=1 2s 0 S +R AN 1986 1989 - Mar Su>=15 2s 0 S +R AN 1986 o - O 19 2s 1 D +R AN 1987 1999 - O lastSu 2s 1 D +R AN 1990 1995 - Mar Su>=1 2s 0 S +R AN 1996 2005 - Mar lastSu 2s 0 S +R AN 2000 o - Au lastSu 2s 1 D +R AN 2001 2007 - O lastSu 2s 1 D +R AN 2006 o - Ap Su>=1 2s 0 S +R AN 2007 o - Mar lastSu 2s 0 S +R AN 2008 ma - Ap Su>=1 2s 0 S +R AN 2008 ma - O Su>=1 2s 1 D +R LH 1981 1984 - O lastSu 2 1 - +R LH 1982 1985 - Mar Su>=1 2 0 - +R LH 1985 o - O lastSu 2 0:30 - +R LH 1986 1989 - Mar Su>=15 2 0 - +R LH 1986 o - O 19 2 0:30 - +R LH 1987 1999 - O lastSu 2 0:30 - +R LH 1990 1995 - Mar Su>=1 2 0 - +R LH 1996 2005 - Mar lastSu 2 0 - +R LH 2000 o - Au lastSu 2 0:30 - +R LH 2001 2007 - O lastSu 2 0:30 - +R LH 2006 o - Ap Su>=1 2 0 - +R LH 2007 o - Mar lastSu 2 0 - +R LH 2008 ma - Ap Su>=1 2 0 - +R LH 2008 ma - O Su>=1 2 0:30 - +R FJ 1998 1999 - N Su>=1 2 1 - +R FJ 1999 2000 - F lastSu 3 0 - +R FJ 2009 o - N 29 2 1 - +R FJ 2010 o - Mar lastSu 3 0 - +R FJ 2010 2013 - O Su>=21 2 1 - +R FJ 2011 o - Mar Su>=1 3 0 - +R FJ 2012 2013 - Ja Su>=18 3 0 - +R FJ 2014 o - Ja Su>=18 2 0 - +R FJ 2014 2018 - N Su>=1 2 1 - +R FJ 2015 2021 - Ja Su>=12 3 0 - +R FJ 2019 o - N Su>=8 2 1 - +R FJ 2020 o - D 20 2 1 - +R Gu 1959 o - Jun 27 2 1 D +R Gu 1961 o - Ja 29 2 0 S +R Gu 1967 o - S 1 2 1 D +R Gu 1969 o - Ja 26 0:1 0 S +R Gu 1969 o - Jun 22 2 1 D +R Gu 1969 o - Au 31 2 0 S +R Gu 1970 1971 - Ap lastSu 2 1 D +R Gu 1970 1971 - S Su>=1 2 0 S +R Gu 1973 o - D 16 2 1 D +R Gu 1974 o - F 24 2 0 S +R Gu 1976 o - May 26 2 1 D +R Gu 1976 o - Au 22 2:1 0 S +R Gu 1977 o - Ap 24 2 1 D +R Gu 1977 o - Au 28 2 0 S +R NC 1977 1978 - D Su>=1 0 1 - +R NC 1978 1979 - F 27 0 0 - +R NC 1996 o - D 1 2s 1 - +R NC 1997 o - Mar 2 2s 0 - +R NZ 1927 o - N 6 2 1 S +R NZ 1928 o - Mar 4 2 0 M +R NZ 1928 1933 - O Su>=8 2 0:30 S +R NZ 1929 1933 - Mar Su>=15 2 0 M +R NZ 1934 1940 - Ap lastSu 2 0 M +R NZ 1934 1940 - S lastSu 2 0:30 S +R NZ 1946 o - Ja 1 0 0 S +R NZ 1974 o - N Su>=1 2s 1 D +R k 1974 o - N Su>=1 2:45s 1 - +R NZ 1975 o - F lastSu 2s 0 S +R k 1975 o - F lastSu 2:45s 0 - +R NZ 1975 1988 - O lastSu 2s 1 D +R k 1975 1988 - O lastSu 2:45s 1 - +R NZ 1976 1989 - Mar Su>=1 2s 0 S +R k 1976 1989 - Mar Su>=1 2:45s 0 - +R NZ 1989 o - O Su>=8 2s 1 D +R k 1989 o - O Su>=8 2:45s 1 - +R NZ 1990 2006 - O Su>=1 2s 1 D +R k 1990 2006 - O Su>=1 2:45s 1 - +R NZ 1990 2007 - Mar Su>=15 2s 0 S +R k 1990 2007 - Mar Su>=15 2:45s 0 - +R NZ 2007 ma - S lastSu 2s 1 D +R k 2007 ma - S lastSu 2:45s 1 - +R NZ 2008 ma - Ap Su>=1 2s 0 S +R k 2008 ma - Ap Su>=1 2:45s 0 - +R CK 1978 o - N 12 0 0:30 - +R CK 1979 1991 - Mar Su>=1 0 0 - +R CK 1979 1990 - O lastSu 0 0:30 - +R WS 2010 o - S lastSu 0 1 - +R WS 2011 o - Ap Sa>=1 4 0 - +R WS 2011 o - S lastSa 3 1 - +R WS 2012 2021 - Ap Su>=1 4 0 - +R WS 2012 2020 - S lastSu 3 1 - +R TO 1999 o - O 7 2s 1 - +R TO 2000 o - Mar 19 2s 0 - +R TO 2000 2001 - N Su>=1 2 1 - +R TO 2001 2002 - Ja lastSu 2 0 - +R TO 2016 o - N Su>=1 2 1 - +R TO 2017 o - Ja Su>=15 3 0 - +R VU 1973 o - D 22 12u 1 - +R VU 1974 o - Mar 30 12u 0 - +R VU 1983 1991 - S Sa>=22 24 1 - +R VU 1984 1991 - Mar Sa>=22 24 0 - +R VU 1992 1993 - Ja Sa>=22 24 0 - +R VU 1992 o - O Sa>=22 24 1 - +R G 1916 o - May 21 2s 1 BST +R G 1916 o - O 1 2s 0 GMT +R G 1917 o - Ap 8 2s 1 BST +R G 1917 o - S 17 2s 0 GMT +R G 1918 o - Mar 24 2s 1 BST +R G 1918 o - S 30 2s 0 GMT +R G 1919 o - Mar 30 2s 1 BST +R G 1919 o - S 29 2s 0 GMT +R G 1920 o - Mar 28 2s 1 BST +R G 1920 o - O 25 2s 0 GMT +R G 1921 o - Ap 3 2s 1 BST +R G 1921 o - O 3 2s 0 GMT +R G 1922 o - Mar 26 2s 1 BST +R G 1922 o - O 8 2s 0 GMT +R G 1923 o - Ap Su>=16 2s 1 BST +R G 1923 1924 - S Su>=16 2s 0 GMT +R G 1924 o - Ap Su>=9 2s 1 BST +R G 1925 1926 - Ap Su>=16 2s 1 BST +R G 1925 1938 - O Su>=2 2s 0 GMT +R G 1927 o - Ap Su>=9 2s 1 BST +R G 1928 1929 - Ap Su>=16 2s 1 BST +R G 1930 o - Ap Su>=9 2s 1 BST +R G 1931 1932 - Ap Su>=16 2s 1 BST +R G 1933 o - Ap Su>=9 2s 1 BST +R G 1934 o - Ap Su>=16 2s 1 BST +R G 1935 o - Ap Su>=9 2s 1 BST +R G 1936 1937 - Ap Su>=16 2s 1 BST +R G 1938 o - Ap Su>=9 2s 1 BST +R G 1939 o - Ap Su>=16 2s 1 BST +R G 1939 o - N Su>=16 2s 0 GMT +R G 1940 o - F Su>=23 2s 1 BST +R G 1941 o - May Su>=2 1s 2 BDST +R G 1941 1943 - Au Su>=9 1s 1 BST +R G 1942 1944 - Ap Su>=2 1s 2 BDST +R G 1944 o - S Su>=16 1s 1 BST +R G 1945 o - Ap M>=2 1s 2 BDST +R G 1945 o - Jul Su>=9 1s 1 BST +R G 1945 1946 - O Su>=2 2s 0 GMT +R G 1946 o - Ap Su>=9 2s 1 BST +R G 1947 o - Mar 16 2s 1 BST +R G 1947 o - Ap 13 1s 2 BDST +R G 1947 o - Au 10 1s 1 BST +R G 1947 o - N 2 2s 0 GMT +R G 1948 o - Mar 14 2s 1 BST +R G 1948 o - O 31 2s 0 GMT +R G 1949 o - Ap 3 2s 1 BST +R G 1949 o - O 30 2s 0 GMT +R G 1950 1952 - Ap Su>=14 2s 1 BST +R G 1950 1952 - O Su>=21 2s 0 GMT +R G 1953 o - Ap Su>=16 2s 1 BST +R G 1953 1960 - O Su>=2 2s 0 GMT +R G 1954 o - Ap Su>=9 2s 1 BST +R G 1955 1956 - Ap Su>=16 2s 1 BST +R G 1957 o - Ap Su>=9 2s 1 BST +R G 1958 1959 - Ap Su>=16 2s 1 BST +R G 1960 o - Ap Su>=9 2s 1 BST +R G 1961 1963 - Mar lastSu 2s 1 BST +R G 1961 1968 - O Su>=23 2s 0 GMT +R G 1964 1967 - Mar Su>=19 2s 1 BST +R G 1968 o - F 18 2s 1 BST +R G 1972 1980 - Mar Su>=16 2s 1 BST +R G 1972 1980 - O Su>=23 2s 0 GMT +R G 1981 1995 - Mar lastSu 1u 1 BST +R G 1981 1989 - O Su>=23 1u 0 GMT +R G 1990 1995 - O Su>=22 1u 0 GMT +R IE 1971 o - O 31 2u -1 - +R IE 1972 1980 - Mar Su>=16 2u 0 - +R IE 1972 1980 - O Su>=23 2u -1 - +R IE 1981 ma - Mar lastSu 1u 0 - +R IE 1981 1989 - O Su>=23 1u -1 - +R IE 1990 1995 - O Su>=22 1u -1 - +R IE 1996 ma - O lastSu 1u -1 - +R E 1977 1980 - Ap Su>=1 1u 1 S +R E 1977 o - S lastSu 1u 0 - +R E 1978 o - O 1 1u 0 - +R E 1979 1995 - S lastSu 1u 0 - +R E 1981 ma - Mar lastSu 1u 1 S +R E 1996 ma - O lastSu 1u 0 - +R W- 1977 1980 - Ap Su>=1 1s 1 S +R W- 1977 o - S lastSu 1s 0 - +R W- 1978 o - O 1 1s 0 - +R W- 1979 1995 - S lastSu 1s 0 - +R W- 1981 ma - Mar lastSu 1s 1 S +R W- 1996 ma - O lastSu 1s 0 - +R c 1916 o - Ap 30 23 1 S +R c 1916 o - O 1 1 0 - +R c 1917 1918 - Ap M>=15 2s 1 S +R c 1917 1918 - S M>=15 2s 0 - +R c 1940 o - Ap 1 2s 1 S +R c 1942 o - N 2 2s 0 - +R c 1943 o - Mar 29 2s 1 S +R c 1943 o - O 4 2s 0 - +R c 1944 1945 - Ap M>=1 2s 1 S +R c 1944 o - O 2 2s 0 - +R c 1945 o - S 16 2s 0 - +R c 1977 1980 - Ap Su>=1 2s 1 S +R c 1977 o - S lastSu 2s 0 - +R c 1978 o - O 1 2s 0 - +R c 1979 1995 - S lastSu 2s 0 - +R c 1981 ma - Mar lastSu 2s 1 S +R c 1996 ma - O lastSu 2s 0 - +R e 1977 1980 - Ap Su>=1 0 1 S +R e 1977 o - S lastSu 0 0 - +R e 1978 o - O 1 0 0 - +R e 1979 1995 - S lastSu 0 0 - +R e 1981 ma - Mar lastSu 0 1 S +R e 1996 ma - O lastSu 0 0 - +R R 1917 o - Jul 1 23 1 MST +R R 1917 o - D 28 0 0 MMT +R R 1918 o - May 31 22 2 MDST +R R 1918 o - S 16 1 1 MST +R R 1919 o - May 31 23 2 MDST +R R 1919 o - Jul 1 0u 1 MSD +R R 1919 o - Au 16 0 0 MSK +R R 1921 o - F 14 23 1 MSD +R R 1921 o - Mar 20 23 2 +05 +R R 1921 o - S 1 0 1 MSD +R R 1921 o - O 1 0 0 - +R R 1981 1984 - Ap 1 0 1 S +R R 1981 1983 - O 1 0 0 - +R R 1984 1995 - S lastSu 2s 0 - +R R 1985 2010 - Mar lastSu 2s 1 S +R R 1996 2010 - O lastSu 2s 0 - +R q 1940 o - Jun 16 0 1 S +R q 1942 o - N 2 3 0 - +R q 1943 o - Mar 29 2 1 S +R q 1943 o - Ap 10 3 0 - +R q 1974 o - May 4 0 1 S +R q 1974 o - O 2 0 0 - +R q 1975 o - May 1 0 1 S +R q 1975 o - O 2 0 0 - +R q 1976 o - May 2 0 1 S +R q 1976 o - O 3 0 0 - +R q 1977 o - May 8 0 1 S +R q 1977 o - O 2 0 0 - +R q 1978 o - May 6 0 1 S +R q 1978 o - O 1 0 0 - +R q 1979 o - May 5 0 1 S +R q 1979 o - S 30 0 0 - +R q 1980 o - May 3 0 1 S +R q 1980 o - O 4 0 0 - +R q 1981 o - Ap 26 0 1 S +R q 1981 o - S 27 0 0 - +R q 1982 o - May 2 0 1 S +R q 1982 o - O 3 0 0 - +R q 1983 o - Ap 18 0 1 S +R q 1983 o - O 1 0 0 - +R q 1984 o - Ap 1 0 1 S +R a 1920 o - Ap 5 2s 1 S +R a 1920 o - S 13 2s 0 - +R a 1946 o - Ap 14 2s 1 S +R a 1946 o - O 7 2s 0 - +R a 1947 1948 - O Su>=1 2s 0 - +R a 1947 o - Ap 6 2s 1 S +R a 1948 o - Ap 18 2s 1 S +R a 1980 o - Ap 6 0 1 S +R a 1980 o - S 28 0 0 - +R b 1918 o - Mar 9 0s 1 S +R b 1918 1919 - O Sa>=1 23s 0 - +R b 1919 o - Mar 1 23s 1 S +R b 1920 o - F 14 23s 1 S +R b 1920 o - O 23 23s 0 - +R b 1921 o - Mar 14 23s 1 S +R b 1921 o - O 25 23s 0 - +R b 1922 o - Mar 25 23s 1 S +R b 1922 1927 - O Sa>=1 23s 0 - +R b 1923 o - Ap 21 23s 1 S +R b 1924 o - Mar 29 23s 1 S +R b 1925 o - Ap 4 23s 1 S +R b 1926 o - Ap 17 23s 1 S +R b 1927 o - Ap 9 23s 1 S +R b 1928 o - Ap 14 23s 1 S +R b 1928 1938 - O Su>=2 2s 0 - +R b 1929 o - Ap 21 2s 1 S +R b 1930 o - Ap 13 2s 1 S +R b 1931 o - Ap 19 2s 1 S +R b 1932 o - Ap 3 2s 1 S +R b 1933 o - Mar 26 2s 1 S +R b 1934 o - Ap 8 2s 1 S +R b 1935 o - Mar 31 2s 1 S +R b 1936 o - Ap 19 2s 1 S +R b 1937 o - Ap 4 2s 1 S +R b 1938 o - Mar 27 2s 1 S +R b 1939 o - Ap 16 2s 1 S +R b 1939 o - N 19 2s 0 - +R b 1940 o - F 25 2s 1 S +R b 1944 o - S 17 2s 0 - +R b 1945 o - Ap 2 2s 1 S +R b 1945 o - S 16 2s 0 - +R b 1946 o - May 19 2s 1 S +R b 1946 o - O 7 2s 0 - +R BG 1979 o - Mar 31 23 1 S +R BG 1979 o - O 1 1 0 - +R BG 1980 1982 - Ap Sa>=1 23 1 S +R BG 1980 o - S 29 1 0 - +R BG 1981 o - S 27 2 0 - +R CZ 1945 o - Ap M>=1 2s 1 S +R CZ 1945 o - O 1 2s 0 - +R CZ 1946 o - May 6 2s 1 S +R CZ 1946 1949 - O Su>=1 2s 0 - +R CZ 1947 1948 - Ap Su>=15 2s 1 S +R CZ 1949 o - Ap 9 2s 1 S +R Th 1991 1992 - Mar lastSu 2 1 D +R Th 1991 1992 - S lastSu 2 0 S +R Th 1993 2006 - Ap Su>=1 2 1 D +R Th 1993 2006 - O lastSu 2 0 S +R Th 2007 ma - Mar Su>=8 2 1 D +R Th 2007 ma - N Su>=1 2 0 S +R FI 1942 o - Ap 2 24 1 S +R FI 1942 o - O 4 1 0 - +R FI 1981 1982 - Mar lastSu 2 1 S +R FI 1981 1982 - S lastSu 3 0 - +R F 1916 o - Jun 14 23s 1 S +R F 1916 1919 - O Su>=1 23s 0 - +R F 1917 o - Mar 24 23s 1 S +R F 1918 o - Mar 9 23s 1 S +R F 1919 o - Mar 1 23s 1 S +R F 1920 o - F 14 23s 1 S +R F 1920 o - O 23 23s 0 - +R F 1921 o - Mar 14 23s 1 S +R F 1921 o - O 25 23s 0 - +R F 1922 o - Mar 25 23s 1 S +R F 1922 1938 - O Sa>=1 23s 0 - +R F 1923 o - May 26 23s 1 S +R F 1924 o - Mar 29 23s 1 S +R F 1925 o - Ap 4 23s 1 S +R F 1926 o - Ap 17 23s 1 S +R F 1927 o - Ap 9 23s 1 S +R F 1928 o - Ap 14 23s 1 S +R F 1929 o - Ap 20 23s 1 S +R F 1930 o - Ap 12 23s 1 S +R F 1931 o - Ap 18 23s 1 S +R F 1932 o - Ap 2 23s 1 S +R F 1933 o - Mar 25 23s 1 S +R F 1934 o - Ap 7 23s 1 S +R F 1935 o - Mar 30 23s 1 S +R F 1936 o - Ap 18 23s 1 S +R F 1937 o - Ap 3 23s 1 S +R F 1938 o - Mar 26 23s 1 S +R F 1939 o - Ap 15 23s 1 S +R F 1939 o - N 18 23s 0 - +R F 1940 o - F 25 2 1 S +R F 1941 o - May 5 0 2 M +R F 1941 o - O 6 0 1 S +R F 1942 o - Mar 9 0 2 M +R F 1942 o - N 2 3 1 S +R F 1943 o - Mar 29 2 2 M +R F 1943 o - O 4 3 1 S +R F 1944 o - Ap 3 2 2 M +R F 1944 o - O 8 1 1 S +R F 1945 o - Ap 2 2 2 M +R F 1945 o - S 16 3 0 - +R F 1976 o - Mar 28 1 1 S +R F 1976 o - S 26 1 0 - +R DE 1946 o - Ap 14 2s 1 S +R DE 1946 o - O 7 2s 0 - +R DE 1947 1949 - O Su>=1 2s 0 - +R DE 1947 o - Ap 6 3s 1 S +R DE 1947 o - May 11 2s 2 M +R DE 1947 o - Jun 29 3 1 S +R DE 1948 o - Ap 18 2s 1 S +R DE 1949 o - Ap 10 2s 1 S +R So 1945 o - May 24 2 2 M +R So 1945 o - S 24 3 1 S +R So 1945 o - N 18 2s 0 - +R g 1932 o - Jul 7 0 1 S +R g 1932 o - S 1 0 0 - +R g 1941 o - Ap 7 0 1 S +R g 1942 o - N 2 3 0 - +R g 1943 o - Mar 30 0 1 S +R g 1943 o - O 4 0 0 - +R g 1952 o - Jul 1 0 1 S +R g 1952 o - N 2 0 0 - +R g 1975 o - Ap 12 0s 1 S +R g 1975 o - N 26 0s 0 - +R g 1976 o - Ap 11 2s 1 S +R g 1976 o - O 10 2s 0 - +R g 1977 1978 - Ap Su>=1 2s 1 S +R g 1977 o - S 26 2s 0 - +R g 1978 o - S 24 4 0 - +R g 1979 o - Ap 1 9 1 S +R g 1979 o - S 29 2 0 - +R g 1980 o - Ap 1 0 1 S +R g 1980 o - S 28 0 0 - +R h 1918 1919 - Ap 15 2 1 S +R h 1918 1920 - S M>=15 3 0 - +R h 1920 o - Ap 5 2 1 S +R h 1945 o - May 1 23 1 S +R h 1945 o - N 1 1 0 - +R h 1946 o - Mar 31 2s 1 S +R h 1946 o - O 7 2 0 - +R h 1947 1949 - Ap Su>=4 2s 1 S +R h 1947 1949 - O Su>=1 2s 0 - +R h 1954 o - May 23 0 1 S +R h 1954 o - O 3 0 0 - +R h 1955 o - May 22 2 1 S +R h 1955 o - O 2 3 0 - +R h 1956 1957 - Jun Su>=1 2 1 S +R h 1956 1957 - S lastSu 3 0 - +R h 1980 o - Ap 6 0 1 S +R h 1980 o - S 28 1 0 - +R h 1981 1983 - Mar lastSu 0 1 S +R h 1981 1983 - S lastSu 1 0 - +R I 1916 o - Jun 3 24 1 S +R I 1916 1917 - S 30 24 0 - +R I 1917 o - Mar 31 24 1 S +R I 1918 o - Mar 9 24 1 S +R I 1918 o - O 6 24 0 - +R I 1919 o - Mar 1 24 1 S +R I 1919 o - O 4 24 0 - +R I 1920 o - Mar 20 24 1 S +R I 1920 o - S 18 24 0 - +R I 1940 o - Jun 14 24 1 S +R I 1942 o - N 2 2s 0 - +R I 1943 o - Mar 29 2s 1 S +R I 1943 o - O 4 2s 0 - +R I 1944 o - Ap 2 2s 1 S +R I 1944 o - S 17 2s 0 - +R I 1945 o - Ap 2 2 1 S +R I 1945 o - S 15 1 0 - +R I 1946 o - Mar 17 2s 1 S +R I 1946 o - O 6 2s 0 - +R I 1947 o - Mar 16 0s 1 S +R I 1947 o - O 5 0s 0 - +R I 1948 o - F 29 2s 1 S +R I 1948 o - O 3 2s 0 - +R I 1966 1968 - May Su>=22 0s 1 S +R I 1966 o - S 24 24 0 - +R I 1967 1969 - S Su>=22 0s 0 - +R I 1969 o - Jun 1 0s 1 S +R I 1970 o - May 31 0s 1 S +R I 1970 o - S lastSu 0s 0 - +R I 1971 1972 - May Su>=22 0s 1 S +R I 1971 o - S lastSu 0s 0 - +R I 1972 o - O 1 0s 0 - +R I 1973 o - Jun 3 0s 1 S +R I 1973 1974 - S lastSu 0s 0 - +R I 1974 o - May 26 0s 1 S +R I 1975 o - Jun 1 0s 1 S +R I 1975 1977 - S lastSu 0s 0 - +R I 1976 o - May 30 0s 1 S +R I 1977 1979 - May Su>=22 0s 1 S +R I 1978 o - O 1 0s 0 - +R I 1979 o - S 30 0s 0 - +R LV 1989 1996 - Mar lastSu 2s 1 S +R LV 1989 1996 - S lastSu 2s 0 - +R MT 1973 o - Mar 31 0s 1 S +R MT 1973 o - S 29 0s 0 - +R MT 1974 o - Ap 21 0s 1 S +R MT 1974 o - S 16 0s 0 - +R MT 1975 1979 - Ap Su>=15 2 1 S +R MT 1975 1980 - S Su>=15 2 0 - +R MT 1980 o - Mar 31 2 1 S +R MD 1997 ma - Mar lastSu 2 1 S +R MD 1997 ma - O lastSu 3 0 - +R O 1918 1919 - S 16 2s 0 - +R O 1919 o - Ap 15 2s 1 S +R O 1944 o - Ap 3 2s 1 S +R O 1944 o - O 4 2 0 - +R O 1945 o - Ap 29 0 1 S +R O 1945 o - N 1 0 0 - +R O 1946 o - Ap 14 0s 1 S +R O 1946 o - O 7 2s 0 - +R O 1947 o - May 4 2s 1 S +R O 1947 1949 - O Su>=1 2s 0 - +R O 1948 o - Ap 18 2s 1 S +R O 1949 o - Ap 10 2s 1 S +R O 1957 o - Jun 2 1s 1 S +R O 1957 1958 - S lastSu 1s 0 - +R O 1958 o - Mar 30 1s 1 S +R O 1959 o - May 31 1s 1 S +R O 1959 1961 - O Su>=1 1s 0 - +R O 1960 o - Ap 3 1s 1 S +R O 1961 1964 - May lastSu 1s 1 S +R O 1962 1964 - S lastSu 1s 0 - +R p 1916 o - Jun 17 23 1 S +R p 1916 o - N 1 1 0 - +R p 1917 1921 - Mar 1 0 1 S +R p 1917 1921 - O 14 24 0 - +R p 1924 o - Ap 16 23s 1 S +R p 1924 o - O 4 23s 0 - +R p 1926 o - Ap 17 23s 1 S +R p 1926 1929 - O Sa>=1 23s 0 - +R p 1927 o - Ap 9 23s 1 S +R p 1928 o - Ap 14 23s 1 S +R p 1929 o - Ap 20 23s 1 S +R p 1931 o - Ap 18 23s 1 S +R p 1931 1932 - O Sa>=1 23s 0 - +R p 1932 o - Ap 2 23s 1 S +R p 1934 o - Ap 7 23s 1 S +R p 1934 1938 - O Sa>=1 23s 0 - +R p 1935 o - Mar 30 23s 1 S +R p 1936 o - Ap 18 23s 1 S +R p 1937 o - Ap 3 23s 1 S +R p 1938 o - Mar 26 23s 1 S +R p 1939 o - Ap 15 23s 1 S +R p 1939 o - N 18 23s 0 - +R p 1940 o - F 24 23s 1 S +R p 1940 o - O 7 23s 0 - +R p 1941 o - Ap 5 23s 1 S +R p 1941 o - O 5 23s 0 - +R p 1942 1945 - Mar Sa>=8 23s 1 S +R p 1942 o - Ap 25 22s 2 M +R p 1942 o - Au 15 22s 1 S +R p 1942 1945 - O Sa>=24 23s 0 - +R p 1943 o - Ap 17 22s 2 M +R p 1943 1945 - Au Sa>=25 22s 1 S +R p 1944 1945 - Ap Sa>=21 22s 2 M +R p 1946 o - Ap Sa>=1 23s 1 S +R p 1946 o - O Sa>=1 23s 0 - +R p 1947 1966 - Ap Su>=1 2s 1 S +R p 1947 1965 - O Su>=1 2s 0 - +R p 1976 o - S lastSu 1 0 - +R p 1977 o - Mar lastSu 0s 1 S +R p 1977 o - S lastSu 0s 0 - +R p 1978 1980 - Ap Su>=1 1s 1 S +R p 1978 o - O 1 1s 0 - +R p 1979 1980 - S lastSu 1s 0 - +R p 1981 1986 - Mar lastSu 0s 1 S +R p 1981 1985 - S lastSu 0s 0 - +R z 1932 o - May 21 0s 1 S +R z 1932 1939 - O Su>=1 0s 0 - +R z 1933 1939 - Ap Su>=2 0s 1 S +R z 1979 o - May 27 0 1 S +R z 1979 o - S lastSu 0 0 - +R z 1980 o - Ap 5 23 1 S +R z 1980 o - S lastSu 1 0 - +R z 1991 1993 - Mar lastSu 0s 1 S +R z 1991 1993 - S lastSu 0s 0 - +R s 1918 o - Ap 15 23 1 S +R s 1918 1919 - O 6 24s 0 - +R s 1919 o - Ap 6 23 1 S +R s 1924 o - Ap 16 23 1 S +R s 1924 o - O 4 24s 0 - +R s 1926 o - Ap 17 23 1 S +R s 1926 1929 - O Sa>=1 24s 0 - +R s 1927 o - Ap 9 23 1 S +R s 1928 o - Ap 15 0 1 S +R s 1929 o - Ap 20 23 1 S +R s 1937 o - Jun 16 23 1 S +R s 1937 o - O 2 24s 0 - +R s 1938 o - Ap 2 23 1 S +R s 1938 o - Ap 30 23 2 M +R s 1938 o - O 2 24 1 S +R s 1939 o - O 7 24s 0 - +R s 1942 o - May 2 23 1 S +R s 1942 o - S 1 1 0 - +R s 1943 1946 - Ap Sa>=13 23 1 S +R s 1943 1944 - O Su>=1 1 0 - +R s 1945 1946 - S lastSu 1 0 - +R s 1949 o - Ap 30 23 1 S +R s 1949 o - O 2 1 0 - +R s 1974 1975 - Ap Sa>=12 23 1 S +R s 1974 1975 - O Su>=1 1 0 - +R s 1976 o - Mar 27 23 1 S +R s 1976 1977 - S lastSu 1 0 - +R s 1977 o - Ap 2 23 1 S +R s 1978 o - Ap 2 2s 1 S +R s 1978 o - O 1 2s 0 - +R Sp 1967 o - Jun 3 12 1 S +R Sp 1967 o - O 1 0 0 - +R Sp 1974 o - Jun 24 0 1 S +R Sp 1974 o - S 1 0 0 - +R Sp 1976 1977 - May 1 0 1 S +R Sp 1976 o - Au 1 0 0 - +R Sp 1977 o - S 28 0 0 - +R Sp 1978 o - Jun 1 0 1 S +R Sp 1978 o - Au 4 0 0 - +R CH 1941 1942 - May M>=1 1 1 S +R CH 1941 1942 - O M>=1 2 0 - +R T 1916 o - May 1 0 1 S +R T 1916 o - O 1 0 0 - +R T 1920 o - Mar 28 0 1 S +R T 1920 o - O 25 0 0 - +R T 1921 o - Ap 3 0 1 S +R T 1921 o - O 3 0 0 - +R T 1922 o - Mar 26 0 1 S +R T 1922 o - O 8 0 0 - +R T 1924 o - May 13 0 1 S +R T 1924 1925 - O 1 0 0 - +R T 1925 o - May 1 0 1 S +R T 1940 o - Jul 1 0 1 S +R T 1940 o - O 6 0 0 - +R T 1940 o - D 1 0 1 S +R T 1941 o - S 21 0 0 - +R T 1942 o - Ap 1 0 1 S +R T 1945 o - O 8 0 0 - +R T 1946 o - Jun 1 0 1 S +R T 1946 o - O 1 0 0 - +R T 1947 1948 - Ap Su>=16 0 1 S +R T 1947 1951 - O Su>=2 0 0 - +R T 1949 o - Ap 10 0 1 S +R T 1950 o - Ap 16 0 1 S +R T 1951 o - Ap 22 0 1 S +R T 1962 o - Jul 15 0 1 S +R T 1963 o - O 30 0 0 - +R T 1964 o - May 15 0 1 S +R T 1964 o - O 1 0 0 - +R T 1973 o - Jun 3 1 1 S +R T 1973 1976 - O Su>=31 2 0 - +R T 1974 o - Mar 31 2 1 S +R T 1975 o - Mar 22 2 1 S +R T 1976 o - Mar 21 2 1 S +R T 1977 1978 - Ap Su>=1 2 1 S +R T 1977 1978 - O Su>=15 2 0 - +R T 1978 o - Jun 29 0 0 - +R T 1983 o - Jul 31 2 1 S +R T 1983 o - O 2 2 0 - +R T 1985 o - Ap 20 1s 1 S +R T 1985 o - S 28 1s 0 - +R T 1986 1993 - Mar lastSu 1s 1 S +R T 1986 1995 - S lastSu 1s 0 - +R T 1994 o - Mar 20 1s 1 S +R T 1995 2006 - Mar lastSu 1s 1 S +R T 1996 2006 - O lastSu 1s 0 - +R u 1918 1919 - Mar lastSu 2 1 D +R u 1918 1919 - O lastSu 2 0 S +R u 1942 o - F 9 2 1 W +R u 1945 o - Au 14 23u 1 P +R u 1945 o - S 30 2 0 S +R u 1967 2006 - O lastSu 2 0 S +R u 1967 1973 - Ap lastSu 2 1 D +R u 1974 o - Ja 6 2 1 D +R u 1975 o - F lastSu 2 1 D +R u 1976 1986 - Ap lastSu 2 1 D +R u 1987 2006 - Ap Su>=1 2 1 D +R u 2007 ma - Mar Su>=8 2 1 D +R u 2007 ma - N Su>=1 2 0 S +R NY 1920 o - Mar lastSu 2 1 D +R NY 1920 o - O lastSu 2 0 S +R NY 1921 1966 - Ap lastSu 2 1 D +R NY 1921 1954 - S lastSu 2 0 S +R NY 1955 1966 - O lastSu 2 0 S +R Ch 1920 o - Jun 13 2 1 D +R Ch 1920 1921 - O lastSu 2 0 S +R Ch 1921 o - Mar lastSu 2 1 D +R Ch 1922 1966 - Ap lastSu 2 1 D +R Ch 1922 1954 - S lastSu 2 0 S +R Ch 1955 1966 - O lastSu 2 0 S +R De 1920 1921 - Mar lastSu 2 1 D +R De 1920 o - O lastSu 2 0 S +R De 1921 o - May 22 2 0 S +R De 1965 1966 - Ap lastSu 2 1 D +R De 1965 1966 - O lastSu 2 0 S +R CA 1948 o - Mar 14 2:1 1 D +R CA 1949 o - Ja 1 2 0 S +R CA 1950 1966 - Ap lastSu 1 1 D +R CA 1950 1961 - S lastSu 2 0 S +R CA 1962 1966 - O lastSu 2 0 S +R In 1941 o - Jun 22 2 1 D +R In 1941 1954 - S lastSu 2 0 S +R In 1946 1954 - Ap lastSu 2 1 D +R Ma 1951 o - Ap lastSu 2 1 D +R Ma 1951 o - S lastSu 2 0 S +R Ma 1954 1960 - Ap lastSu 2 1 D +R Ma 1954 1960 - S lastSu 2 0 S +R V 1946 o - Ap lastSu 2 1 D +R V 1946 o - S lastSu 2 0 S +R V 1953 1954 - Ap lastSu 2 1 D +R V 1953 1959 - S lastSu 2 0 S +R V 1955 o - May 1 0 1 D +R V 1956 1963 - Ap lastSu 2 1 D +R V 1960 o - O lastSu 2 0 S +R V 1961 o - S lastSu 2 0 S +R V 1962 1963 - O lastSu 2 0 S +R Pe 1955 o - May 1 0 1 D +R Pe 1955 1960 - S lastSu 2 0 S +R Pe 1956 1963 - Ap lastSu 2 1 D +R Pe 1961 1963 - O lastSu 2 0 S +R Pi 1955 o - May 1 0 1 D +R Pi 1955 1960 - S lastSu 2 0 S +R Pi 1956 1964 - Ap lastSu 2 1 D +R Pi 1961 1964 - O lastSu 2 0 S +R St 1947 1961 - Ap lastSu 2 1 D +R St 1947 1954 - S lastSu 2 0 S +R St 1955 1956 - O lastSu 2 0 S +R St 1957 1958 - S lastSu 2 0 S +R St 1959 1961 - O lastSu 2 0 S +R Pu 1946 1960 - Ap lastSu 2 1 D +R Pu 1946 1954 - S lastSu 2 0 S +R Pu 1955 1956 - O lastSu 2 0 S +R Pu 1957 1960 - S lastSu 2 0 S +R v 1921 o - May 1 2 1 D +R v 1921 o - S 1 2 0 S +R v 1941 o - Ap lastSu 2 1 D +R v 1941 o - S lastSu 2 0 S +R v 1946 o - Ap lastSu 0:1 1 D +R v 1946 o - Jun 2 2 0 S +R v 1950 1961 - Ap lastSu 2 1 D +R v 1950 1955 - S lastSu 2 0 S +R v 1956 1961 - O lastSu 2 0 S +R Dt 1948 o - Ap lastSu 2 1 D +R Dt 1948 o - S lastSu 2 0 S +R Me 1946 o - Ap lastSu 2 1 D +R Me 1946 o - S lastSu 2 0 S +R Me 1966 o - Ap lastSu 2 1 D +R Me 1966 o - O lastSu 2 0 S +R C 1918 o - Ap 14 2 1 D +R C 1918 o - O 27 2 0 S +R C 1942 o - F 9 2 1 W +R C 1945 o - Au 14 23u 1 P +R C 1945 o - S 30 2 0 S +R C 1974 1986 - Ap lastSu 2 1 D +R C 1974 2006 - O lastSu 2 0 S +R C 1987 2006 - Ap Su>=1 2 1 D +R C 2007 ma - Mar Su>=8 2 1 D +R C 2007 ma - N Su>=1 2 0 S +R j 1917 o - Ap 8 2 1 D +R j 1917 o - S 17 2 0 S +R j 1919 o - May 5 23 1 D +R j 1919 o - Au 12 23 0 S +R j 1920 1935 - May Su>=1 23 1 D +R j 1920 1935 - O lastSu 23 0 S +R j 1936 1941 - May M>=9 0 1 D +R j 1936 1941 - O M>=2 0 0 S +R j 1946 1950 - May Su>=8 2 1 D +R j 1946 1950 - O Su>=2 2 0 S +R j 1951 1986 - Ap lastSu 2 1 D +R j 1951 1959 - S lastSu 2 0 S +R j 1960 1986 - O lastSu 2 0 S +R j 1987 o - Ap Su>=1 0:1 1 D +R j 1987 2006 - O lastSu 0:1 0 S +R j 1988 o - Ap Su>=1 0:1 2 DD +R j 1989 2006 - Ap Su>=1 0:1 1 D +R j 2007 2011 - Mar Su>=8 0:1 1 D +R j 2007 2010 - N Su>=1 0:1 0 S +R H 1916 o - Ap 1 0 1 D +R H 1916 o - O 1 0 0 S +R H 1920 o - May 9 0 1 D +R H 1920 o - Au 29 0 0 S +R H 1921 o - May 6 0 1 D +R H 1921 1922 - S 5 0 0 S +R H 1922 o - Ap 30 0 1 D +R H 1923 1925 - May Su>=1 0 1 D +R H 1923 o - S 4 0 0 S +R H 1924 o - S 15 0 0 S +R H 1925 o - S 28 0 0 S +R H 1926 o - May 16 0 1 D +R H 1926 o - S 13 0 0 S +R H 1927 o - May 1 0 1 D +R H 1927 o - S 26 0 0 S +R H 1928 1931 - May Su>=8 0 1 D +R H 1928 o - S 9 0 0 S +R H 1929 o - S 3 0 0 S +R H 1930 o - S 15 0 0 S +R H 1931 1932 - S M>=24 0 0 S +R H 1932 o - May 1 0 1 D +R H 1933 o - Ap 30 0 1 D +R H 1933 o - O 2 0 0 S +R H 1934 o - May 20 0 1 D +R H 1934 o - S 16 0 0 S +R H 1935 o - Jun 2 0 1 D +R H 1935 o - S 30 0 0 S +R H 1936 o - Jun 1 0 1 D +R H 1936 o - S 14 0 0 S +R H 1937 1938 - May Su>=1 0 1 D +R H 1937 1941 - S M>=24 0 0 S +R H 1939 o - May 28 0 1 D +R H 1940 1941 - May Su>=1 0 1 D +R H 1946 1949 - Ap lastSu 2 1 D +R H 1946 1949 - S lastSu 2 0 S +R H 1951 1954 - Ap lastSu 2 1 D +R H 1951 1954 - S lastSu 2 0 S +R H 1956 1959 - Ap lastSu 2 1 D +R H 1956 1959 - S lastSu 2 0 S +R H 1962 1973 - Ap lastSu 2 1 D +R H 1962 1973 - O lastSu 2 0 S +R o 1933 1935 - Jun Su>=8 1 1 D +R o 1933 1935 - S Su>=8 1 0 S +R o 1936 1938 - Jun Su>=1 1 1 D +R o 1936 1938 - S Su>=1 1 0 S +R o 1939 o - May 27 1 1 D +R o 1939 1941 - S Sa>=21 1 0 S +R o 1940 o - May 19 1 1 D +R o 1941 o - May 4 1 1 D +R o 1946 1972 - Ap lastSu 2 1 D +R o 1946 1956 - S lastSu 2 0 S +R o 1957 1972 - O lastSu 2 0 S +R o 1993 2006 - Ap Su>=1 0:1 1 D +R o 1993 2006 - O lastSu 0:1 0 S +R t 1919 o - Mar 30 23:30 1 D +R t 1919 o - O 26 0 0 S +R t 1920 o - May 2 2 1 D +R t 1920 o - S 26 0 0 S +R t 1921 o - May 15 2 1 D +R t 1921 o - S 15 2 0 S +R t 1922 1923 - May Su>=8 2 1 D +R t 1922 1926 - S Su>=15 2 0 S +R t 1924 1927 - May Su>=1 2 1 D +R t 1927 1937 - S Su>=25 2 0 S +R t 1928 1937 - Ap Su>=25 2 1 D +R t 1938 1940 - Ap lastSu 2 1 D +R t 1938 1939 - S lastSu 2 0 S +R t 1945 1948 - S lastSu 2 0 S +R t 1946 1973 - Ap lastSu 2 1 D +R t 1949 1950 - N lastSu 2 0 S +R t 1951 1956 - S lastSu 2 0 S +R t 1957 1973 - O lastSu 2 0 S +R W 1916 o - Ap 23 0 1 D +R W 1916 o - S 17 0 0 S +R W 1918 o - Ap 14 2 1 D +R W 1918 o - O 27 2 0 S +R W 1937 o - May 16 2 1 D +R W 1937 o - S 26 2 0 S +R W 1942 o - F 9 2 1 W +R W 1945 o - Au 14 23u 1 P +R W 1945 o - S lastSu 2 0 S +R W 1946 o - May 12 2 1 D +R W 1946 o - O 13 2 0 S +R W 1947 1949 - Ap lastSu 2 1 D +R W 1947 1949 - S lastSu 2 0 S +R W 1950 o - May 1 2 1 D +R W 1950 o - S 30 2 0 S +R W 1951 1960 - Ap lastSu 2 1 D +R W 1951 1958 - S lastSu 2 0 S +R W 1959 o - O lastSu 2 0 S +R W 1960 o - S lastSu 2 0 S +R W 1963 o - Ap lastSu 2 1 D +R W 1963 o - S 22 2 0 S +R W 1966 1986 - Ap lastSu 2s 1 D +R W 1966 2005 - O lastSu 2s 0 S +R W 1987 2005 - Ap Su>=1 2s 1 D +R r 1918 o - Ap 14 2 1 D +R r 1918 o - O 27 2 0 S +R r 1930 1934 - May Su>=1 0 1 D +R r 1930 1934 - O Su>=1 0 0 S +R r 1937 1941 - Ap Su>=8 0 1 D +R r 1937 o - O Su>=8 0 0 S +R r 1938 o - O Su>=1 0 0 S +R r 1939 1941 - O Su>=8 0 0 S +R r 1942 o - F 9 2 1 W +R r 1945 o - Au 14 23u 1 P +R r 1945 o - S lastSu 2 0 S +R r 1946 o - Ap Su>=8 2 1 D +R r 1946 o - O Su>=8 2 0 S +R r 1947 1957 - Ap lastSu 2 1 D +R r 1947 1957 - S lastSu 2 0 S +R r 1959 o - Ap lastSu 2 1 D +R r 1959 o - O lastSu 2 0 S +R Sw 1957 o - Ap lastSu 2 1 D +R Sw 1957 o - O lastSu 2 0 S +R Sw 1959 1961 - Ap lastSu 2 1 D +R Sw 1959 o - O lastSu 2 0 S +R Sw 1960 1961 - S lastSu 2 0 S +R Ed 1918 1919 - Ap Su>=8 2 1 D +R Ed 1918 o - O 27 2 0 S +R Ed 1919 o - May 27 2 0 S +R Ed 1920 1923 - Ap lastSu 2 1 D +R Ed 1920 o - O lastSu 2 0 S +R Ed 1921 1923 - S lastSu 2 0 S +R Ed 1942 o - F 9 2 1 W +R Ed 1945 o - Au 14 23u 1 P +R Ed 1945 o - S lastSu 2 0 S +R Ed 1947 o - Ap lastSu 2 1 D +R Ed 1947 o - S lastSu 2 0 S +R Ed 1972 1986 - Ap lastSu 2 1 D +R Ed 1972 2006 - O lastSu 2 0 S +R Va 1918 o - Ap 14 2 1 D +R Va 1918 o - O 27 2 0 S +R Va 1942 o - F 9 2 1 W +R Va 1945 o - Au 14 23u 1 P +R Va 1945 o - S 30 2 0 S +R Va 1946 1986 - Ap lastSu 2 1 D +R Va 1946 o - S 29 2 0 S +R Va 1947 1961 - S lastSu 2 0 S +R Va 1962 2006 - O lastSu 2 0 S +R Y 1918 o - Ap 14 2 1 D +R Y 1918 o - O 27 2 0 S +R Y 1919 o - May 25 2 1 D +R Y 1919 o - N 1 0 0 S +R Y 1942 o - F 9 2 1 W +R Y 1945 o - Au 14 23u 1 P +R Y 1945 o - S 30 2 0 S +R Y 1972 1986 - Ap lastSu 2 1 D +R Y 1972 2006 - O lastSu 2 0 S +R Y 1987 2006 - Ap Su>=1 2 1 D +R Yu 1965 o - Ap lastSu 0 2 DD +R Yu 1965 o - O lastSu 2 0 S +R m 1931 o - Ap 30 0 1 D +R m 1931 o - O 1 0 0 S +R m 1939 o - F 5 0 1 D +R m 1939 o - Jun 25 0 0 S +R m 1940 o - D 9 0 1 D +R m 1941 o - Ap 1 0 0 S +R m 1943 o - D 16 0 1 W +R m 1944 o - May 1 0 0 S +R m 1950 o - F 12 0 1 D +R m 1950 o - Jul 30 0 0 S +R m 1996 2000 - Ap Su>=1 2 1 D +R m 1996 2000 - O lastSu 2 0 S +R m 2001 o - May Su>=1 2 1 D +R m 2001 o - S lastSu 2 0 S +R m 2002 2022 - Ap Su>=1 2 1 D +R m 2002 2022 - O lastSu 2 0 S +R BB 1942 o - Ap 19 5u 1 D +R BB 1942 o - Au 31 6u 0 S +R BB 1943 o - May 2 5u 1 D +R BB 1943 o - S 5 6u 0 S +R BB 1944 o - Ap 10 5u 0:30 - +R BB 1944 o - S 10 6u 0 S +R BB 1977 o - Jun 12 2 1 D +R BB 1977 1978 - O Su>=1 2 0 S +R BB 1978 1980 - Ap Su>=15 2 1 D +R BB 1979 o - S 30 2 0 S +R BB 1980 o - S 25 2 0 S +R BZ 1918 1941 - O Sa>=1 24 0:30 -0530 +R BZ 1919 1942 - F Sa>=8 24 0 CST +R BZ 1942 o - Jun 27 24 1 CWT +R BZ 1945 o - Au 14 23u 1 CPT +R BZ 1945 o - D 15 24 0 CST +R BZ 1947 1967 - O Sa>=1 24 0:30 -0530 +R BZ 1948 1968 - F Sa>=8 24 0 CST +R BZ 1973 o - D 5 0 1 CDT +R BZ 1974 o - F 9 0 0 CST +R BZ 1982 o - D 18 0 1 CDT +R BZ 1983 o - F 12 0 0 CST +R Be 1917 o - Ap 5 24 1 - +R Be 1917 o - S 30 24 0 - +R Be 1918 o - Ap 13 24 1 - +R Be 1918 o - S 15 24 0 S +R Be 1942 o - Ja 11 2 1 D +R Be 1942 o - O 18 2 0 S +R Be 1943 o - Mar 21 2 1 D +R Be 1943 o - O 31 2 0 S +R Be 1944 1945 - Mar Su>=8 2 1 D +R Be 1944 1945 - N Su>=1 2 0 S +R Be 1947 o - May Su>=15 2 1 D +R Be 1947 o - S Su>=8 2 0 S +R Be 1948 1952 - May Su>=22 2 1 D +R Be 1948 1952 - S Su>=1 2 0 S +R Be 1956 o - May Su>=22 2 1 D +R Be 1956 o - O lastSu 2 0 S +R CR 1979 1980 - F lastSu 0 1 D +R CR 1979 1980 - Jun Su>=1 0 0 S +R CR 1991 1992 - Ja Sa>=15 0 1 D +R CR 1991 o - Jul 1 0 0 S +R CR 1992 o - Mar 15 0 0 S +R Q 1928 o - Jun 10 0 1 D +R Q 1928 o - O 10 0 0 S +R Q 1940 1942 - Jun Su>=1 0 1 D +R Q 1940 1942 - S Su>=1 0 0 S +R Q 1945 1946 - Jun Su>=1 0 1 D +R Q 1945 1946 - S Su>=1 0 0 S +R Q 1965 o - Jun 1 0 1 D +R Q 1965 o - S 30 0 0 S +R Q 1966 o - May 29 0 1 D +R Q 1966 o - O 2 0 0 S +R Q 1967 o - Ap 8 0 1 D +R Q 1967 1968 - S Su>=8 0 0 S +R Q 1968 o - Ap 14 0 1 D +R Q 1969 1977 - Ap lastSu 0 1 D +R Q 1969 1971 - O lastSu 0 0 S +R Q 1972 1974 - O 8 0 0 S +R Q 1975 1977 - O lastSu 0 0 S +R Q 1978 o - May 7 0 1 D +R Q 1978 1990 - O Su>=8 0 0 S +R Q 1979 1980 - Mar Su>=15 0 1 D +R Q 1981 1985 - May Su>=5 0 1 D +R Q 1986 1989 - Mar Su>=14 0 1 D +R Q 1990 1997 - Ap Su>=1 0 1 D +R Q 1991 1995 - O Su>=8 0s 0 S +R Q 1996 o - O 6 0s 0 S +R Q 1997 o - O 12 0s 0 S +R Q 1998 1999 - Mar lastSu 0s 1 D +R Q 1998 2003 - O lastSu 0s 0 S +R Q 2000 2003 - Ap Su>=1 0s 1 D +R Q 2004 o - Mar lastSu 0s 1 D +R Q 2006 2010 - O lastSu 0s 0 S +R Q 2007 o - Mar Su>=8 0s 1 D +R Q 2008 o - Mar Su>=15 0s 1 D +R Q 2009 2010 - Mar Su>=8 0s 1 D +R Q 2011 o - Mar Su>=15 0s 1 D +R Q 2011 o - N 13 0s 0 S +R Q 2012 o - Ap 1 0s 1 D +R Q 2012 ma - N Su>=1 0s 0 S +R Q 2013 ma - Mar Su>=8 0s 1 D +R DO 1966 o - O 30 0 1 EDT +R DO 1967 o - F 28 0 0 EST +R DO 1969 1973 - O lastSu 0 0:30 -0430 +R DO 1970 o - F 21 0 0 EST +R DO 1971 o - Ja 20 0 0 EST +R DO 1972 1974 - Ja 21 0 0 EST +R SV 1987 1988 - May Su>=1 0 1 D +R SV 1987 1988 - S lastSu 0 0 S +R GT 1973 o - N 25 0 1 D +R GT 1974 o - F 24 0 0 S +R GT 1983 o - May 21 0 1 D +R GT 1983 o - S 22 0 0 S +R GT 1991 o - Mar 23 0 1 D +R GT 1991 o - S 7 0 0 S +R GT 2006 o - Ap 30 0 1 D +R GT 2006 o - O 1 0 0 S +R HT 1983 o - May 8 0 1 D +R HT 1984 1987 - Ap lastSu 0 1 D +R HT 1983 1987 - O lastSu 0 0 S +R HT 1988 1997 - Ap Su>=1 1s 1 D +R HT 1988 1997 - O lastSu 1s 0 S +R HT 2005 2006 - Ap Su>=1 0 1 D +R HT 2005 2006 - O lastSu 0 0 S +R HT 2012 2015 - Mar Su>=8 2 1 D +R HT 2012 2015 - N Su>=1 2 0 S +R HT 2017 ma - Mar Su>=8 2 1 D +R HT 2017 ma - N Su>=1 2 0 S +R HN 1987 1988 - May Su>=1 0 1 D +R HN 1987 1988 - S lastSu 0 0 S +R HN 2006 o - May Su>=1 0 1 D +R HN 2006 o - Au M>=1 0 0 S +R NI 1979 1980 - Mar Su>=16 0 1 D +R NI 1979 1980 - Jun M>=23 0 0 S +R NI 2005 o - Ap 10 0 1 D +R NI 2005 o - O Su>=1 0 0 S +R NI 2006 o - Ap 30 2 1 D +R NI 2006 o - O Su>=1 1 0 S +R A 1930 o - D 1 0 1 - +R A 1931 o - Ap 1 0 0 - +R A 1931 o - O 15 0 1 - +R A 1932 1940 - Mar 1 0 0 - +R A 1932 1939 - N 1 0 1 - +R A 1940 o - Jul 1 0 1 - +R A 1941 o - Jun 15 0 0 - +R A 1941 o - O 15 0 1 - +R A 1943 o - Au 1 0 0 - +R A 1943 o - O 15 0 1 - +R A 1946 o - Mar 1 0 0 - +R A 1946 o - O 1 0 1 - +R A 1963 o - O 1 0 0 - +R A 1963 o - D 15 0 1 - +R A 1964 1966 - Mar 1 0 0 - +R A 1964 1966 - O 15 0 1 - +R A 1967 o - Ap 2 0 0 - +R A 1967 1968 - O Su>=1 0 1 - +R A 1968 1969 - Ap Su>=1 0 0 - +R A 1974 o - Ja 23 0 1 - +R A 1974 o - May 1 0 0 - +R A 1988 o - D 1 0 1 - +R A 1989 1993 - Mar Su>=1 0 0 - +R A 1989 1992 - O Su>=15 0 1 - +R A 1999 o - O Su>=1 0 1 - +R A 2000 o - Mar 3 0 0 - +R A 2007 o - D 30 0 1 - +R A 2008 2009 - Mar Su>=15 0 0 - +R A 2008 o - O Su>=15 0 1 - +R Sa 2008 2009 - Mar Su>=8 0 0 - +R Sa 2007 2008 - O Su>=8 0 1 - +R B 1931 o - O 3 11 1 - +R B 1932 1933 - Ap 1 0 0 - +R B 1932 o - O 3 0 1 - +R B 1949 1952 - D 1 0 1 - +R B 1950 o - Ap 16 1 0 - +R B 1951 1952 - Ap 1 0 0 - +R B 1953 o - Mar 1 0 0 - +R B 1963 o - D 9 0 1 - +R B 1964 o - Mar 1 0 0 - +R B 1965 o - Ja 31 0 1 - +R B 1965 o - Mar 31 0 0 - +R B 1965 o - D 1 0 1 - +R B 1966 1968 - Mar 1 0 0 - +R B 1966 1967 - N 1 0 1 - +R B 1985 o - N 2 0 1 - +R B 1986 o - Mar 15 0 0 - +R B 1986 o - O 25 0 1 - +R B 1987 o - F 14 0 0 - +R B 1987 o - O 25 0 1 - +R B 1988 o - F 7 0 0 - +R B 1988 o - O 16 0 1 - +R B 1989 o - Ja 29 0 0 - +R B 1989 o - O 15 0 1 - +R B 1990 o - F 11 0 0 - +R B 1990 o - O 21 0 1 - +R B 1991 o - F 17 0 0 - +R B 1991 o - O 20 0 1 - +R B 1992 o - F 9 0 0 - +R B 1992 o - O 25 0 1 - +R B 1993 o - Ja 31 0 0 - +R B 1993 1995 - O Su>=11 0 1 - +R B 1994 1995 - F Su>=15 0 0 - +R B 1996 o - F 11 0 0 - +R B 1996 o - O 6 0 1 - +R B 1997 o - F 16 0 0 - +R B 1997 o - O 6 0 1 - +R B 1998 o - Mar 1 0 0 - +R B 1998 o - O 11 0 1 - +R B 1999 o - F 21 0 0 - +R B 1999 o - O 3 0 1 - +R B 2000 o - F 27 0 0 - +R B 2000 2001 - O Su>=8 0 1 - +R B 2001 2006 - F Su>=15 0 0 - +R B 2002 o - N 3 0 1 - +R B 2003 o - O 19 0 1 - +R B 2004 o - N 2 0 1 - +R B 2005 o - O 16 0 1 - +R B 2006 o - N 5 0 1 - +R B 2007 o - F 25 0 0 - +R B 2007 o - O Su>=8 0 1 - +R B 2008 2017 - O Su>=15 0 1 - +R B 2008 2011 - F Su>=15 0 0 - +R B 2012 o - F Su>=22 0 0 - +R B 2013 2014 - F Su>=15 0 0 - +R B 2015 o - F Su>=22 0 0 - +R B 2016 2019 - F Su>=15 0 0 - +R B 2018 o - N Su>=1 0 1 - +R x 1927 1931 - S 1 0 1 - +R x 1928 1932 - Ap 1 0 0 - +R x 1968 o - N 3 4u 1 - +R x 1969 o - Mar 30 3u 0 - +R x 1969 o - N 23 4u 1 - +R x 1970 o - Mar 29 3u 0 - +R x 1971 o - Mar 14 3u 0 - +R x 1970 1972 - O Su>=9 4u 1 - +R x 1972 1986 - Mar Su>=9 3u 0 - +R x 1973 o - S 30 4u 1 - +R x 1974 1987 - O Su>=9 4u 1 - +R x 1987 o - Ap 12 3u 0 - +R x 1988 1990 - Mar Su>=9 3u 0 - +R x 1988 1989 - O Su>=9 4u 1 - +R x 1990 o - S 16 4u 1 - +R x 1991 1996 - Mar Su>=9 3u 0 - +R x 1991 1997 - O Su>=9 4u 1 - +R x 1997 o - Mar 30 3u 0 - +R x 1998 o - Mar Su>=9 3u 0 - +R x 1998 o - S 27 4u 1 - +R x 1999 o - Ap 4 3u 0 - +R x 1999 2010 - O Su>=9 4u 1 - +R x 2000 2007 - Mar Su>=9 3u 0 - +R x 2008 o - Mar 30 3u 0 - +R x 2009 o - Mar Su>=9 3u 0 - +R x 2010 o - Ap Su>=1 3u 0 - +R x 2011 o - May Su>=2 3u 0 - +R x 2011 o - Au Su>=16 4u 1 - +R x 2012 2014 - Ap Su>=23 3u 0 - +R x 2012 2014 - S Su>=2 4u 1 - +R x 2016 2018 - May Su>=9 3u 0 - +R x 2016 2018 - Au Su>=9 4u 1 - +R x 2019 ma - Ap Su>=2 3u 0 - +R x 2019 2021 - S Su>=2 4u 1 - +R x 2022 o - S Su>=9 4u 1 - +R x 2023 ma - S Su>=2 4u 1 - +R CO 1992 o - May 3 0 1 - +R CO 1993 o - F 6 24 0 - +R EC 1992 o - N 28 0 1 - +R EC 1993 o - F 5 0 0 - +R FK 1937 1938 - S lastSu 0 1 - +R FK 1938 1942 - Mar Su>=19 0 0 - +R FK 1939 o - O 1 0 1 - +R FK 1940 1942 - S lastSu 0 1 - +R FK 1943 o - Ja 1 0 0 - +R FK 1983 o - S lastSu 0 1 - +R FK 1984 1985 - Ap lastSu 0 0 - +R FK 1984 o - S 16 0 1 - +R FK 1985 2000 - S Su>=9 0 1 - +R FK 1986 2000 - Ap Su>=16 0 0 - +R FK 2001 2010 - Ap Su>=15 2 0 - +R FK 2001 2010 - S Su>=1 2 1 - +R y 1975 1988 - O 1 0 1 - +R y 1975 1978 - Mar 1 0 0 - +R y 1979 1991 - Ap 1 0 0 - +R y 1989 o - O 22 0 1 - +R y 1990 o - O 1 0 1 - +R y 1991 o - O 6 0 1 - +R y 1992 o - Mar 1 0 0 - +R y 1992 o - O 5 0 1 - +R y 1993 o - Mar 31 0 0 - +R y 1993 1995 - O 1 0 1 - +R y 1994 1995 - F lastSu 0 0 - +R y 1996 o - Mar 1 0 0 - +R y 1996 2001 - O Su>=1 0 1 - +R y 1997 o - F lastSu 0 0 - +R y 1998 2001 - Mar Su>=1 0 0 - +R y 2002 2004 - Ap Su>=1 0 0 - +R y 2002 2003 - S Su>=1 0 1 - +R y 2004 2009 - O Su>=15 0 1 - +R y 2005 2009 - Mar Su>=8 0 0 - +R y 2010 2024 - O Su>=1 0 1 - +R y 2010 2012 - Ap Su>=8 0 0 - +R y 2013 2024 - Mar Su>=22 0 0 - +R PE 1938 o - Ja 1 0 1 - +R PE 1938 o - Ap 1 0 0 - +R PE 1938 1939 - S lastSu 0 1 - +R PE 1939 1940 - Mar Su>=24 0 0 - +R PE 1986 1987 - Ja 1 0 1 - +R PE 1986 1987 - Ap 1 0 0 - +R PE 1990 o - Ja 1 0 1 - +R PE 1990 o - Ap 1 0 0 - +R PE 1994 o - Ja 1 0 1 - +R PE 1994 o - Ap 1 0 0 - +R U 1923 1925 - O 1 0 0:30 - +R U 1924 1926 - Ap 1 0 0 - +R U 1933 1938 - O lastSu 0 0:30 - +R U 1934 1941 - Mar lastSa 24 0 - +R U 1939 o - O 1 0 0:30 - +R U 1940 o - O 27 0 0:30 - +R U 1941 o - Au 1 0 0:30 - +R U 1942 o - D 14 0 0:30 - +R U 1943 o - Mar 14 0 0 - +R U 1959 o - May 24 0 0:30 - +R U 1959 o - N 15 0 0 - +R U 1960 o - Ja 17 0 1 - +R U 1960 o - Mar 6 0 0 - +R U 1965 o - Ap 4 0 1 - +R U 1965 o - S 26 0 0 - +R U 1968 o - May 27 0 0:30 - +R U 1968 o - D 1 0 0 - +R U 1970 o - Ap 25 0 1 - +R U 1970 o - Jun 14 0 0 - +R U 1972 o - Ap 23 0 1 - +R U 1972 o - Jul 16 0 0 - +R U 1974 o - Ja 13 0 1:30 - +R U 1974 o - Mar 10 0 0:30 - +R U 1974 o - S 1 0 0 - +R U 1974 o - D 22 0 1 - +R U 1975 o - Mar 30 0 0 - +R U 1976 o - D 19 0 1 - +R U 1977 o - Mar 6 0 0 - +R U 1977 o - D 4 0 1 - +R U 1978 1979 - Mar Su>=1 0 0 - +R U 1978 o - D 17 0 1 - +R U 1979 o - Ap 29 0 1 - +R U 1980 o - Mar 16 0 0 - +R U 1987 o - D 14 0 1 - +R U 1988 o - F 28 0 0 - +R U 1988 o - D 11 0 1 - +R U 1989 o - Mar 5 0 0 - +R U 1989 o - O 29 0 1 - +R U 1990 o - F 25 0 0 - +R U 1990 1991 - O Su>=21 0 1 - +R U 1991 1992 - Mar Su>=1 0 0 - +R U 1992 o - O 18 0 1 - +R U 1993 o - F 28 0 0 - +R U 2004 o - S 19 0 1 - +R U 2005 o - Mar 27 2 0 - +R U 2005 o - O 9 2 1 - +R U 2006 2015 - Mar Su>=8 2 0 - +R U 2006 2014 - O Su>=1 2 1 - +Z Africa/Abidjan -0:16:8 - LMT 1912 +0 - GMT +Z Africa/Algiers 0:12:12 - LMT 1891 Mar 16 +0:9:21 - PMT 1911 Mar 11 +0 d WE%sT 1940 F 25 2 +1 d CE%sT 1946 O 7 +0 - WET 1956 Ja 29 +1 - CET 1963 Ap 14 +0 d WE%sT 1977 O 21 +1 d CE%sT 1979 O 26 +0 d WE%sT 1981 May +1 - CET +Z Africa/Bissau -1:2:20 - LMT 1912 Ja 1 1u +-1 - %z 1975 +0 - GMT +Z Africa/Cairo 2:5:9 - LMT 1900 O +2 K EE%sT +Z Africa/Casablanca -0:30:20 - LMT 1913 O 26 +0 M %z 1984 Mar 16 +1 - %z 1986 +0 M %z 2018 O 28 3 +1 M %z +Z Africa/Ceuta -0:21:16 - LMT 1901 Ja 1 0u +0 - WET 1918 May 6 23 +0 1 WEST 1918 O 7 23 +0 - WET 1924 +0 s WE%sT 1929 +0 - WET 1967 +0 Sp WE%sT 1984 Mar 16 +1 - CET 1986 +1 E CE%sT +Z Africa/El_Aaiun -0:52:48 - LMT 1934 +-1 - %z 1976 Ap 14 +0 M %z 2018 O 28 3 +1 M %z +Z Africa/Johannesburg 1:52 - LMT 1892 F 8 +1:30 - SAST 1903 Mar +2 SA SAST +Z Africa/Juba 2:6:28 - LMT 1931 +2 SD CA%sT 2000 Ja 15 12 +3 - EAT 2021 F +2 - CAT +Z Africa/Khartoum 2:10:8 - LMT 1931 +2 SD CA%sT 2000 Ja 15 12 +3 - EAT 2017 N +2 - CAT +Z Africa/Lagos 0:13:35 - LMT 1905 Jul +0 - GMT 1908 Jul +0:13:35 - LMT 1914 +0:30 - %z 1919 S +1 - WAT +Z Africa/Maputo 2:10:18 - LMT 1909 +2 - CAT +Z Africa/Monrovia -0:43:8 - LMT 1882 +-0:43:8 - MMT 1919 Mar +-0:44:30 - MMT 1972 Ja 7 +0 - GMT +Z Africa/Nairobi 2:27:16 - LMT 1908 May +2:30 - %z 1928 Jun 30 24 +3 - EAT 1930 Ja 4 24 +2:30 - %z 1936 D 31 24 +2:45 - %z 1942 Jul 31 24 +3 - EAT +Z Africa/Ndjamena 1:0:12 - LMT 1912 +1 - WAT 1979 O 14 +1 1 WAST 1980 Mar 8 +1 - WAT +Z Africa/Sao_Tome 0:26:56 - LMT 1884 +-0:36:45 - LMT 1912 Ja 1 0u +0 - GMT 2018 Ja 1 1 +1 - WAT 2019 Ja 1 2 +0 - GMT +Z Africa/Tripoli 0:52:44 - LMT 1920 +1 L CE%sT 1959 +2 - EET 1982 +1 L CE%sT 1990 May 4 +2 - EET 1996 S 30 +1 L CE%sT 1997 O 4 +2 - EET 2012 N 10 2 +1 L CE%sT 2013 O 25 2 +2 - EET +Z Africa/Tunis 0:40:44 - LMT 1881 May 12 +0:9:21 - PMT 1911 Mar 11 +1 n CE%sT +Z Africa/Windhoek 1:8:24 - LMT 1892 F 8 +1:30 - %z 1903 Mar +2 - SAST 1942 S 20 2 +2 1 SAST 1943 Mar 21 2 +2 - SAST 1990 Mar 21 +2 NA %s +Z America/Adak 12:13:22 - LMT 1867 O 19 12:44:35 +-11:46:38 - LMT 1900 Au 20 12 +-11 - NST 1942 +-11 u N%sT 1946 +-11 - NST 1967 Ap +-11 - BST 1969 +-11 u B%sT 1983 O 30 2 +-10 u AH%sT 1983 N 30 +-10 u H%sT +Z America/Anchorage 14:0:24 - LMT 1867 O 19 14:31:37 +-9:59:36 - LMT 1900 Au 20 12 +-10 - AST 1942 +-10 u A%sT 1967 Ap +-10 - AHST 1969 +-10 u AH%sT 1983 O 30 2 +-9 u Y%sT 1983 N 30 +-9 u AK%sT +Z America/Araguaina -3:12:48 - LMT 1914 +-3 B %z 1990 S 17 +-3 - %z 1995 S 14 +-3 B %z 2003 S 24 +-3 - %z 2012 O 21 +-3 B %z 2013 S +-3 - %z +Z America/Argentina/Buenos_Aires -3:53:48 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 A %z +Z America/Argentina/Catamarca -4:23:8 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1991 Mar 3 +-4 - %z 1991 O 20 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 - %z 2004 Jun +-4 - %z 2004 Jun 20 +-3 A %z 2008 O 18 +-3 - %z +Z America/Argentina/Cordoba -4:16:48 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1991 Mar 3 +-4 - %z 1991 O 20 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 A %z +Z America/Argentina/Jujuy -4:21:12 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1990 Mar 4 +-4 - %z 1990 O 28 +-4 1 %z 1991 Mar 17 +-4 - %z 1991 O 6 +-3 1 %z 1992 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 A %z 2008 O 18 +-3 - %z +Z America/Argentina/La_Rioja -4:27:24 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1991 Mar +-4 - %z 1991 May 7 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 - %z 2004 Jun +-4 - %z 2004 Jun 20 +-3 A %z 2008 O 18 +-3 - %z +Z America/Argentina/Mendoza -4:35:16 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1990 Mar 4 +-4 - %z 1990 O 15 +-4 1 %z 1991 Mar +-4 - %z 1991 O 15 +-4 1 %z 1992 Mar +-4 - %z 1992 O 18 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 - %z 2004 May 23 +-4 - %z 2004 S 26 +-3 A %z 2008 O 18 +-3 - %z +Z America/Argentina/Rio_Gallegos -4:36:52 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 - %z 2004 Jun +-4 - %z 2004 Jun 20 +-3 A %z 2008 O 18 +-3 - %z +Z America/Argentina/Salta -4:21:40 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1991 Mar 3 +-4 - %z 1991 O 20 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 A %z 2008 O 18 +-3 - %z +Z America/Argentina/San_Juan -4:34:4 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1991 Mar +-4 - %z 1991 May 7 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 - %z 2004 May 31 +-4 - %z 2004 Jul 25 +-3 A %z 2008 O 18 +-3 - %z +Z America/Argentina/San_Luis -4:25:24 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1990 +-3 1 %z 1990 Mar 14 +-4 - %z 1990 O 15 +-4 1 %z 1991 Mar +-4 - %z 1991 Jun +-3 - %z 1999 O 3 +-4 1 %z 2000 Mar 3 +-3 - %z 2004 May 31 +-4 - %z 2004 Jul 25 +-3 A %z 2008 Ja 21 +-4 Sa %z 2009 O 11 +-3 - %z +Z America/Argentina/Tucuman -4:20:52 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1991 Mar 3 +-4 - %z 1991 O 20 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 - %z 2004 Jun +-4 - %z 2004 Jun 13 +-3 A %z +Z America/Argentina/Ushuaia -4:33:12 - LMT 1894 O 31 +-4:16:48 - CMT 1920 May +-4 - %z 1930 D +-4 A %z 1969 O 5 +-3 A %z 1999 O 3 +-4 A %z 2000 Mar 3 +-3 - %z 2004 May 30 +-4 - %z 2004 Jun 20 +-3 A %z 2008 O 18 +-3 - %z +Z America/Asuncion -3:50:40 - LMT 1890 +-3:50:40 - AMT 1931 O 10 +-4 - %z 1972 O +-3 - %z 1974 Ap +-4 y %z 2024 O 15 +-3 - %z +Z America/Bahia -2:34:4 - LMT 1914 +-3 B %z 2003 S 24 +-3 - %z 2011 O 16 +-3 B %z 2012 O 21 +-3 - %z +Z America/Bahia_Banderas -7:1 - LMT 1922 Ja 1 7u +-7 - MST 1927 Jun 10 +-6 - CST 1930 N 15 +-7 m M%sT 1932 Ap +-6 - CST 1942 Ap 24 +-7 - MST 1970 +-7 m M%sT 2010 Ap 4 2 +-6 m C%sT +Z America/Barbados -3:58:29 - LMT 1911 Au 28 +-4 BB A%sT 1944 +-4 BB AST/-0330 1945 +-4 BB A%sT +Z America/Belem -3:13:56 - LMT 1914 +-3 B %z 1988 S 12 +-3 - %z +Z America/Belize -5:52:48 - LMT 1912 Ap +-6 BZ %s +Z America/Boa_Vista -4:2:40 - LMT 1914 +-4 B %z 1988 S 12 +-4 - %z 1999 S 30 +-4 B %z 2000 O 15 +-4 - %z +Z America/Bogota -4:56:16 - LMT 1884 Mar 13 +-4:56:16 - BMT 1914 N 23 +-5 CO %z +Z America/Boise -7:44:49 - LMT 1883 N 18 20u +-8 u P%sT 1923 May 13 2 +-7 u M%sT 1974 +-7 - MST 1974 F 3 2 +-7 u M%sT +Z America/Cambridge_Bay 0 - -00 1920 +-7 Y M%sT 1999 O 31 2 +-6 C C%sT 2000 O 29 2 +-5 - EST 2000 N 5 +-6 - CST 2001 Ap 1 3 +-7 C M%sT +Z America/Campo_Grande -3:38:28 - LMT 1914 +-4 B %z +Z America/Cancun -5:47:4 - LMT 1922 Ja 1 6u +-6 - CST 1981 D 26 2 +-5 - EST 1983 Ja 4 +-6 m C%sT 1997 O 26 2 +-5 m E%sT 1998 Au 2 2 +-6 m C%sT 2015 F 1 2 +-5 - EST +Z America/Caracas -4:27:44 - LMT 1890 +-4:27:40 - CMT 1912 F 12 +-4:30 - %z 1965 +-4 - %z 2007 D 9 3 +-4:30 - %z 2016 May 1 2:30 +-4 - %z +Z America/Cayenne -3:29:20 - LMT 1911 Jul +-4 - %z 1967 O +-3 - %z +Z America/Chicago -5:50:36 - LMT 1883 N 18 18u +-6 u C%sT 1920 +-6 Ch C%sT 1936 Mar 1 2 +-5 - EST 1936 N 15 2 +-6 Ch C%sT 1942 +-6 u C%sT 1946 +-6 Ch C%sT 1967 +-6 u C%sT +Z America/Chihuahua -7:4:20 - LMT 1922 Ja 1 7u +-7 - MST 1927 Jun 10 +-6 - CST 1930 N 15 +-7 m M%sT 1932 Ap +-6 - CST 1996 +-6 m C%sT 1998 +-6 - CST 1998 Ap Su>=1 3 +-7 m M%sT 2022 O 30 2 +-6 - CST +Z America/Ciudad_Juarez -7:5:56 - LMT 1922 Ja 1 7u +-7 - MST 1927 Jun 10 +-6 - CST 1930 N 15 +-7 m M%sT 1932 Ap +-6 - CST 1996 +-6 m C%sT 1998 +-6 - CST 1998 Ap Su>=1 3 +-7 m M%sT 2010 +-7 u M%sT 2022 O 30 2 +-6 - CST 2022 N 30 +-7 u M%sT +Z America/Costa_Rica -5:36:13 - LMT 1890 +-5:36:13 - SJMT 1921 Ja 15 +-6 CR C%sT +Z America/Coyhaique -4:48:16 - LMT 1890 +-4:42:45 - SMT 1910 Ja 10 +-5 - %z 1916 Jul +-4:42:45 - SMT 1918 S 10 +-4 - %z 1919 Jul +-4:42:45 - SMT 1927 S +-5 x %z 1932 S +-4 - %z 1942 Jun +-5 - %z 1942 Au +-4 - %z 1946 Au 28 24 +-5 1 %z 1947 Mar 31 24 +-5 - %z 1947 May 21 23 +-4 x %z 2025 Mar 20 +-3 - %z +Z America/Cuiaba -3:44:20 - LMT 1914 +-4 B %z 2003 S 24 +-4 - %z 2004 O +-4 B %z +Z America/Danmarkshavn -1:14:40 - LMT 1916 Jul 28 +-3 - %z 1980 Ap 6 2 +-3 E %z 1996 +0 - GMT +Z America/Dawson -9:17:40 - LMT 1900 Au 20 +-9 Y Y%sT 1965 +-9 Yu Y%sT 1973 O 28 +-8 - PST 1980 +-8 C P%sT 2020 N +-7 - MST +Z America/Dawson_Creek -8:0:56 - LMT 1884 +-8 C P%sT 1947 +-8 Va P%sT 1972 Au 30 2 +-7 - MST +Z America/Denver -6:59:56 - LMT 1883 N 18 19u +-7 u M%sT 1920 +-7 De M%sT 1942 +-7 u M%sT 1946 +-7 De M%sT 1967 +-7 u M%sT +Z America/Detroit -5:32:11 - LMT 1905 +-6 - CST 1915 May 15 2 +-5 - EST 1942 +-5 u E%sT 1946 +-5 Dt E%sT 1967 Jun 14 0:1 +-5 u E%sT 1969 +-5 - EST 1973 +-5 u E%sT 1975 +-5 - EST 1975 Ap 27 2 +-5 u E%sT +Z America/Edmonton -7:33:52 - LMT 1906 S +-7 Ed M%sT 1987 +-7 C M%sT +Z America/Eirunepe -4:39:28 - LMT 1914 +-5 B %z 1988 S 12 +-5 - %z 1993 S 28 +-5 B %z 1994 S 22 +-5 - %z 2008 Jun 24 +-4 - %z 2013 N 10 +-5 - %z +Z America/El_Salvador -5:56:48 - LMT 1921 +-6 SV C%sT +Z America/Fort_Nelson -8:10:47 - LMT 1884 +-8 Va P%sT 1946 +-8 - PST 1947 +-8 Va P%sT 1987 +-8 C P%sT 2015 Mar 8 2 +-7 - MST +Z America/Fortaleza -2:34 - LMT 1914 +-3 B %z 1990 S 17 +-3 - %z 1999 S 30 +-3 B %z 2000 O 22 +-3 - %z 2001 S 13 +-3 B %z 2002 O +-3 - %z +Z America/Glace_Bay -3:59:48 - LMT 1902 Jun 15 +-4 C A%sT 1953 +-4 H A%sT 1954 +-4 - AST 1972 +-4 H A%sT 1974 +-4 C A%sT +Z America/Goose_Bay -4:1:40 - LMT 1884 +-3:30:52 - NST 1918 +-3:30:52 C N%sT 1919 +-3:30:52 - NST 1935 Mar 30 +-3:30 - NST 1936 +-3:30 j N%sT 1942 May 11 +-3:30 C N%sT 1946 +-3:30 j N%sT 1966 Mar 15 2 +-4 j A%sT 2011 N +-4 C A%sT +Z America/Grand_Turk -4:44:32 - LMT 1890 +-5:7:10 - KMT 1912 F +-5 - EST 1979 +-5 u E%sT 2015 Mar 8 2 +-4 - AST 2018 Mar 11 3 +-5 u E%sT +Z America/Guatemala -6:2:4 - LMT 1918 O 5 +-6 GT C%sT +Z America/Guayaquil -5:19:20 - LMT 1890 +-5:14 - QMT 1931 +-5 EC %z +Z America/Guyana -3:52:39 - LMT 1911 Au +-4 - %z 1915 Mar +-3:45 - %z 1975 Au +-3 - %z 1992 Mar 29 1 +-4 - %z +Z America/Halifax -4:14:24 - LMT 1902 Jun 15 +-4 H A%sT 1918 +-4 C A%sT 1919 +-4 H A%sT 1942 F 9 2s +-4 C A%sT 1946 +-4 H A%sT 1974 +-4 C A%sT +Z America/Havana -5:29:28 - LMT 1890 +-5:29:36 - HMT 1925 Jul 19 12 +-5 Q C%sT +Z America/Hermosillo -7:23:52 - LMT 1922 Ja 1 7u +-7 - MST 1927 Jun 10 +-6 - CST 1930 N 15 +-7 m M%sT 1932 Ap +-6 - CST 1942 Ap 24 +-7 - MST 1996 +-7 m M%sT 1999 +-7 - MST +Z America/Indiana/Indianapolis -5:44:38 - LMT 1883 N 18 18u +-6 u C%sT 1920 +-6 In C%sT 1942 +-6 u C%sT 1946 +-6 In C%sT 1955 Ap 24 2 +-5 - EST 1957 S 29 2 +-6 - CST 1958 Ap 27 2 +-5 - EST 1969 +-5 u E%sT 1971 +-5 - EST 2006 +-5 u E%sT +Z America/Indiana/Knox -5:46:30 - LMT 1883 N 18 18u +-6 u C%sT 1947 +-6 St C%sT 1962 Ap 29 2 +-5 - EST 1963 O 27 2 +-6 u C%sT 1991 O 27 2 +-5 - EST 2006 Ap 2 2 +-6 u C%sT +Z America/Indiana/Marengo -5:45:23 - LMT 1883 N 18 18u +-6 u C%sT 1951 +-6 Ma C%sT 1961 Ap 30 2 +-5 - EST 1969 +-5 u E%sT 1974 Ja 6 2 +-6 1 CDT 1974 O 27 2 +-5 u E%sT 1976 +-5 - EST 2006 +-5 u E%sT +Z America/Indiana/Petersburg -5:49:7 - LMT 1883 N 18 18u +-6 u C%sT 1955 +-6 Pi C%sT 1965 Ap 25 2 +-5 - EST 1966 O 30 2 +-6 u C%sT 1977 O 30 2 +-5 - EST 2006 Ap 2 2 +-6 u C%sT 2007 N 4 2 +-5 u E%sT +Z America/Indiana/Tell_City -5:47:3 - LMT 1883 N 18 18u +-6 u C%sT 1946 +-6 Pe C%sT 1964 Ap 26 2 +-5 - EST 1967 O 29 2 +-6 u C%sT 1969 Ap 27 2 +-5 u E%sT 1971 +-5 - EST 2006 Ap 2 2 +-6 u C%sT +Z America/Indiana/Vevay -5:40:16 - LMT 1883 N 18 18u +-6 u C%sT 1954 Ap 25 2 +-5 - EST 1969 +-5 u E%sT 1973 +-5 - EST 2006 +-5 u E%sT +Z America/Indiana/Vincennes -5:50:7 - LMT 1883 N 18 18u +-6 u C%sT 1946 +-6 V C%sT 1964 Ap 26 2 +-5 - EST 1969 +-5 u E%sT 1971 +-5 - EST 2006 Ap 2 2 +-6 u C%sT 2007 N 4 2 +-5 u E%sT +Z America/Indiana/Winamac -5:46:25 - LMT 1883 N 18 18u +-6 u C%sT 1946 +-6 Pu C%sT 1961 Ap 30 2 +-5 - EST 1969 +-5 u E%sT 1971 +-5 - EST 2006 Ap 2 2 +-6 u C%sT 2007 Mar 11 2 +-5 u E%sT +Z America/Inuvik 0 - -00 1953 +-8 Y P%sT 1979 Ap lastSu 2 +-7 Y M%sT 1980 +-7 C M%sT +Z America/Iqaluit 0 - -00 1942 Au +-5 Y E%sT 1999 O 31 2 +-6 C C%sT 2000 O 29 2 +-5 C E%sT +Z America/Jamaica -5:7:10 - LMT 1890 +-5:7:10 - KMT 1912 F +-5 - EST 1974 +-5 u E%sT 1984 +-5 - EST +Z America/Juneau 15:2:19 - LMT 1867 O 19 15:33:32 +-8:57:41 - LMT 1900 Au 20 12 +-8 - PST 1942 +-8 u P%sT 1946 +-8 - PST 1969 +-8 u P%sT 1980 Ap 27 2 +-9 u Y%sT 1980 O 26 2 +-8 u P%sT 1983 O 30 2 +-9 u Y%sT 1983 N 30 +-9 u AK%sT +Z America/Kentucky/Louisville -5:43:2 - LMT 1883 N 18 18u +-6 u C%sT 1921 +-6 v C%sT 1942 +-6 u C%sT 1946 +-6 v C%sT 1961 Jul 23 2 +-5 - EST 1968 +-5 u E%sT 1974 Ja 6 2 +-6 1 CDT 1974 O 27 2 +-5 u E%sT +Z America/Kentucky/Monticello -5:39:24 - LMT 1883 N 18 18u +-6 u C%sT 1946 +-6 - CST 1968 +-6 u C%sT 2000 O 29 2 +-5 u E%sT +Z America/La_Paz -4:32:36 - LMT 1890 +-4:32:36 - CMT 1931 O 15 +-4:32:36 1 BST 1932 Mar 21 +-4 - %z +Z America/Lima -5:8:12 - LMT 1890 +-5:8:36 - LMT 1908 Jul 28 +-5 PE %z +Z America/Los_Angeles -7:52:58 - LMT 1883 N 18 20u +-8 u P%sT 1946 +-8 CA P%sT 1967 +-8 u P%sT +Z America/Maceio -2:22:52 - LMT 1914 +-3 B %z 1990 S 17 +-3 - %z 1995 O 13 +-3 B %z 1996 S 4 +-3 - %z 1999 S 30 +-3 B %z 2000 O 22 +-3 - %z 2001 S 13 +-3 B %z 2002 O +-3 - %z +Z America/Managua -5:45:8 - LMT 1890 +-5:45:12 - MMT 1934 Jun 23 +-6 - CST 1973 May +-5 - EST 1975 F 16 +-6 NI C%sT 1992 Ja 1 4 +-5 - EST 1992 S 24 +-6 - CST 1993 +-5 - EST 1997 +-6 NI C%sT +Z America/Manaus -4:0:4 - LMT 1914 +-4 B %z 1988 S 12 +-4 - %z 1993 S 28 +-4 B %z 1994 S 22 +-4 - %z +Z America/Martinique -4:4:20 - LMT 1890 +-4:4:20 - FFMT 1911 May +-4 - AST 1980 Ap 6 +-4 1 ADT 1980 S 28 +-4 - AST +Z America/Matamoros -6:30 - LMT 1922 Ja 1 6u +-6 - CST 1988 +-6 u C%sT 1989 +-6 m C%sT 2010 +-6 u C%sT +Z America/Mazatlan -7:5:40 - LMT 1922 Ja 1 7u +-7 - MST 1927 Jun 10 +-6 - CST 1930 N 15 +-7 m M%sT 1932 Ap +-6 - CST 1942 Ap 24 +-7 - MST 1970 +-7 m M%sT +Z America/Menominee -5:50:27 - LMT 1885 S 18 12 +-6 u C%sT 1946 +-6 Me C%sT 1969 Ap 27 2 +-5 - EST 1973 Ap 29 2 +-6 u C%sT +Z America/Merida -5:58:28 - LMT 1922 Ja 1 6u +-6 - CST 1981 D 26 2 +-5 - EST 1982 N 2 2 +-6 m C%sT +Z America/Metlakatla 15:13:42 - LMT 1867 O 19 15:44:55 +-8:46:18 - LMT 1900 Au 20 12 +-8 - PST 1942 +-8 u P%sT 1946 +-8 - PST 1969 +-8 u P%sT 1983 O 30 2 +-8 - PST 2015 N 1 2 +-9 u AK%sT 2018 N 4 2 +-8 - PST 2019 Ja 20 2 +-9 u AK%sT +Z America/Mexico_City -6:36:36 - LMT 1922 Ja 1 7u +-7 - MST 1927 Jun 10 +-6 - CST 1930 N 15 +-7 m M%sT 1932 Ap +-6 m C%sT 2001 S 30 2 +-6 - CST 2002 F 20 +-6 m C%sT +Z America/Miquelon -3:44:40 - LMT 1911 Jun 15 +-4 - AST 1980 May +-3 - %z 1987 +-3 C %z +Z America/Moncton -4:19:8 - LMT 1883 D 9 +-5 - EST 1902 Jun 15 +-4 C A%sT 1933 +-4 o A%sT 1942 +-4 C A%sT 1946 +-4 o A%sT 1973 +-4 C A%sT 1993 +-4 o A%sT 2007 +-4 C A%sT +Z America/Monterrey -6:41:16 - LMT 1922 Ja 1 6u +-7 - MST 1927 Jun 10 +-6 - CST 1930 N 15 +-7 m M%sT 1932 Ap +-6 - CST 1988 +-6 u C%sT 1989 +-6 m C%sT +Z America/Montevideo -3:44:51 - LMT 1908 Jun 10 +-3:44:51 - MMT 1920 May +-4 - %z 1923 O +-3:30 U %z 1942 D 14 +-3 U %z 1960 +-3 U %z 1968 +-3 U %z 1970 +-3 U %z 1974 +-3 U %z 1974 Mar 10 +-3 U %z 1974 D 22 +-3 U %z +Z America/New_York -4:56:2 - LMT 1883 N 18 17u +-5 u E%sT 1920 +-5 NY E%sT 1942 +-5 u E%sT 1946 +-5 NY E%sT 1967 +-5 u E%sT +Z America/Nome 12:58:22 - LMT 1867 O 19 13:29:35 +-11:1:38 - LMT 1900 Au 20 12 +-11 - NST 1942 +-11 u N%sT 1946 +-11 - NST 1967 Ap +-11 - BST 1969 +-11 u B%sT 1983 O 30 2 +-9 u Y%sT 1983 N 30 +-9 u AK%sT +Z America/Noronha -2:9:40 - LMT 1914 +-2 B %z 1990 S 17 +-2 - %z 1999 S 30 +-2 B %z 2000 O 15 +-2 - %z 2001 S 13 +-2 B %z 2002 O +-2 - %z +Z America/North_Dakota/Beulah -6:47:7 - LMT 1883 N 18 19u +-7 u M%sT 2010 N 7 2 +-6 u C%sT +Z America/North_Dakota/Center -6:45:12 - LMT 1883 N 18 19u +-7 u M%sT 1992 O 25 2 +-6 u C%sT +Z America/North_Dakota/New_Salem -6:45:39 - LMT 1883 N 18 19u +-7 u M%sT 2003 O 26 2 +-6 u C%sT +Z America/Nuuk -3:26:56 - LMT 1916 Jul 28 +-3 - %z 1980 Ap 6 2 +-3 E %z 2023 Mar 26 1u +-2 - %z 2023 O 29 1u +-2 E %z +Z America/Ojinaga -6:57:40 - LMT 1922 Ja 1 7u +-7 - MST 1927 Jun 10 +-6 - CST 1930 N 15 +-7 m M%sT 1932 Ap +-6 - CST 1996 +-6 m C%sT 1998 +-6 - CST 1998 Ap Su>=1 3 +-7 m M%sT 2010 +-7 u M%sT 2022 O 30 2 +-6 - CST 2022 N 30 +-6 u C%sT +Z America/Panama -5:18:8 - LMT 1890 +-5:19:36 - CMT 1908 Ap 22 +-5 - EST +Z America/Paramaribo -3:40:40 - LMT 1911 +-3:40:52 - PMT 1935 +-3:40:36 - PMT 1945 O +-3:30 - %z 1984 O +-3 - %z +Z America/Phoenix -7:28:18 - LMT 1883 N 18 19u +-7 u M%sT 1944 Ja 1 0:1 +-7 - MST 1944 Ap 1 0:1 +-7 u M%sT 1944 O 1 0:1 +-7 - MST 1967 +-7 u M%sT 1968 Mar 21 +-7 - MST +Z America/Port-au-Prince -4:49:20 - LMT 1890 +-4:49 - PPMT 1917 Ja 24 12 +-5 HT E%sT +Z America/Porto_Velho -4:15:36 - LMT 1914 +-4 B %z 1988 S 12 +-4 - %z +Z America/Puerto_Rico -4:24:25 - LMT 1899 Mar 28 12 +-4 - AST 1942 May 3 +-4 u A%sT 1946 +-4 - AST +Z America/Punta_Arenas -4:43:40 - LMT 1890 +-4:42:45 - SMT 1910 Ja 10 +-5 - %z 1916 Jul +-4:42:45 - SMT 1918 S 10 +-4 - %z 1919 Jul +-4:42:45 - SMT 1927 S +-5 x %z 1932 S +-4 - %z 1942 Jun +-5 - %z 1942 Au +-4 - %z 1946 Au 28 24 +-5 1 %z 1947 Mar 31 24 +-5 - %z 1947 May 21 23 +-4 x %z 2016 D 4 +-3 - %z +Z America/Rankin_Inlet 0 - -00 1957 +-6 Y C%sT 2000 O 29 2 +-5 - EST 2001 Ap 1 3 +-6 C C%sT +Z America/Recife -2:19:36 - LMT 1914 +-3 B %z 1990 S 17 +-3 - %z 1999 S 30 +-3 B %z 2000 O 15 +-3 - %z 2001 S 13 +-3 B %z 2002 O +-3 - %z +Z America/Regina -6:58:36 - LMT 1905 S +-7 r M%sT 1960 Ap lastSu 2 +-6 - CST +Z America/Resolute 0 - -00 1947 Au 31 +-6 Y C%sT 2000 O 29 2 +-5 - EST 2001 Ap 1 3 +-6 C C%sT 2006 O 29 2 +-5 - EST 2007 Mar 11 3 +-6 C C%sT +Z America/Rio_Branco -4:31:12 - LMT 1914 +-5 B %z 1988 S 12 +-5 - %z 2008 Jun 24 +-4 - %z 2013 N 10 +-5 - %z +Z America/Santarem -3:38:48 - LMT 1914 +-4 B %z 1988 S 12 +-4 - %z 2008 Jun 24 +-3 - %z +Z America/Santiago -4:42:45 - LMT 1890 +-4:42:45 - SMT 1910 Ja 10 +-5 - %z 1916 Jul +-4:42:45 - SMT 1918 S 10 +-4 - %z 1919 Jul +-4:42:45 - SMT 1927 S +-5 x %z 1932 S +-4 - %z 1942 Jun +-5 - %z 1942 Au +-4 - %z 1946 Jul 14 24 +-4 1 %z 1946 Au 28 24 +-5 1 %z 1947 Mar 31 24 +-5 - %z 1947 May 21 23 +-4 x %z +Z America/Santo_Domingo -4:39:36 - LMT 1890 +-4:40 - SDMT 1933 Ap 1 12 +-5 DO %s 1974 O 27 +-4 - AST 2000 O 29 2 +-5 u E%sT 2000 D 3 1 +-4 - AST +Z America/Sao_Paulo -3:6:28 - LMT 1914 +-3 B %z 1963 O 23 +-3 1 %z 1964 +-3 B %z +Z America/Scoresbysund -1:27:52 - LMT 1916 Jul 28 +-2 - %z 1980 Ap 6 2 +-2 c %z 1981 Mar 29 +-1 E %z 2024 Mar 31 +-2 E %z +Z America/Sitka 14:58:47 - LMT 1867 O 19 15:30 +-9:1:13 - LMT 1900 Au 20 12 +-8 - PST 1942 +-8 u P%sT 1946 +-8 - PST 1969 +-8 u P%sT 1983 O 30 2 +-9 u Y%sT 1983 N 30 +-9 u AK%sT +Z America/St_Johns -3:30:52 - LMT 1884 +-3:30:52 j N%sT 1918 +-3:30:52 C N%sT 1919 +-3:30:52 j N%sT 1935 Mar 30 +-3:30 j N%sT 1942 May 11 +-3:30 C N%sT 1946 +-3:30 j N%sT 2011 N +-3:30 C N%sT +Z America/Swift_Current -7:11:20 - LMT 1905 S +-7 C M%sT 1946 Ap lastSu 2 +-7 r M%sT 1950 +-7 Sw M%sT 1972 Ap lastSu 2 +-6 - CST +Z America/Tegucigalpa -5:48:52 - LMT 1921 Ap +-6 HN C%sT +Z America/Thule -4:35:8 - LMT 1916 Jul 28 +-4 Th A%sT +Z America/Tijuana -7:48:4 - LMT 1922 Ja 1 7u +-7 - MST 1924 +-8 - PST 1927 Jun 10 +-7 - MST 1930 N 15 +-8 - PST 1931 Ap +-8 1 PDT 1931 S 30 +-8 - PST 1942 Ap 24 +-8 1 PWT 1945 Au 14 23u +-8 1 PPT 1945 N 15 +-8 - PST 1948 Ap 5 +-8 1 PDT 1949 Ja 14 +-8 - PST 1950 May +-8 1 PDT 1950 S 24 +-8 - PST 1951 Ap 29 2 +-8 1 PDT 1951 S 30 2 +-8 - PST 1952 Ap 27 2 +-8 1 PDT 1952 S 28 2 +-8 - PST 1954 +-8 CA P%sT 1961 +-8 - PST 1976 +-8 u P%sT 1996 +-8 m P%sT 2001 +-8 u P%sT 2002 F 20 +-8 m P%sT 2010 +-8 u P%sT +Z America/Toronto -5:17:32 - LMT 1895 +-5 C E%sT 1919 +-5 t E%sT 1942 F 9 2s +-5 C E%sT 1946 +-5 t E%sT 1974 +-5 C E%sT +Z America/Vancouver -8:12:28 - LMT 1884 +-8 Va P%sT 1987 +-8 C P%sT +Z America/Whitehorse -9:0:12 - LMT 1900 Au 20 +-9 Y Y%sT 1965 +-9 Yu Y%sT 1966 F 27 +-8 - PST 1980 +-8 C P%sT 2020 N +-7 - MST +Z America/Winnipeg -6:28:36 - LMT 1887 Jul 16 +-6 W C%sT 2006 +-6 C C%sT +Z America/Yakutat 14:41:5 - LMT 1867 O 19 15:12:18 +-9:18:55 - LMT 1900 Au 20 12 +-9 - YST 1942 +-9 u Y%sT 1946 +-9 - YST 1969 +-9 u Y%sT 1983 N 30 +-9 u AK%sT +Z Antarctica/Casey 0 - -00 1969 +8 - %z 2009 O 18 2 +11 - %z 2010 Mar 5 2 +8 - %z 2011 O 28 2 +11 - %z 2012 F 21 17u +8 - %z 2016 O 22 +11 - %z 2018 Mar 11 4 +8 - %z 2018 O 7 4 +11 - %z 2019 Mar 17 3 +8 - %z 2019 O 4 3 +11 - %z 2020 Mar 8 3 +8 - %z 2020 O 4 0:1 +11 - %z 2021 Mar 14 +8 - %z 2021 O 3 0:1 +11 - %z 2022 Mar 13 +8 - %z 2022 O 2 0:1 +11 - %z 2023 Mar 9 3 +8 - %z +Z Antarctica/Davis 0 - -00 1957 Ja 13 +7 - %z 1964 N +0 - -00 1969 F +7 - %z 2009 O 18 2 +5 - %z 2010 Mar 10 20u +7 - %z 2011 O 28 2 +5 - %z 2012 F 21 20u +7 - %z +Z Antarctica/Macquarie 0 - -00 1899 N +10 - AEST 1916 O 1 2 +10 1 AEDT 1917 F +10 AU AE%sT 1919 Ap 1 0s +0 - -00 1948 Mar 25 +10 AU AE%sT 1967 +10 AT AE%sT 2010 +10 1 AEDT 2011 +10 AT AE%sT +Z Antarctica/Mawson 0 - -00 1954 F 13 +6 - %z 2009 O 18 2 +5 - %z +Z Antarctica/Palmer 0 - -00 1965 +-4 A %z 1969 O 5 +-3 A %z 1982 May +-4 x %z 2016 D 4 +-3 - %z +Z Antarctica/Rothera 0 - -00 1976 D +-3 - %z +Z Antarctica/Troll 0 - -00 2005 F 12 +0 Tr %s +Z Antarctica/Vostok 0 - -00 1957 D 16 +7 - %z 1994 F +0 - -00 1994 N +7 - %z 2023 D 18 2 +5 - %z +Z Asia/Almaty 5:7:48 - LMT 1924 May 2 +5 - %z 1930 Jun 21 +6 R %z 1991 Mar 31 2s +5 R %z 1992 Ja 19 2s +6 R %z 2004 O 31 2s +6 - %z 2024 Mar +5 - %z +Z Asia/Amman 2:23:44 - LMT 1931 +2 J EE%sT 2022 O 28 0s +3 - %z +Z Asia/Anadyr 11:49:56 - LMT 1924 May 2 +12 - %z 1930 Jun 21 +13 R %z 1982 Ap 1 0s +12 R %z 1991 Mar 31 2s +11 R %z 1992 Ja 19 2s +12 R %z 2010 Mar 28 2s +11 R %z 2011 Mar 27 2s +12 - %z +Z Asia/Aqtau 3:21:4 - LMT 1924 May 2 +4 - %z 1930 Jun 21 +5 - %z 1981 O +6 - %z 1982 Ap +5 R %z 1991 Mar 31 2s +4 R %z 1992 Ja 19 2s +5 R %z 1994 S 25 2s +4 R %z 2004 O 31 2s +5 - %z +Z Asia/Aqtobe 3:48:40 - LMT 1924 May 2 +4 - %z 1930 Jun 21 +5 - %z 1981 Ap +5 1 %z 1981 O +6 - %z 1982 Ap +5 R %z 1991 Mar 31 2s +4 R %z 1992 Ja 19 2s +5 R %z 2004 O 31 2s +5 - %z +Z Asia/Ashgabat 3:53:32 - LMT 1924 May 2 +4 - %z 1930 Jun 21 +5 R %z 1991 Mar 31 2 +4 R %z 1992 Ja 19 2 +5 - %z +Z Asia/Atyrau 3:27:44 - LMT 1924 May 2 +3 - %z 1930 Jun 21 +5 - %z 1981 O +6 - %z 1982 Ap +5 R %z 1991 Mar 31 2s +4 R %z 1992 Ja 19 2s +5 R %z 1999 Mar 28 2s +4 R %z 2004 O 31 2s +5 - %z +Z Asia/Baghdad 2:57:40 - LMT 1890 +2:57:36 - BMT 1918 +3 - %z 1982 May +3 IQ %z +Z Asia/Baku 3:19:24 - LMT 1924 May 2 +3 - %z 1957 Mar +4 R %z 1991 Mar 31 2s +3 R %z 1992 S lastSu 2s +4 - %z 1996 +4 E %z 1997 +4 AZ %z +Z Asia/Bangkok 6:42:4 - LMT 1880 +6:42:4 - BMT 1920 Ap +7 - %z +Z Asia/Barnaul 5:35 - LMT 1919 D 10 +6 - %z 1930 Jun 21 +7 R %z 1991 Mar 31 2s +6 R %z 1992 Ja 19 2s +7 R %z 1995 May 28 +6 R %z 2011 Mar 27 2s +7 - %z 2014 O 26 2s +6 - %z 2016 Mar 27 2s +7 - %z +Z Asia/Beirut 2:22 - LMT 1880 +2 l EE%sT +Z Asia/Bishkek 4:58:24 - LMT 1924 May 2 +5 - %z 1930 Jun 21 +6 R %z 1991 Mar 31 2s +5 R %z 1991 Au 31 2 +5 KG %z 2005 Au 12 +6 - %z +Z Asia/Chita 7:33:52 - LMT 1919 D 15 +8 - %z 1930 Jun 21 +9 R %z 1991 Mar 31 2s +8 R %z 1992 Ja 19 2s +9 R %z 2011 Mar 27 2s +10 - %z 2014 O 26 2s +8 - %z 2016 Mar 27 2 +9 - %z +Z Asia/Colombo 5:19:24 - LMT 1880 +5:19:32 - MMT 1906 +5:30 - %z 1942 Ja 5 +5:30 0:30 %z 1942 S +5:30 1 %z 1945 O 16 2 +5:30 - %z 1996 May 25 +6:30 - %z 1996 O 26 0:30 +6 - %z 2006 Ap 15 0:30 +5:30 - %z +Z Asia/Damascus 2:25:12 - LMT 1920 +2 S EE%sT 2022 O 28 +3 - %z +Z Asia/Dhaka 6:1:40 - LMT 1890 +5:53:20 - HMT 1941 O +6:30 - %z 1942 May 15 +5:30 - %z 1942 S +6:30 - %z 1951 S 30 +6 - %z 2009 +6 BD %z +Z Asia/Dili 8:22:20 - LMT 1911 D 31 16u +8 - %z 1942 F 21 23 +9 - %z 1976 May 3 +8 - %z 2000 S 17 +9 - %z +Z Asia/Dubai 3:41:12 - LMT 1920 +4 - %z +Z Asia/Dushanbe 4:35:12 - LMT 1924 May 2 +5 - %z 1930 Jun 21 +6 R %z 1991 Mar 31 2s +5 1 %z 1991 S 9 2s +5 - %z +Z Asia/Famagusta 2:15:48 - LMT 1921 N 14 +2 CY EE%sT 1998 S +2 E EE%sT 2016 S 8 +3 - %z 2017 O 29 1u +2 E EE%sT +Z Asia/Gaza 2:17:52 - LMT 1900 O +2 Z EET/EEST 1948 May 15 +2 K EE%sT 1967 Jun 5 +2 Z I%sT 1996 +2 J EE%sT 1999 +2 P EE%sT 2008 Au 29 +2 - EET 2008 S +2 P EE%sT 2010 +2 - EET 2010 Mar 27 0:1 +2 P EE%sT 2011 Au +2 - EET 2012 +2 P EE%sT +Z Asia/Hebron 2:20:23 - LMT 1900 O +2 Z EET/EEST 1948 May 15 +2 K EE%sT 1967 Jun 5 +2 Z I%sT 1996 +2 J EE%sT 1999 +2 P EE%sT +Z Asia/Ho_Chi_Minh 7:6:30 - LMT 1906 Jul +7:6:30 - PLMT 1911 May +7 - %z 1942 D 31 23 +8 - %z 1945 Mar 14 23 +9 - %z 1945 S 1 24 +7 - %z 1947 Ap +8 - %z 1955 Jul 1 1 +7 - %z 1959 D 31 23 +8 - %z 1975 Jun 13 +7 - %z +Z Asia/Hong_Kong 7:36:42 - LMT 1904 O 29 17u +8 - HKT 1941 Jun 15 3 +8 1 HKST 1941 O 1 4 +8 0:30 HKWT 1941 D 25 +9 - JST 1945 N 18 2 +8 HK HK%sT +Z Asia/Hovd 6:6:36 - LMT 1905 Au +6 - %z 1978 +7 X %z +Z Asia/Irkutsk 6:57:5 - LMT 1880 +6:57:5 - IMT 1920 Ja 25 +7 - %z 1930 Jun 21 +8 R %z 1991 Mar 31 2s +7 R %z 1992 Ja 19 2s +8 R %z 2011 Mar 27 2s +9 - %z 2014 O 26 2s +8 - %z +Z Asia/Jakarta 7:7:12 - LMT 1867 Au 10 +7:7:12 - BMT 1923 D 31 16:40u +7:20 - %z 1932 N +7:30 - %z 1942 Mar 23 +9 - %z 1945 S 23 +7:30 - %z 1948 May +8 - %z 1950 May +7:30 - %z 1964 +7 - WIB +Z Asia/Jayapura 9:22:48 - LMT 1932 N +9 - %z 1944 S +9:30 - %z 1964 +9 - WIT +Z Asia/Jerusalem 2:20:54 - LMT 1880 +2:20:40 - JMT 1918 +2 Z I%sT +Z Asia/Kabul 4:36:48 - LMT 1890 +4 - %z 1945 +4:30 - %z +Z Asia/Kamchatka 10:34:36 - LMT 1922 N 10 +11 - %z 1930 Jun 21 +12 R %z 1991 Mar 31 2s +11 R %z 1992 Ja 19 2s +12 R %z 2010 Mar 28 2s +11 R %z 2011 Mar 27 2s +12 - %z +Z Asia/Karachi 4:28:12 - LMT 1907 +5:30 - %z 1942 S +5:30 1 %z 1945 O 15 +5:30 - %z 1951 S 30 +5 - %z 1971 Mar 26 +5 PK PK%sT +Z Asia/Kathmandu 5:41:16 - LMT 1920 +5:30 - %z 1986 +5:45 - %z +Z Asia/Khandyga 9:2:13 - LMT 1919 D 15 +8 - %z 1930 Jun 21 +9 R %z 1991 Mar 31 2s +8 R %z 1992 Ja 19 2s +9 R %z 2004 +10 R %z 2011 Mar 27 2s +11 - %z 2011 S 13 0s +10 - %z 2014 O 26 2s +9 - %z +Z Asia/Kolkata 5:53:28 - LMT 1854 Jun 28 +5:53:20 - HMT 1870 +5:21:10 - MMT 1906 +5:30 - IST 1941 O +5:30 1 %z 1942 May 15 +5:30 - IST 1942 S +5:30 1 %z 1945 O 15 +5:30 - IST +Z Asia/Krasnoyarsk 6:11:26 - LMT 1920 Ja 6 +6 - %z 1930 Jun 21 +7 R %z 1991 Mar 31 2s +6 R %z 1992 Ja 19 2s +7 R %z 2011 Mar 27 2s +8 - %z 2014 O 26 2s +7 - %z +Z Asia/Kuching 7:21:20 - LMT 1926 Mar +7:30 - %z 1933 +8 NB %z 1942 F 16 +9 - %z 1945 S 12 +8 - %z +Z Asia/Macau 7:34:10 - LMT 1904 O 30 +8 - CST 1941 D 21 23 +9 _ %z 1945 S 30 24 +8 _ C%sT +Z Asia/Magadan 10:3:12 - LMT 1924 May 2 +10 - %z 1930 Jun 21 +11 R %z 1991 Mar 31 2s +10 R %z 1992 Ja 19 2s +11 R %z 2011 Mar 27 2s +12 - %z 2014 O 26 2s +10 - %z 2016 Ap 24 2s +11 - %z +Z Asia/Makassar 7:57:36 - LMT 1920 +7:57:36 - MMT 1932 N +8 - %z 1942 F 9 +9 - %z 1945 S 23 +8 - WITA +Z Asia/Manila -15:56:8 - LMT 1844 D 31 +8:3:52 - LMT 1899 S 6 4u +8 PH P%sT 1942 F 11 24 +9 - JST 1945 Mar 4 +8 PH P%sT +Z Asia/Nicosia 2:13:28 - LMT 1921 N 14 +2 CY EE%sT 1998 S +2 E EE%sT +Z Asia/Novokuznetsk 5:48:48 - LMT 1924 May +6 - %z 1930 Jun 21 +7 R %z 1991 Mar 31 2s +6 R %z 1992 Ja 19 2s +7 R %z 2010 Mar 28 2s +6 R %z 2011 Mar 27 2s +7 - %z +Z Asia/Novosibirsk 5:31:40 - LMT 1919 D 14 6 +6 - %z 1930 Jun 21 +7 R %z 1991 Mar 31 2s +6 R %z 1992 Ja 19 2s +7 R %z 1993 May 23 +6 R %z 2011 Mar 27 2s +7 - %z 2014 O 26 2s +6 - %z 2016 Jul 24 2s +7 - %z +Z Asia/Omsk 4:53:30 - LMT 1919 N 14 +5 - %z 1930 Jun 21 +6 R %z 1991 Mar 31 2s +5 R %z 1992 Ja 19 2s +6 R %z 2011 Mar 27 2s +7 - %z 2014 O 26 2s +6 - %z +Z Asia/Oral 3:25:24 - LMT 1924 May 2 +3 - %z 1930 Jun 21 +5 - %z 1981 Ap +5 1 %z 1981 O +6 - %z 1982 Ap +5 R %z 1989 Mar 26 2s +4 R %z 1992 Ja 19 2s +5 R %z 1992 Mar 29 2s +4 R %z 2004 O 31 2s +5 - %z +Z Asia/Pontianak 7:17:20 - LMT 1908 May +7:17:20 - PMT 1932 N +7:30 - %z 1942 Ja 29 +9 - %z 1945 S 23 +7:30 - %z 1948 May +8 - %z 1950 May +7:30 - %z 1964 +8 - WITA 1988 +7 - WIB +Z Asia/Pyongyang 8:23 - LMT 1908 Ap +8:30 - KST 1912 +9 - JST 1945 Au 24 +9 - KST 2015 Au 15 +8:30 - KST 2018 May 4 23:30 +9 - KST +Z Asia/Qatar 3:26:8 - LMT 1920 +4 - %z 1972 Jun +3 - %z +Z Asia/Qostanay 4:14:28 - LMT 1924 May 2 +4 - %z 1930 Jun 21 +5 - %z 1981 Ap +5 1 %z 1981 O +6 - %z 1982 Ap +5 R %z 1991 Mar 31 2s +4 R %z 1992 Ja 19 2s +5 R %z 2004 O 31 2s +6 - %z 2024 Mar +5 - %z +Z Asia/Qyzylorda 4:21:52 - LMT 1924 May 2 +4 - %z 1930 Jun 21 +5 - %z 1981 Ap +5 1 %z 1981 O +6 - %z 1982 Ap +5 R %z 1991 Mar 31 2s +4 R %z 1991 S 29 2s +5 R %z 1992 Ja 19 2s +6 R %z 1992 Mar 29 2s +5 R %z 2004 O 31 2s +6 - %z 2018 D 21 +5 - %z +Z Asia/Riyadh 3:6:52 - LMT 1947 Mar 14 +3 - %z +Z Asia/Sakhalin 9:30:48 - LMT 1905 Au 23 +9 - %z 1945 Au 25 +11 R %z 1991 Mar 31 2s +10 R %z 1992 Ja 19 2s +11 R %z 1997 Mar lastSu 2s +10 R %z 2011 Mar 27 2s +11 - %z 2014 O 26 2s +10 - %z 2016 Mar 27 2s +11 - %z +Z Asia/Samarkand 4:27:53 - LMT 1924 May 2 +4 - %z 1930 Jun 21 +5 - %z 1981 Ap +5 1 %z 1981 O +6 - %z 1982 Ap +5 R %z 1992 +5 - %z +Z Asia/Seoul 8:27:52 - LMT 1908 Ap +8:30 - KST 1912 +9 - JST 1945 S 8 +9 KR K%sT 1954 Mar 21 +8:30 KR K%sT 1961 Au 10 +9 KR K%sT +Z Asia/Shanghai 8:5:43 - LMT 1901 +8 Sh C%sT 1949 May 28 +8 CN C%sT +Z Asia/Singapore 6:55:25 - LMT 1901 +6:55:25 - SMT 1905 Jun +7 - %z 1933 +7 0:20 %z 1936 +7:20 - %z 1941 S +7:30 - %z 1942 F 16 +9 - %z 1945 S 12 +7:30 - %z 1981 D 31 16u +8 - %z +Z Asia/Srednekolymsk 10:14:52 - LMT 1924 May 2 +10 - %z 1930 Jun 21 +11 R %z 1991 Mar 31 2s +10 R %z 1992 Ja 19 2s +11 R %z 2011 Mar 27 2s +12 - %z 2014 O 26 2s +11 - %z +Z Asia/Taipei 8:6 - LMT 1896 +8 - CST 1937 O +9 - JST 1945 S 21 1 +8 f C%sT +Z Asia/Tashkent 4:37:11 - LMT 1924 May 2 +5 - %z 1930 Jun 21 +6 R %z 1991 Mar 31 2 +5 R %z 1992 +5 - %z +Z Asia/Tbilisi 2:59:11 - LMT 1880 +2:59:11 - TBMT 1924 May 2 +3 - %z 1957 Mar +4 R %z 1991 Mar 31 2s +3 R %z 1992 +3 e %z 1994 S lastSu +4 e %z 1996 O lastSu +4 1 %z 1997 Mar lastSu +4 e %z 2004 Jun 27 +3 R %z 2005 Mar lastSu 2 +4 - %z +Z Asia/Tehran 3:25:44 - LMT 1916 +3:25:44 - TMT 1935 Jun 13 +3:30 i %z 1977 O 20 24 +4 i %z 1978 N 10 24 +3:30 i %z +Z Asia/Thimphu 5:58:36 - LMT 1947 Au 15 +5:30 - %z 1987 O +6 - %z +Z Asia/Tokyo 9:18:59 - LMT 1887 D 31 15u +9 JP J%sT +Z Asia/Tomsk 5:39:51 - LMT 1919 D 22 +6 - %z 1930 Jun 21 +7 R %z 1991 Mar 31 2s +6 R %z 1992 Ja 19 2s +7 R %z 2002 May 1 3 +6 R %z 2011 Mar 27 2s +7 - %z 2014 O 26 2s +6 - %z 2016 May 29 2s +7 - %z +Z Asia/Ulaanbaatar 7:7:32 - LMT 1905 Au +7 - %z 1978 +8 X %z +Z Asia/Urumqi 5:50:20 - LMT 1928 +6 - %z +Z Asia/Ust-Nera 9:32:54 - LMT 1919 D 15 +8 - %z 1930 Jun 21 +9 R %z 1981 Ap +11 R %z 1991 Mar 31 2s +10 R %z 1992 Ja 19 2s +11 R %z 2011 Mar 27 2s +12 - %z 2011 S 13 0s +11 - %z 2014 O 26 2s +10 - %z +Z Asia/Vladivostok 8:47:31 - LMT 1922 N 15 +9 - %z 1930 Jun 21 +10 R %z 1991 Mar 31 2s +9 R %z 1992 Ja 19 2s +10 R %z 2011 Mar 27 2s +11 - %z 2014 O 26 2s +10 - %z +Z Asia/Yakutsk 8:38:58 - LMT 1919 D 15 +8 - %z 1930 Jun 21 +9 R %z 1991 Mar 31 2s +8 R %z 1992 Ja 19 2s +9 R %z 2011 Mar 27 2s +10 - %z 2014 O 26 2s +9 - %z +Z Asia/Yangon 6:24:47 - LMT 1880 +6:24:47 - RMT 1920 +6:30 - %z 1942 May +9 - %z 1945 May 3 +6:30 - %z +Z Asia/Yekaterinburg 4:2:33 - LMT 1916 Jul 3 +3:45:5 - PMT 1919 Jul 15 4 +4 - %z 1930 Jun 21 +5 R %z 1991 Mar 31 2s +4 R %z 1992 Ja 19 2s +5 R %z 2011 Mar 27 2s +6 - %z 2014 O 26 2s +5 - %z +Z Asia/Yerevan 2:58 - LMT 1924 May 2 +3 - %z 1957 Mar +4 R %z 1991 Mar 31 2s +3 R %z 1995 S 24 2s +4 - %z 1997 +4 R %z 2011 +4 AM %z +Z Atlantic/Azores -1:42:40 - LMT 1884 +-1:54:32 - HMT 1912 Ja 1 2u +-2 p %z 1966 O 2 2s +-1 - %z 1982 Mar 28 0s +-1 p %z 1986 +-1 E %z 1992 D 27 1s +0 E WE%sT 1993 Jun 17 1u +-1 E %z +Z Atlantic/Bermuda -4:19:18 - LMT 1890 +-4:19:18 Be BMT/BST 1930 Ja 1 2 +-4 Be A%sT 1974 Ap 28 2 +-4 C A%sT 1976 +-4 u A%sT +Z Atlantic/Canary -1:1:36 - LMT 1922 Mar +-1 - %z 1946 S 30 1 +0 - WET 1980 Ap 6 0s +0 1 WEST 1980 S 28 1u +0 E WE%sT +Z Atlantic/Cape_Verde -1:34:4 - LMT 1912 Ja 1 2u +-2 - %z 1942 S +-2 1 %z 1945 O 15 +-2 - %z 1975 N 25 2 +-1 - %z +Z Atlantic/Faroe -0:27:4 - LMT 1908 Ja 11 +0 - WET 1981 +0 E WE%sT +Z Atlantic/Madeira -1:7:36 - LMT 1884 +-1:7:36 - FMT 1912 Ja 1 1u +-1 p %z 1966 O 2 2s +0 - WET 1982 Ap 4 +0 p WE%sT 1986 Jul 31 +0 E WE%sT +Z Atlantic/South_Georgia -2:26:8 - LMT 1890 +-2 - %z +Z Atlantic/Stanley -3:51:24 - LMT 1890 +-3:51:24 - SMT 1912 Mar 12 +-4 FK %z 1983 May +-3 FK %z 1985 S 15 +-4 FK %z 2010 S 5 2 +-3 - %z +Z Australia/Adelaide 9:14:20 - LMT 1895 F +9 - ACST 1899 May +9:30 AU AC%sT 1971 +9:30 AS AC%sT +Z Australia/Brisbane 10:12:8 - LMT 1895 +10 AU AE%sT 1971 +10 AQ AE%sT +Z Australia/Broken_Hill 9:25:48 - LMT 1895 F +10 - AEST 1896 Au 23 +9 - ACST 1899 May +9:30 AU AC%sT 1971 +9:30 AN AC%sT 2000 +9:30 AS AC%sT +Z Australia/Darwin 8:43:20 - LMT 1895 F +9 - ACST 1899 May +9:30 AU AC%sT +Z Australia/Eucla 8:35:28 - LMT 1895 D +8:45 AU %z 1943 Jul +8:45 AW %z +Z Australia/Hobart 9:49:16 - LMT 1895 S +10 AT AE%sT 1919 O 24 +10 AU AE%sT 1967 +10 AT AE%sT +Z Australia/Lindeman 9:55:56 - LMT 1895 +10 AU AE%sT 1971 +10 AQ AE%sT 1992 Jul +10 Ho AE%sT +Z Australia/Lord_Howe 10:36:20 - LMT 1895 F +10 - AEST 1981 Mar +10:30 LH %z 1985 Jul +10:30 LH %z +Z Australia/Melbourne 9:39:52 - LMT 1895 F +10 AU AE%sT 1971 +10 AV AE%sT +Z Australia/Perth 7:43:24 - LMT 1895 D +8 AU AW%sT 1943 Jul +8 AW AW%sT +Z Australia/Sydney 10:4:52 - LMT 1895 F +10 AU AE%sT 1971 +10 AN AE%sT +Z Etc/GMT 0 - GMT +Z Etc/GMT+1 -1 - %z +Z Etc/GMT+10 -10 - %z +Z Etc/GMT+11 -11 - %z +Z Etc/GMT+12 -12 - %z +Z Etc/GMT+2 -2 - %z +Z Etc/GMT+3 -3 - %z +Z Etc/GMT+4 -4 - %z +Z Etc/GMT+5 -5 - %z +Z Etc/GMT+6 -6 - %z +Z Etc/GMT+7 -7 - %z +Z Etc/GMT+8 -8 - %z +Z Etc/GMT+9 -9 - %z +Z Etc/GMT-1 1 - %z +Z Etc/GMT-10 10 - %z +Z Etc/GMT-11 11 - %z +Z Etc/GMT-12 12 - %z +Z Etc/GMT-13 13 - %z +Z Etc/GMT-14 14 - %z +Z Etc/GMT-2 2 - %z +Z Etc/GMT-3 3 - %z +Z Etc/GMT-4 4 - %z +Z Etc/GMT-5 5 - %z +Z Etc/GMT-6 6 - %z +Z Etc/GMT-7 7 - %z +Z Etc/GMT-8 8 - %z +Z Etc/GMT-9 9 - %z +Z Etc/UTC 0 - UTC +Z Europe/Andorra 0:6:4 - LMT 1901 +0 - WET 1946 S 30 +1 - CET 1985 Mar 31 2 +1 E CE%sT +Z Europe/Astrakhan 3:12:12 - LMT 1924 May +3 - %z 1930 Jun 21 +4 R %z 1989 Mar 26 2s +3 R %z 1991 Mar 31 2s +4 - %z 1992 Mar 29 2s +3 R %z 2011 Mar 27 2s +4 - %z 2014 O 26 2s +3 - %z 2016 Mar 27 2s +4 - %z +Z Europe/Athens 1:34:52 - LMT 1895 S 14 +1:34:52 - AMT 1916 Jul 28 0:1 +2 g EE%sT 1941 Ap 30 +1 g CE%sT 1944 Ap 4 +2 g EE%sT 1981 +2 E EE%sT +Z Europe/Belgrade 1:22 - LMT 1884 +1 - CET 1941 Ap 18 23 +1 c CE%sT 1945 +1 - CET 1945 May 8 2s +1 1 CEST 1945 S 16 2s +1 - CET 1982 N 27 +1 E CE%sT +Z Europe/Berlin 0:53:28 - LMT 1893 Ap +1 c CE%sT 1945 May 24 2 +1 So CE%sT 1946 +1 DE CE%sT 1980 +1 E CE%sT +Z Europe/Brussels 0:17:30 - LMT 1880 +0:17:30 - BMT 1892 May 1 0:17:30 +0 - WET 1914 N 8 +1 - CET 1916 May +1 c CE%sT 1918 N 11 11u +0 b WE%sT 1940 May 20 2s +1 c CE%sT 1944 S 3 +1 b CE%sT 1977 +1 E CE%sT +Z Europe/Bucharest 1:44:24 - LMT 1891 O +1:44:24 - BMT 1931 Jul 24 +2 z EE%sT 1981 Mar 29 2s +2 c EE%sT 1991 +2 z EE%sT 1994 +2 e EE%sT 1997 +2 E EE%sT +Z Europe/Budapest 1:16:20 - LMT 1890 N +1 c CE%sT 1918 +1 h CE%sT 1941 Ap 7 23 +1 c CE%sT 1945 +1 h CE%sT 1984 +1 E CE%sT +Z Europe/Chisinau 1:55:20 - LMT 1880 +1:55 - CMT 1918 F 15 +1:44:24 - BMT 1931 Jul 24 +2 z EE%sT 1940 Au 15 +2 1 EEST 1941 Jul 17 +1 c CE%sT 1944 Au 24 +3 R MSK/MSD 1990 May 6 2 +2 R EE%sT 1992 +2 e EE%sT 1997 +2 MD EE%sT +Z Europe/Dublin -0:25:21 - LMT 1880 Au 2 +-0:25:21 - DMT 1916 May 21 2s +-0:25:21 1 IST 1916 O 1 2s +0 G %s 1921 D 6 +0 G GMT/IST 1940 F 25 2s +0 1 IST 1946 O 6 2s +0 - GMT 1947 Mar 16 2s +0 1 IST 1947 N 2 2s +0 - GMT 1948 Ap 18 2s +0 G GMT/IST 1968 O 27 +1 IE IST/GMT +Z Europe/Gibraltar -0:21:24 - LMT 1880 Au 2 +0 G %s 1957 Ap 14 2 +1 - CET 1982 +1 E CE%sT +Z Europe/Helsinki 1:39:49 - LMT 1878 May 31 +1:39:49 - HMT 1921 May +2 FI EE%sT 1983 +2 E EE%sT +Z Europe/Istanbul 1:55:52 - LMT 1880 +1:56:56 - IMT 1910 O +2 T EE%sT 1978 Jun 29 +3 T %z 1984 N 1 2 +2 T EE%sT 2007 +2 E EE%sT 2011 Mar 27 1u +2 - EET 2011 Mar 28 1u +2 E EE%sT 2014 Mar 30 1u +2 - EET 2014 Mar 31 1u +2 E EE%sT 2015 O 25 1u +2 1 EEST 2015 N 8 1u +2 E EE%sT 2016 S 7 +3 - %z +Z Europe/Kaliningrad 1:22 - LMT 1893 Ap +1 c CE%sT 1945 Ap 10 +2 O EE%sT 1946 Ap 7 +3 R MSK/MSD 1989 Mar 26 2s +2 R EE%sT 2011 Mar 27 2s +3 - %z 2014 O 26 2s +2 - EET +Z Europe/Kirov 3:18:48 - LMT 1919 Jul 1 0u +3 - %z 1930 Jun 21 +4 R %z 1989 Mar 26 2s +3 R MSK/MSD 1991 Mar 31 2s +4 - %z 1992 Mar 29 2s +3 R MSK/MSD 2011 Mar 27 2s +4 - MSK 2014 O 26 2s +3 - MSK +Z Europe/Kyiv 2:2:4 - LMT 1880 +2:2:4 - KMT 1924 May 2 +2 - EET 1930 Jun 21 +3 - MSK 1941 S 20 +1 c CE%sT 1943 N 6 +3 R MSK/MSD 1990 Jul 1 2 +2 1 EEST 1991 S 29 3 +2 c EE%sT 1996 May 13 +2 E EE%sT +Z Europe/Lisbon -0:36:45 - LMT 1884 +-0:36:45 - LMT 1912 Ja 1 0u +0 p WE%sT 1966 O 2 2s +1 - CET 1976 S 26 1 +0 p WE%sT 1986 +0 E WE%sT 1992 S 27 1u +1 E CE%sT 1996 Mar 31 1u +0 E WE%sT +Z Europe/London -0:1:15 - LMT 1847 D +0 G %s 1968 O 27 +1 - BST 1971 O 31 2u +0 G %s 1996 +0 E GMT/BST +Z Europe/Madrid -0:14:44 - LMT 1901 Ja 1 0u +0 s WE%sT 1940 Mar 16 23 +1 s CE%sT 1979 +1 E CE%sT +Z Europe/Malta 0:58:4 - LMT 1893 N 2 +1 I CE%sT 1973 Mar 31 +1 MT CE%sT 1981 +1 E CE%sT +Z Europe/Minsk 1:50:16 - LMT 1880 +1:50 - MMT 1924 May 2 +2 - EET 1930 Jun 21 +3 - MSK 1941 Jun 28 +1 c CE%sT 1944 Jul 3 +3 R MSK/MSD 1990 +3 - MSK 1991 Mar 31 2s +2 R EE%sT 2011 Mar 27 2s +3 - %z +Z Europe/Moscow 2:30:17 - LMT 1880 +2:30:17 - MMT 1916 Jul 3 +2:31:19 R %s 1919 Jul 1 0u +3 R %s 1921 O +3 R MSK/MSD 1922 O +2 - EET 1930 Jun 21 +3 R MSK/MSD 1991 Mar 31 2s +2 R EE%sT 1992 Ja 19 2s +3 R MSK/MSD 2011 Mar 27 2s +4 - MSK 2014 O 26 2s +3 - MSK +Z Europe/Paris 0:9:21 - LMT 1891 Mar 16 +0:9:21 - PMT 1911 Mar 11 +0 F WE%sT 1940 Jun 14 23 +1 c CE%sT 1944 Au 25 +0 F WE%sT 1945 S 16 3 +1 F CE%sT 1977 +1 E CE%sT +Z Europe/Prague 0:57:44 - LMT 1850 +0:57:44 - PMT 1891 O +1 c CE%sT 1945 May 9 +1 CZ CE%sT 1946 D 1 3 +1 -1 GMT 1947 F 23 2 +1 CZ CE%sT 1979 +1 E CE%sT +Z Europe/Riga 1:36:34 - LMT 1880 +1:36:34 - RMT 1918 Ap 15 2 +1:36:34 1 LST 1918 S 16 3 +1:36:34 - RMT 1919 Ap 1 2 +1:36:34 1 LST 1919 May 22 3 +1:36:34 - RMT 1926 May 11 +2 - EET 1940 Au 5 +3 - MSK 1941 Jul +1 c CE%sT 1944 O 13 +3 R MSK/MSD 1989 Mar lastSu 2s +2 1 EEST 1989 S lastSu 2s +2 LV EE%sT 1997 Ja 21 +2 E EE%sT 2000 F 29 +2 - EET 2001 Ja 2 +2 E EE%sT +Z Europe/Rome 0:49:56 - LMT 1866 D 12 +0:49:56 - RMT 1893 O 31 23u +1 I CE%sT 1943 S 10 +1 c CE%sT 1944 Jun 4 +1 I CE%sT 1980 +1 E CE%sT +Z Europe/Samara 3:20:20 - LMT 1919 Jul 1 0u +3 - %z 1930 Jun 21 +4 - %z 1935 Ja 27 +4 R %z 1989 Mar 26 2s +3 R %z 1991 Mar 31 2s +2 R %z 1991 S 29 2s +3 - %z 1991 O 20 3 +4 R %z 2010 Mar 28 2s +3 R %z 2011 Mar 27 2s +4 - %z +Z Europe/Saratov 3:4:18 - LMT 1919 Jul 1 0u +3 - %z 1930 Jun 21 +4 R %z 1988 Mar 27 2s +3 R %z 1991 Mar 31 2s +4 - %z 1992 Mar 29 2s +3 R %z 2011 Mar 27 2s +4 - %z 2014 O 26 2s +3 - %z 2016 D 4 2s +4 - %z +Z Europe/Simferopol 2:16:24 - LMT 1880 +2:16 - SMT 1924 May 2 +2 - EET 1930 Jun 21 +3 - MSK 1941 N +1 c CE%sT 1944 Ap 13 +3 R MSK/MSD 1990 +3 - MSK 1990 Jul 1 2 +2 - EET 1992 Mar 20 +2 c EE%sT 1994 May +3 c MSK/MSD 1996 Mar 31 0s +3 1 MSD 1996 O 27 3s +3 - MSK 1997 Mar lastSu 1u +2 E EE%sT 2014 Mar 30 2 +4 - MSK 2014 O 26 2s +3 - MSK +Z Europe/Sofia 1:33:16 - LMT 1880 +1:56:56 - IMT 1894 N 30 +2 - EET 1942 N 2 3 +1 c CE%sT 1945 +1 - CET 1945 Ap 2 3 +2 - EET 1979 Mar 31 23 +2 BG EE%sT 1982 S 26 3 +2 c EE%sT 1991 +2 e EE%sT 1997 +2 E EE%sT +Z Europe/Tallinn 1:39 - LMT 1880 +1:39 - TMT 1918 F +1 c CE%sT 1919 Jul +1:39 - TMT 1921 May +2 - EET 1940 Au 6 +3 - MSK 1941 S 15 +1 c CE%sT 1944 S 22 +3 R MSK/MSD 1989 Mar 26 2s +2 1 EEST 1989 S 24 2s +2 c EE%sT 1998 S 22 +2 E EE%sT 1999 O 31 4 +2 - EET 2002 F 21 +2 E EE%sT +Z Europe/Tirane 1:19:20 - LMT 1914 +1 - CET 1940 Jun 16 +1 q CE%sT 1984 Jul +1 E CE%sT +Z Europe/Ulyanovsk 3:13:36 - LMT 1919 Jul 1 0u +3 - %z 1930 Jun 21 +4 R %z 1989 Mar 26 2s +3 R %z 1991 Mar 31 2s +2 R %z 1992 Ja 19 2s +3 R %z 2011 Mar 27 2s +4 - %z 2014 O 26 2s +3 - %z 2016 Mar 27 2s +4 - %z +Z Europe/Vienna 1:5:21 - LMT 1893 Ap +1 c CE%sT 1920 +1 a CE%sT 1940 Ap 1 2s +1 c CE%sT 1945 Ap 2 2s +1 1 CEST 1945 Ap 12 2s +1 - CET 1946 +1 a CE%sT 1981 +1 E CE%sT +Z Europe/Vilnius 1:41:16 - LMT 1880 +1:24 - WMT 1917 +1:35:36 - KMT 1919 O 10 +1 - CET 1920 Jul 12 +2 - EET 1920 O 9 +1 - CET 1940 Au 3 +3 - MSK 1941 Jun 24 +1 c CE%sT 1944 Au +3 R MSK/MSD 1989 Mar 26 2s +2 R EE%sT 1991 S 29 2s +2 c EE%sT 1998 +2 - EET 1998 Mar 29 1u +1 E CE%sT 1999 O 31 1u +2 - EET 2003 +2 E EE%sT +Z Europe/Volgograd 2:57:40 - LMT 1920 Ja 3 +3 - %z 1930 Jun 21 +4 - %z 1961 N 11 +4 R %z 1988 Mar 27 2s +3 R MSK/MSD 1991 Mar 31 2s +4 - %z 1992 Mar 29 2s +3 R MSK/MSD 2011 Mar 27 2s +4 - MSK 2014 O 26 2s +3 - MSK 2018 O 28 2s +4 - %z 2020 D 27 2s +3 - MSK +Z Europe/Warsaw 1:24 - LMT 1880 +1:24 - WMT 1915 Au 5 +1 c CE%sT 1918 S 16 3 +2 O EE%sT 1922 Jun +1 O CE%sT 1940 Jun 23 2 +1 c CE%sT 1944 O +1 O CE%sT 1977 +1 W- CE%sT 1988 +1 E CE%sT +Z Europe/Zurich 0:34:8 - LMT 1853 Jul 16 +0:29:46 - BMT 1894 Jun +1 CH CE%sT 1981 +1 E CE%sT +Z Factory 0 - -00 +Z Indian/Chagos 4:49:40 - LMT 1907 +5 - %z 1996 +6 - %z +Z Indian/Maldives 4:54 - LMT 1880 +4:54 - MMT 1960 +5 - %z +Z Indian/Mauritius 3:50 - LMT 1907 +4 MU %z +Z Pacific/Apia 12:33:4 - LMT 1892 Jul 5 +-11:26:56 - LMT 1911 +-11:30 - %z 1950 +-11 WS %z 2011 D 29 24 +13 WS %z +Z Pacific/Auckland 11:39:4 - LMT 1868 N 2 +11:30 NZ NZ%sT 1946 +12 NZ NZ%sT +Z Pacific/Bougainville 10:22:16 - LMT 1880 +9:48:32 - PMMT 1895 +10 - %z 1942 Jul +9 - %z 1945 Au 21 +10 - %z 2014 D 28 2 +11 - %z +Z Pacific/Chatham 12:13:48 - LMT 1868 N 2 +12:15 - %z 1946 +12:45 k %z +Z Pacific/Easter -7:17:28 - LMT 1890 +-7:17:28 - EMT 1932 S +-7 x %z 1982 Mar 14 3u +-6 x %z +Z Pacific/Efate 11:13:16 - LMT 1912 Ja 13 +11 VU %z +Z Pacific/Fakaofo -11:24:56 - LMT 1901 +-11 - %z 2011 D 30 +13 - %z +Z Pacific/Fiji 11:55:44 - LMT 1915 O 26 +12 FJ %z +Z Pacific/Galapagos -5:58:24 - LMT 1931 +-5 - %z 1986 +-6 EC %z +Z Pacific/Gambier -8:59:48 - LMT 1912 O +-9 - %z +Z Pacific/Guadalcanal 10:39:48 - LMT 1912 O +11 - %z +Z Pacific/Guam -14:21 - LMT 1844 D 31 +9:39 - LMT 1901 +10 - GST 1941 D 10 +9 - %z 1944 Jul 31 +10 Gu G%sT 2000 D 23 +10 - ChST +Z Pacific/Honolulu -10:31:26 - LMT 1896 Ja 13 12 +-10:30 - HST 1933 Ap 30 2 +-10:30 1 HDT 1933 May 21 12 +-10:30 u H%sT 1947 Jun 8 2 +-10 - HST +Z Pacific/Kanton 0 - -00 1937 Au 31 +-12 - %z 1979 O +-11 - %z 1994 D 31 +13 - %z +Z Pacific/Kiritimati -10:29:20 - LMT 1901 +-10:40 - %z 1979 O +-10 - %z 1994 D 31 +14 - %z +Z Pacific/Kosrae -13:8:4 - LMT 1844 D 31 +10:51:56 - LMT 1901 +11 - %z 1914 O +9 - %z 1919 F +11 - %z 1937 +10 - %z 1941 Ap +9 - %z 1945 Au +11 - %z 1969 O +12 - %z 1999 +11 - %z +Z Pacific/Kwajalein 11:9:20 - LMT 1901 +11 - %z 1937 +10 - %z 1941 Ap +9 - %z 1944 F 6 +11 - %z 1969 O +-12 - %z 1993 Au 20 24 +12 - %z +Z Pacific/Marquesas -9:18 - LMT 1912 O +-9:30 - %z +Z Pacific/Nauru 11:7:40 - LMT 1921 Ja 15 +11:30 - %z 1942 Au 29 +9 - %z 1945 S 8 +11:30 - %z 1979 F 10 2 +12 - %z +Z Pacific/Niue -11:19:40 - LMT 1952 O 16 +-11:20 - %z 1964 Jul +-11 - %z +Z Pacific/Norfolk 11:11:52 - LMT 1901 +11:12 - %z 1951 +11:30 - %z 1974 O 27 2s +11:30 1 %z 1975 Mar 2 2s +11:30 - %z 2015 O 4 2s +11 - %z 2019 Jul +11 AN %z +Z Pacific/Noumea 11:5:48 - LMT 1912 Ja 13 +11 NC %z +Z Pacific/Pago_Pago 12:37:12 - LMT 1892 Jul 5 +-11:22:48 - LMT 1911 +-11 - SST +Z Pacific/Palau -15:2:4 - LMT 1844 D 31 +8:57:56 - LMT 1901 +9 - %z +Z Pacific/Pitcairn -8:40:20 - LMT 1901 +-8:30 - %z 1998 Ap 27 +-8 - %z +Z Pacific/Port_Moresby 9:48:40 - LMT 1880 +9:48:32 - PMMT 1895 +10 - %z +Z Pacific/Rarotonga 13:20:56 - LMT 1899 D 26 +-10:39:4 - LMT 1952 O 16 +-10:30 - %z 1978 N 12 +-10 CK %z +Z Pacific/Tahiti -9:58:16 - LMT 1912 O +-10 - %z +Z Pacific/Tarawa 11:32:4 - LMT 1901 +12 - %z +Z Pacific/Tongatapu 12:19:12 - LMT 1945 S 10 +12:20 - %z 1961 +13 - %z 1999 +13 TO %z +L Etc/GMT GMT +L Australia/Sydney Australia/ACT +L Australia/Lord_Howe Australia/LHI +L Australia/Sydney Australia/NSW +L Australia/Darwin Australia/North +L Australia/Brisbane Australia/Queensland +L Australia/Adelaide Australia/South +L Australia/Hobart Australia/Tasmania +L Australia/Melbourne Australia/Victoria +L Australia/Perth Australia/West +L Australia/Broken_Hill Australia/Yancowinna +L America/Rio_Branco Brazil/Acre +L America/Noronha Brazil/DeNoronha +L America/Sao_Paulo Brazil/East +L America/Manaus Brazil/West +L Europe/Brussels CET +L America/Chicago CST6CDT +L America/Halifax Canada/Atlantic +L America/Winnipeg Canada/Central +L America/Toronto Canada/Eastern +L America/Edmonton Canada/Mountain +L America/St_Johns Canada/Newfoundland +L America/Vancouver Canada/Pacific +L America/Regina Canada/Saskatchewan +L America/Whitehorse Canada/Yukon +L America/Santiago Chile/Continental +L Pacific/Easter Chile/EasterIsland +L America/Havana Cuba +L Europe/Athens EET +L America/Panama EST +L America/New_York EST5EDT +L Africa/Cairo Egypt +L Europe/Dublin Eire +L Etc/GMT Etc/GMT+0 +L Etc/GMT Etc/GMT-0 +L Etc/GMT Etc/GMT0 +L Etc/GMT Etc/Greenwich +L Etc/UTC Etc/UCT +L Etc/UTC Etc/Universal +L Etc/UTC Etc/Zulu +L Europe/London GB +L Europe/London GB-Eire +L Etc/GMT GMT+0 +L Etc/GMT GMT-0 +L Etc/GMT GMT0 +L Etc/GMT Greenwich +L Asia/Hong_Kong Hongkong +L Africa/Abidjan Iceland +L Asia/Tehran Iran +L Asia/Jerusalem Israel +L America/Jamaica Jamaica +L Asia/Tokyo Japan +L Pacific/Kwajalein Kwajalein +L Africa/Tripoli Libya +L Europe/Brussels MET +L America/Phoenix MST +L America/Denver MST7MDT +L America/Tijuana Mexico/BajaNorte +L America/Mazatlan Mexico/BajaSur +L America/Mexico_City Mexico/General +L Pacific/Auckland NZ +L Pacific/Chatham NZ-CHAT +L America/Denver Navajo +L Asia/Shanghai PRC +L Europe/Warsaw Poland +L Europe/Lisbon Portugal +L Asia/Taipei ROC +L Asia/Seoul ROK +L Asia/Singapore Singapore +L Europe/Istanbul Turkey +L Etc/UTC UCT +L America/Anchorage US/Alaska +L America/Adak US/Aleutian +L America/Phoenix US/Arizona +L America/Chicago US/Central +L America/Indiana/Indianapolis US/East-Indiana +L America/New_York US/Eastern +L Pacific/Honolulu US/Hawaii +L America/Indiana/Knox US/Indiana-Starke +L America/Detroit US/Michigan +L America/Denver US/Mountain +L America/Los_Angeles US/Pacific +L Pacific/Pago_Pago US/Samoa +L Etc/UTC UTC +L Etc/UTC Universal +L Europe/Moscow W-SU +L Etc/UTC Zulu +L America/Argentina/Buenos_Aires America/Buenos_Aires +L America/Argentina/Catamarca America/Catamarca +L America/Argentina/Cordoba America/Cordoba +L America/Indiana/Indianapolis America/Indianapolis +L America/Argentina/Jujuy America/Jujuy +L America/Indiana/Knox America/Knox_IN +L America/Kentucky/Louisville America/Louisville +L America/Argentina/Mendoza America/Mendoza +L America/Puerto_Rico America/Virgin +L Pacific/Pago_Pago Pacific/Samoa +L Africa/Abidjan Africa/Accra +L Africa/Nairobi Africa/Addis_Ababa +L Africa/Nairobi Africa/Asmara +L Africa/Abidjan Africa/Bamako +L Africa/Lagos Africa/Bangui +L Africa/Abidjan Africa/Banjul +L Africa/Maputo Africa/Blantyre +L Africa/Lagos Africa/Brazzaville +L Africa/Maputo Africa/Bujumbura +L Africa/Abidjan Africa/Conakry +L Africa/Abidjan Africa/Dakar +L Africa/Nairobi Africa/Dar_es_Salaam +L Africa/Nairobi Africa/Djibouti +L Africa/Lagos Africa/Douala +L Africa/Abidjan Africa/Freetown +L Africa/Maputo Africa/Gaborone +L Africa/Maputo Africa/Harare +L Africa/Nairobi Africa/Kampala +L Africa/Maputo Africa/Kigali +L Africa/Lagos Africa/Kinshasa +L Africa/Lagos Africa/Libreville +L Africa/Abidjan Africa/Lome +L Africa/Lagos Africa/Luanda +L Africa/Maputo Africa/Lubumbashi +L Africa/Maputo Africa/Lusaka +L Africa/Lagos Africa/Malabo +L Africa/Johannesburg Africa/Maseru +L Africa/Johannesburg Africa/Mbabane +L Africa/Nairobi Africa/Mogadishu +L Africa/Lagos Africa/Niamey +L Africa/Abidjan Africa/Nouakchott +L Africa/Abidjan Africa/Ouagadougou +L Africa/Lagos Africa/Porto-Novo +L America/Puerto_Rico America/Anguilla +L America/Puerto_Rico America/Antigua +L America/Puerto_Rico America/Aruba +L America/Panama America/Atikokan +L America/Puerto_Rico America/Blanc-Sablon +L America/Panama America/Cayman +L America/Phoenix America/Creston +L America/Puerto_Rico America/Curacao +L America/Puerto_Rico America/Dominica +L America/Puerto_Rico America/Grenada +L America/Puerto_Rico America/Guadeloupe +L America/Puerto_Rico America/Kralendijk +L America/Puerto_Rico America/Lower_Princes +L America/Puerto_Rico America/Marigot +L America/Puerto_Rico America/Montserrat +L America/Toronto America/Nassau +L America/Puerto_Rico America/Port_of_Spain +L America/Puerto_Rico America/St_Barthelemy +L America/Puerto_Rico America/St_Kitts +L America/Puerto_Rico America/St_Lucia +L America/Puerto_Rico America/St_Thomas +L America/Puerto_Rico America/St_Vincent +L America/Puerto_Rico America/Tortola +L Pacific/Port_Moresby Antarctica/DumontDUrville +L Pacific/Auckland Antarctica/McMurdo +L Asia/Riyadh Antarctica/Syowa +L Europe/Berlin Arctic/Longyearbyen +L Asia/Riyadh Asia/Aden +L Asia/Qatar Asia/Bahrain +L Asia/Kuching Asia/Brunei +L Asia/Singapore Asia/Kuala_Lumpur +L Asia/Riyadh Asia/Kuwait +L Asia/Dubai Asia/Muscat +L Asia/Bangkok Asia/Phnom_Penh +L Asia/Bangkok Asia/Vientiane +L Africa/Abidjan Atlantic/Reykjavik +L Africa/Abidjan Atlantic/St_Helena +L Europe/Brussels Europe/Amsterdam +L Europe/Prague Europe/Bratislava +L Europe/Zurich Europe/Busingen +L Europe/Berlin Europe/Copenhagen +L Europe/London Europe/Guernsey +L Europe/London Europe/Isle_of_Man +L Europe/London Europe/Jersey +L Europe/Belgrade Europe/Ljubljana +L Europe/Brussels Europe/Luxembourg +L Europe/Helsinki Europe/Mariehamn +L Europe/Paris Europe/Monaco +L Europe/Berlin Europe/Oslo +L Europe/Belgrade Europe/Podgorica +L Europe/Rome Europe/San_Marino +L Europe/Belgrade Europe/Sarajevo +L Europe/Belgrade Europe/Skopje +L Europe/Berlin Europe/Stockholm +L Europe/Zurich Europe/Vaduz +L Europe/Rome Europe/Vatican +L Europe/Belgrade Europe/Zagreb +L Africa/Nairobi Indian/Antananarivo +L Asia/Bangkok Indian/Christmas +L Asia/Yangon Indian/Cocos +L Africa/Nairobi Indian/Comoro +L Indian/Maldives Indian/Kerguelen +L Asia/Dubai Indian/Mahe +L Africa/Nairobi Indian/Mayotte +L Asia/Dubai Indian/Reunion +L Pacific/Port_Moresby Pacific/Chuuk +L Pacific/Tarawa Pacific/Funafuti +L Pacific/Tarawa Pacific/Majuro +L Pacific/Pago_Pago Pacific/Midway +L Pacific/Guadalcanal Pacific/Pohnpei +L Pacific/Guam Pacific/Saipan +L Pacific/Tarawa Pacific/Wake +L Pacific/Tarawa Pacific/Wallis +L Africa/Abidjan Africa/Timbuktu +L America/Argentina/Catamarca America/Argentina/ComodRivadavia +L America/Adak America/Atka +L America/Panama America/Coral_Harbour +L America/Tijuana America/Ensenada +L America/Indiana/Indianapolis America/Fort_Wayne +L America/Toronto America/Montreal +L America/Toronto America/Nipigon +L America/Iqaluit America/Pangnirtung +L America/Rio_Branco America/Porto_Acre +L America/Winnipeg America/Rainy_River +L America/Argentina/Cordoba America/Rosario +L America/Tijuana America/Santa_Isabel +L America/Denver America/Shiprock +L America/Toronto America/Thunder_Bay +L America/Edmonton America/Yellowknife +L Pacific/Auckland Antarctica/South_Pole +L Asia/Ulaanbaatar Asia/Choibalsan +L Asia/Shanghai Asia/Chongqing +L Asia/Shanghai Asia/Harbin +L Asia/Urumqi Asia/Kashgar +L Asia/Jerusalem Asia/Tel_Aviv +L Europe/Berlin Atlantic/Jan_Mayen +L Australia/Sydney Australia/Canberra +L Australia/Hobart Australia/Currie +L Europe/London Europe/Belfast +L Europe/Chisinau Europe/Tiraspol +L Europe/Kyiv Europe/Uzhgorod +L Europe/Kyiv Europe/Zaporozhye +L Pacific/Kanton Pacific/Enderbury +L Pacific/Honolulu Pacific/Johnston +L Pacific/Port_Moresby Pacific/Yap +L Europe/Lisbon WET +L Africa/Nairobi Africa/Asmera +L America/Nuuk America/Godthab +L Asia/Ashgabat Asia/Ashkhabad +L Asia/Kolkata Asia/Calcutta +L Asia/Shanghai Asia/Chungking +L Asia/Dhaka Asia/Dacca +L Europe/Istanbul Asia/Istanbul +L Asia/Kathmandu Asia/Katmandu +L Asia/Macau Asia/Macao +L Asia/Yangon Asia/Rangoon +L Asia/Ho_Chi_Minh Asia/Saigon +L Asia/Thimphu Asia/Thimbu +L Asia/Makassar Asia/Ujung_Pandang +L Asia/Ulaanbaatar Asia/Ulan_Bator +L Atlantic/Faroe Atlantic/Faeroe +L Europe/Kyiv Europe/Kiev +L Asia/Nicosia Europe/Nicosia +L Pacific/Honolulu HST +L America/Los_Angeles PST8PDT +L Pacific/Guadalcanal Pacific/Ponape +L Pacific/Port_Moresby Pacific/Truk diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/zone.tab b/lib/python3.12/site-packages/pytz/zoneinfo/zone.tab new file mode 100644 index 0000000000000000000000000000000000000000..2626b0550341a0605087f33cac6952d5fbb24e67 --- /dev/null +++ b/lib/python3.12/site-packages/pytz/zoneinfo/zone.tab @@ -0,0 +1,448 @@ +# tzdb timezone descriptions (deprecated version) +# +# This file is in the public domain, so clarified as of +# 2009-05-17 by Arthur David Olson. +# +# From Paul Eggert (2021-09-20): +# This file is intended as a backward-compatibility aid for older programs. +# New programs should use zone1970.tab. This file is like zone1970.tab (see +# zone1970.tab's comments), but with the following additional restrictions: +# +# 1. This file contains only ASCII characters. +# 2. The first data column contains exactly one country code. +# +# Because of (2), each row stands for an area that is the intersection +# of a region identified by a country code and of a timezone where civil +# clocks have agreed since 1970; this is a narrower definition than +# that of zone1970.tab. +# +# Unlike zone1970.tab, a row's third column can be a Link from +# 'backward' instead of a Zone. +# +# This table is intended as an aid for users, to help them select timezones +# appropriate for their practical needs. It is not intended to take or +# endorse any position on legal or territorial claims. +# +#country- +#code coordinates TZ comments +AD +4230+00131 Europe/Andorra +AE +2518+05518 Asia/Dubai +AF +3431+06912 Asia/Kabul +AG +1703-06148 America/Antigua +AI +1812-06304 America/Anguilla +AL +4120+01950 Europe/Tirane +AM +4011+04430 Asia/Yerevan +AO -0848+01314 Africa/Luanda +AQ -7750+16636 Antarctica/McMurdo New Zealand time - McMurdo, South Pole +AQ -6617+11031 Antarctica/Casey Casey +AQ -6835+07758 Antarctica/Davis Davis +AQ -6640+14001 Antarctica/DumontDUrville Dumont-d'Urville +AQ -6736+06253 Antarctica/Mawson Mawson +AQ -6448-06406 Antarctica/Palmer Palmer +AQ -6734-06808 Antarctica/Rothera Rothera +AQ -690022+0393524 Antarctica/Syowa Syowa +AQ -720041+0023206 Antarctica/Troll Troll +AQ -7824+10654 Antarctica/Vostok Vostok +AR -3436-05827 America/Argentina/Buenos_Aires Buenos Aires (BA, CF) +AR -3124-06411 America/Argentina/Cordoba Argentina (most areas: CB, CC, CN, ER, FM, MN, SE, SF) +AR -2447-06525 America/Argentina/Salta Salta (SA, LP, NQ, RN) +AR -2411-06518 America/Argentina/Jujuy Jujuy (JY) +AR -2649-06513 America/Argentina/Tucuman Tucuman (TM) +AR -2828-06547 America/Argentina/Catamarca Catamarca (CT), Chubut (CH) +AR -2926-06651 America/Argentina/La_Rioja La Rioja (LR) +AR -3132-06831 America/Argentina/San_Juan San Juan (SJ) +AR -3253-06849 America/Argentina/Mendoza Mendoza (MZ) +AR -3319-06621 America/Argentina/San_Luis San Luis (SL) +AR -5138-06913 America/Argentina/Rio_Gallegos Santa Cruz (SC) +AR -5448-06818 America/Argentina/Ushuaia Tierra del Fuego (TF) +AS -1416-17042 Pacific/Pago_Pago +AT +4813+01620 Europe/Vienna +AU -3133+15905 Australia/Lord_Howe Lord Howe Island +AU -5430+15857 Antarctica/Macquarie Macquarie Island +AU -4253+14719 Australia/Hobart Tasmania +AU -3749+14458 Australia/Melbourne Victoria +AU -3352+15113 Australia/Sydney New South Wales (most areas) +AU -3157+14127 Australia/Broken_Hill New South Wales (Yancowinna) +AU -2728+15302 Australia/Brisbane Queensland (most areas) +AU -2016+14900 Australia/Lindeman Queensland (Whitsunday Islands) +AU -3455+13835 Australia/Adelaide South Australia +AU -1228+13050 Australia/Darwin Northern Territory +AU -3157+11551 Australia/Perth Western Australia (most areas) +AU -3143+12852 Australia/Eucla Western Australia (Eucla) +AW +1230-06958 America/Aruba +AX +6006+01957 Europe/Mariehamn +AZ +4023+04951 Asia/Baku +BA +4352+01825 Europe/Sarajevo +BB +1306-05937 America/Barbados +BD +2343+09025 Asia/Dhaka +BE +5050+00420 Europe/Brussels +BF +1222-00131 Africa/Ouagadougou +BG +4241+02319 Europe/Sofia +BH +2623+05035 Asia/Bahrain +BI -0323+02922 Africa/Bujumbura +BJ +0629+00237 Africa/Porto-Novo +BL +1753-06251 America/St_Barthelemy +BM +3217-06446 Atlantic/Bermuda +BN +0456+11455 Asia/Brunei +BO -1630-06809 America/La_Paz +BQ +120903-0681636 America/Kralendijk +BR -0351-03225 America/Noronha Atlantic islands +BR -0127-04829 America/Belem Para (east), Amapa +BR -0343-03830 America/Fortaleza Brazil (northeast: MA, PI, CE, RN, PB) +BR -0803-03454 America/Recife Pernambuco +BR -0712-04812 America/Araguaina Tocantins +BR -0940-03543 America/Maceio Alagoas, Sergipe +BR -1259-03831 America/Bahia Bahia +BR -2332-04637 America/Sao_Paulo Brazil (southeast: GO, DF, MG, ES, RJ, SP, PR, SC, RS) +BR -2027-05437 America/Campo_Grande Mato Grosso do Sul +BR -1535-05605 America/Cuiaba Mato Grosso +BR -0226-05452 America/Santarem Para (west) +BR -0846-06354 America/Porto_Velho Rondonia +BR +0249-06040 America/Boa_Vista Roraima +BR -0308-06001 America/Manaus Amazonas (east) +BR -0640-06952 America/Eirunepe Amazonas (west) +BR -0958-06748 America/Rio_Branco Acre +BS +2505-07721 America/Nassau +BT +2728+08939 Asia/Thimphu +BW -2439+02555 Africa/Gaborone +BY +5354+02734 Europe/Minsk +BZ +1730-08812 America/Belize +CA +4734-05243 America/St_Johns Newfoundland, Labrador (SE) +CA +4439-06336 America/Halifax Atlantic - NS (most areas), PE +CA +4612-05957 America/Glace_Bay Atlantic - NS (Cape Breton) +CA +4606-06447 America/Moncton Atlantic - New Brunswick +CA +5320-06025 America/Goose_Bay Atlantic - Labrador (most areas) +CA +5125-05707 America/Blanc-Sablon AST - QC (Lower North Shore) +CA +4339-07923 America/Toronto Eastern - ON & QC (most areas) +CA +6344-06828 America/Iqaluit Eastern - NU (most areas) +CA +484531-0913718 America/Atikokan EST - ON (Atikokan), NU (Coral H) +CA +4953-09709 America/Winnipeg Central - ON (west), Manitoba +CA +744144-0944945 America/Resolute Central - NU (Resolute) +CA +624900-0920459 America/Rankin_Inlet Central - NU (central) +CA +5024-10439 America/Regina CST - SK (most areas) +CA +5017-10750 America/Swift_Current CST - SK (midwest) +CA +5333-11328 America/Edmonton Mountain - AB, BC(E), NT(E), SK(W) +CA +690650-1050310 America/Cambridge_Bay Mountain - NU (west) +CA +682059-1334300 America/Inuvik Mountain - NT (west) +CA +4906-11631 America/Creston MST - BC (Creston) +CA +5546-12014 America/Dawson_Creek MST - BC (Dawson Cr, Ft St John) +CA +5848-12242 America/Fort_Nelson MST - BC (Ft Nelson) +CA +6043-13503 America/Whitehorse MST - Yukon (east) +CA +6404-13925 America/Dawson MST - Yukon (west) +CA +4916-12307 America/Vancouver Pacific - BC (most areas) +CC -1210+09655 Indian/Cocos +CD -0418+01518 Africa/Kinshasa Dem. Rep. of Congo (west) +CD -1140+02728 Africa/Lubumbashi Dem. Rep. of Congo (east) +CF +0422+01835 Africa/Bangui +CG -0416+01517 Africa/Brazzaville +CH +4723+00832 Europe/Zurich +CI +0519-00402 Africa/Abidjan +CK -2114-15946 Pacific/Rarotonga +CL -3327-07040 America/Santiago most of Chile +CL -4534-07204 America/Coyhaique Aysen Region +CL -5309-07055 America/Punta_Arenas Magallanes Region +CL -2709-10926 Pacific/Easter Easter Island +CM +0403+00942 Africa/Douala +CN +3114+12128 Asia/Shanghai Beijing Time +CN +4348+08735 Asia/Urumqi Xinjiang Time +CO +0436-07405 America/Bogota +CR +0956-08405 America/Costa_Rica +CU +2308-08222 America/Havana +CV +1455-02331 Atlantic/Cape_Verde +CW +1211-06900 America/Curacao +CX -1025+10543 Indian/Christmas +CY +3510+03322 Asia/Nicosia most of Cyprus +CY +3507+03357 Asia/Famagusta Northern Cyprus +CZ +5005+01426 Europe/Prague +DE +5230+01322 Europe/Berlin most of Germany +DE +4742+00841 Europe/Busingen Busingen +DJ +1136+04309 Africa/Djibouti +DK +5540+01235 Europe/Copenhagen +DM +1518-06124 America/Dominica +DO +1828-06954 America/Santo_Domingo +DZ +3647+00303 Africa/Algiers +EC -0210-07950 America/Guayaquil Ecuador (mainland) +EC -0054-08936 Pacific/Galapagos Galapagos Islands +EE +5925+02445 Europe/Tallinn +EG +3003+03115 Africa/Cairo +EH +2709-01312 Africa/El_Aaiun +ER +1520+03853 Africa/Asmara +ES +4024-00341 Europe/Madrid Spain (mainland) +ES +3553-00519 Africa/Ceuta Ceuta, Melilla +ES +2806-01524 Atlantic/Canary Canary Islands +ET +0902+03842 Africa/Addis_Ababa +FI +6010+02458 Europe/Helsinki +FJ -1808+17825 Pacific/Fiji +FK -5142-05751 Atlantic/Stanley +FM +0725+15147 Pacific/Chuuk Chuuk/Truk, Yap +FM +0658+15813 Pacific/Pohnpei Pohnpei/Ponape +FM +0519+16259 Pacific/Kosrae Kosrae +FO +6201-00646 Atlantic/Faroe +FR +4852+00220 Europe/Paris +GA +0023+00927 Africa/Libreville +GB +513030-0000731 Europe/London +GD +1203-06145 America/Grenada +GE +4143+04449 Asia/Tbilisi +GF +0456-05220 America/Cayenne +GG +492717-0023210 Europe/Guernsey +GH +0533-00013 Africa/Accra +GI +3608-00521 Europe/Gibraltar +GL +6411-05144 America/Nuuk most of Greenland +GL +7646-01840 America/Danmarkshavn National Park (east coast) +GL +7029-02158 America/Scoresbysund Scoresbysund/Ittoqqortoormiit +GL +7634-06847 America/Thule Thule/Pituffik +GM +1328-01639 Africa/Banjul +GN +0931-01343 Africa/Conakry +GP +1614-06132 America/Guadeloupe +GQ +0345+00847 Africa/Malabo +GR +3758+02343 Europe/Athens +GS -5416-03632 Atlantic/South_Georgia +GT +1438-09031 America/Guatemala +GU +1328+14445 Pacific/Guam +GW +1151-01535 Africa/Bissau +GY +0648-05810 America/Guyana +HK +2217+11409 Asia/Hong_Kong +HN +1406-08713 America/Tegucigalpa +HR +4548+01558 Europe/Zagreb +HT +1832-07220 America/Port-au-Prince +HU +4730+01905 Europe/Budapest +ID -0610+10648 Asia/Jakarta Java, Sumatra +ID -0002+10920 Asia/Pontianak Borneo (west, central) +ID -0507+11924 Asia/Makassar Borneo (east, south), Sulawesi/Celebes, Bali, Nusa Tengarra, Timor (west) +ID -0232+14042 Asia/Jayapura New Guinea (West Papua / Irian Jaya), Malukus/Moluccas +IE +5320-00615 Europe/Dublin +IL +314650+0351326 Asia/Jerusalem +IM +5409-00428 Europe/Isle_of_Man +IN +2232+08822 Asia/Kolkata +IO -0720+07225 Indian/Chagos +IQ +3321+04425 Asia/Baghdad +IR +3540+05126 Asia/Tehran +IS +6409-02151 Atlantic/Reykjavik +IT +4154+01229 Europe/Rome +JE +491101-0020624 Europe/Jersey +JM +175805-0764736 America/Jamaica +JO +3157+03556 Asia/Amman +JP +353916+1394441 Asia/Tokyo +KE -0117+03649 Africa/Nairobi +KG +4254+07436 Asia/Bishkek +KH +1133+10455 Asia/Phnom_Penh +KI +0125+17300 Pacific/Tarawa Gilbert Islands +KI -0247-17143 Pacific/Kanton Phoenix Islands +KI +0152-15720 Pacific/Kiritimati Line Islands +KM -1141+04316 Indian/Comoro +KN +1718-06243 America/St_Kitts +KP +3901+12545 Asia/Pyongyang +KR +3733+12658 Asia/Seoul +KW +2920+04759 Asia/Kuwait +KY +1918-08123 America/Cayman +KZ +4315+07657 Asia/Almaty most of Kazakhstan +KZ +4448+06528 Asia/Qyzylorda Qyzylorda/Kyzylorda/Kzyl-Orda +KZ +5312+06337 Asia/Qostanay Qostanay/Kostanay/Kustanay +KZ +5017+05710 Asia/Aqtobe Aqtobe/Aktobe +KZ +4431+05016 Asia/Aqtau Mangghystau/Mankistau +KZ +4707+05156 Asia/Atyrau Atyrau/Atirau/Gur'yev +KZ +5113+05121 Asia/Oral West Kazakhstan +LA +1758+10236 Asia/Vientiane +LB +3353+03530 Asia/Beirut +LC +1401-06100 America/St_Lucia +LI +4709+00931 Europe/Vaduz +LK +0656+07951 Asia/Colombo +LR +0618-01047 Africa/Monrovia +LS -2928+02730 Africa/Maseru +LT +5441+02519 Europe/Vilnius +LU +4936+00609 Europe/Luxembourg +LV +5657+02406 Europe/Riga +LY +3254+01311 Africa/Tripoli +MA +3339-00735 Africa/Casablanca +MC +4342+00723 Europe/Monaco +MD +4700+02850 Europe/Chisinau +ME +4226+01916 Europe/Podgorica +MF +1804-06305 America/Marigot +MG -1855+04731 Indian/Antananarivo +MH +0709+17112 Pacific/Majuro most of Marshall Islands +MH +0905+16720 Pacific/Kwajalein Kwajalein +MK +4159+02126 Europe/Skopje +ML +1239-00800 Africa/Bamako +MM +1647+09610 Asia/Yangon +MN +4755+10653 Asia/Ulaanbaatar most of Mongolia +MN +4801+09139 Asia/Hovd Bayan-Olgii, Hovd, Uvs +MO +221150+1133230 Asia/Macau +MP +1512+14545 Pacific/Saipan +MQ +1436-06105 America/Martinique +MR +1806-01557 Africa/Nouakchott +MS +1643-06213 America/Montserrat +MT +3554+01431 Europe/Malta +MU -2010+05730 Indian/Mauritius +MV +0410+07330 Indian/Maldives +MW -1547+03500 Africa/Blantyre +MX +1924-09909 America/Mexico_City Central Mexico +MX +2105-08646 America/Cancun Quintana Roo +MX +2058-08937 America/Merida Campeche, Yucatan +MX +2540-10019 America/Monterrey Durango; Coahuila, Nuevo Leon, Tamaulipas (most areas) +MX +2550-09730 America/Matamoros Coahuila, Nuevo Leon, Tamaulipas (US border) +MX +2838-10605 America/Chihuahua Chihuahua (most areas) +MX +3144-10629 America/Ciudad_Juarez Chihuahua (US border - west) +MX +2934-10425 America/Ojinaga Chihuahua (US border - east) +MX +2313-10625 America/Mazatlan Baja California Sur, Nayarit (most areas), Sinaloa +MX +2048-10515 America/Bahia_Banderas Bahia de Banderas +MX +2904-11058 America/Hermosillo Sonora +MX +3232-11701 America/Tijuana Baja California +MY +0310+10142 Asia/Kuala_Lumpur Malaysia (peninsula) +MY +0133+11020 Asia/Kuching Sabah, Sarawak +MZ -2558+03235 Africa/Maputo +NA -2234+01706 Africa/Windhoek +NC -2216+16627 Pacific/Noumea +NE +1331+00207 Africa/Niamey +NF -2903+16758 Pacific/Norfolk +NG +0627+00324 Africa/Lagos +NI +1209-08617 America/Managua +NL +5222+00454 Europe/Amsterdam +NO +5955+01045 Europe/Oslo +NP +2743+08519 Asia/Kathmandu +NR -0031+16655 Pacific/Nauru +NU -1901-16955 Pacific/Niue +NZ -3652+17446 Pacific/Auckland most of New Zealand +NZ -4357-17633 Pacific/Chatham Chatham Islands +OM +2336+05835 Asia/Muscat +PA +0858-07932 America/Panama +PE -1203-07703 America/Lima +PF -1732-14934 Pacific/Tahiti Society Islands +PF -0900-13930 Pacific/Marquesas Marquesas Islands +PF -2308-13457 Pacific/Gambier Gambier Islands +PG -0930+14710 Pacific/Port_Moresby most of Papua New Guinea +PG -0613+15534 Pacific/Bougainville Bougainville +PH +143512+1205804 Asia/Manila +PK +2452+06703 Asia/Karachi +PL +5215+02100 Europe/Warsaw +PM +4703-05620 America/Miquelon +PN -2504-13005 Pacific/Pitcairn +PR +182806-0660622 America/Puerto_Rico +PS +3130+03428 Asia/Gaza Gaza Strip +PS +313200+0350542 Asia/Hebron West Bank +PT +3843-00908 Europe/Lisbon Portugal (mainland) +PT +3238-01654 Atlantic/Madeira Madeira Islands +PT +3744-02540 Atlantic/Azores Azores +PW +0720+13429 Pacific/Palau +PY -2516-05740 America/Asuncion +QA +2517+05132 Asia/Qatar +RE -2052+05528 Indian/Reunion +RO +4426+02606 Europe/Bucharest +RS +4450+02030 Europe/Belgrade +RU +5443+02030 Europe/Kaliningrad MSK-01 - Kaliningrad +RU +554521+0373704 Europe/Moscow MSK+00 - Moscow area +# The obsolescent zone.tab format cannot represent Europe/Simferopol well. +# Put it in RU section and list as UA. See "territorial claims" above. +# Programs should use zone1970.tab instead; see above. +UA +4457+03406 Europe/Simferopol Crimea +RU +5836+04939 Europe/Kirov MSK+00 - Kirov +RU +4844+04425 Europe/Volgograd MSK+00 - Volgograd +RU +4621+04803 Europe/Astrakhan MSK+01 - Astrakhan +RU +5134+04602 Europe/Saratov MSK+01 - Saratov +RU +5420+04824 Europe/Ulyanovsk MSK+01 - Ulyanovsk +RU +5312+05009 Europe/Samara MSK+01 - Samara, Udmurtia +RU +5651+06036 Asia/Yekaterinburg MSK+02 - Urals +RU +5500+07324 Asia/Omsk MSK+03 - Omsk +RU +5502+08255 Asia/Novosibirsk MSK+04 - Novosibirsk +RU +5322+08345 Asia/Barnaul MSK+04 - Altai +RU +5630+08458 Asia/Tomsk MSK+04 - Tomsk +RU +5345+08707 Asia/Novokuznetsk MSK+04 - Kemerovo +RU +5601+09250 Asia/Krasnoyarsk MSK+04 - Krasnoyarsk area +RU +5216+10420 Asia/Irkutsk MSK+05 - Irkutsk, Buryatia +RU +5203+11328 Asia/Chita MSK+06 - Zabaykalsky +RU +6200+12940 Asia/Yakutsk MSK+06 - Lena River +RU +623923+1353314 Asia/Khandyga MSK+06 - Tomponsky, Ust-Maysky +RU +4310+13156 Asia/Vladivostok MSK+07 - Amur River +RU +643337+1431336 Asia/Ust-Nera MSK+07 - Oymyakonsky +RU +5934+15048 Asia/Magadan MSK+08 - Magadan +RU +4658+14242 Asia/Sakhalin MSK+08 - Sakhalin Island +RU +6728+15343 Asia/Srednekolymsk MSK+08 - Sakha (E), N Kuril Is +RU +5301+15839 Asia/Kamchatka MSK+09 - Kamchatka +RU +6445+17729 Asia/Anadyr MSK+09 - Bering Sea +RW -0157+03004 Africa/Kigali +SA +2438+04643 Asia/Riyadh +SB -0932+16012 Pacific/Guadalcanal +SC -0440+05528 Indian/Mahe +SD +1536+03232 Africa/Khartoum +SE +5920+01803 Europe/Stockholm +SG +0117+10351 Asia/Singapore +SH -1555-00542 Atlantic/St_Helena +SI +4603+01431 Europe/Ljubljana +SJ +7800+01600 Arctic/Longyearbyen +SK +4809+01707 Europe/Bratislava +SL +0830-01315 Africa/Freetown +SM +4355+01228 Europe/San_Marino +SN +1440-01726 Africa/Dakar +SO +0204+04522 Africa/Mogadishu +SR +0550-05510 America/Paramaribo +SS +0451+03137 Africa/Juba +ST +0020+00644 Africa/Sao_Tome +SV +1342-08912 America/El_Salvador +SX +180305-0630250 America/Lower_Princes +SY +3330+03618 Asia/Damascus +SZ -2618+03106 Africa/Mbabane +TC +2128-07108 America/Grand_Turk +TD +1207+01503 Africa/Ndjamena +TF -492110+0701303 Indian/Kerguelen +TG +0608+00113 Africa/Lome +TH +1345+10031 Asia/Bangkok +TJ +3835+06848 Asia/Dushanbe +TK -0922-17114 Pacific/Fakaofo +TL -0833+12535 Asia/Dili +TM +3757+05823 Asia/Ashgabat +TN +3648+01011 Africa/Tunis +TO -210800-1751200 Pacific/Tongatapu +TR +4101+02858 Europe/Istanbul +TT +1039-06131 America/Port_of_Spain +TV -0831+17913 Pacific/Funafuti +TW +2503+12130 Asia/Taipei +TZ -0648+03917 Africa/Dar_es_Salaam +UA +5026+03031 Europe/Kyiv most of Ukraine +UG +0019+03225 Africa/Kampala +UM +2813-17722 Pacific/Midway Midway Islands +UM +1917+16637 Pacific/Wake Wake Island +US +404251-0740023 America/New_York Eastern (most areas) +US +421953-0830245 America/Detroit Eastern - MI (most areas) +US +381515-0854534 America/Kentucky/Louisville Eastern - KY (Louisville area) +US +364947-0845057 America/Kentucky/Monticello Eastern - KY (Wayne) +US +394606-0860929 America/Indiana/Indianapolis Eastern - IN (most areas) +US +384038-0873143 America/Indiana/Vincennes Eastern - IN (Da, Du, K, Mn) +US +410305-0863611 America/Indiana/Winamac Eastern - IN (Pulaski) +US +382232-0862041 America/Indiana/Marengo Eastern - IN (Crawford) +US +382931-0871643 America/Indiana/Petersburg Eastern - IN (Pike) +US +384452-0850402 America/Indiana/Vevay Eastern - IN (Switzerland) +US +415100-0873900 America/Chicago Central (most areas) +US +375711-0864541 America/Indiana/Tell_City Central - IN (Perry) +US +411745-0863730 America/Indiana/Knox Central - IN (Starke) +US +450628-0873651 America/Menominee Central - MI (Wisconsin border) +US +470659-1011757 America/North_Dakota/Center Central - ND (Oliver) +US +465042-1012439 America/North_Dakota/New_Salem Central - ND (Morton rural) +US +471551-1014640 America/North_Dakota/Beulah Central - ND (Mercer) +US +394421-1045903 America/Denver Mountain (most areas) +US +433649-1161209 America/Boise Mountain - ID (south), OR (east) +US +332654-1120424 America/Phoenix MST - AZ (except Navajo) +US +340308-1181434 America/Los_Angeles Pacific +US +611305-1495401 America/Anchorage Alaska (most areas) +US +581807-1342511 America/Juneau Alaska - Juneau area +US +571035-1351807 America/Sitka Alaska - Sitka area +US +550737-1313435 America/Metlakatla Alaska - Annette Island +US +593249-1394338 America/Yakutat Alaska - Yakutat +US +643004-1652423 America/Nome Alaska (west) +US +515248-1763929 America/Adak Alaska - western Aleutians +US +211825-1575130 Pacific/Honolulu Hawaii +UY -345433-0561245 America/Montevideo +UZ +3940+06648 Asia/Samarkand Uzbekistan (west) +UZ +4120+06918 Asia/Tashkent Uzbekistan (east) +VA +415408+0122711 Europe/Vatican +VC +1309-06114 America/St_Vincent +VE +1030-06656 America/Caracas +VG +1827-06437 America/Tortola +VI +1821-06456 America/St_Thomas +VN +1045+10640 Asia/Ho_Chi_Minh +VU -1740+16825 Pacific/Efate +WF -1318-17610 Pacific/Wallis +WS -1350-17144 Pacific/Apia +YE +1245+04512 Asia/Aden +YT -1247+04514 Indian/Mayotte +ZA -2615+02800 Africa/Johannesburg +ZM -1525+02817 Africa/Lusaka +ZW -1750+03103 Africa/Harare diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/zone1970.tab b/lib/python3.12/site-packages/pytz/zoneinfo/zone1970.tab new file mode 100644 index 0000000000000000000000000000000000000000..36535bdf5cfbdedfc70b4d7899de7351b8eb5070 --- /dev/null +++ b/lib/python3.12/site-packages/pytz/zoneinfo/zone1970.tab @@ -0,0 +1,375 @@ +# tzdb timezone descriptions +# +# This file is in the public domain. +# +# From Paul Eggert (2018-06-27): +# This file contains a table where each row stands for a timezone where +# civil timestamps have agreed since 1970. Columns are separated by +# a single tab. Lines beginning with '#' are comments. All text uses +# UTF-8 encoding. The columns of the table are as follows: +# +# 1. The countries that overlap the timezone, as a comma-separated list +# of ISO 3166 2-character country codes. See the file 'iso3166.tab'. +# 2. Latitude and longitude of the timezone's principal location +# in ISO 6709 sign-degrees-minutes-seconds format, +# either ±DDMM±DDDMM or ±DDMMSS±DDDMMSS, +# first latitude (+ is north), then longitude (+ is east). +# 3. Timezone name used in value of TZ environment variable. +# Please see the theory.html file for how these names are chosen. +# If multiple timezones overlap a country, each has a row in the +# table, with each column 1 containing the country code. +# 4. Comments; present if and only if countries have multiple timezones, +# and useful only for those countries. For example, the comments +# for the row with countries CH,DE,LI and name Europe/Zurich +# are useful only for DE, since CH and LI have no other timezones. +# +# If a timezone covers multiple countries, the most-populous city is used, +# and that country is listed first in column 1; any other countries +# are listed alphabetically by country code. The table is sorted +# first by country code, then (if possible) by an order within the +# country that (1) makes some geographical sense, and (2) puts the +# most populous timezones first, where that does not contradict (1). +# +# This table is intended as an aid for users, to help them select timezones +# appropriate for their practical needs. It is not intended to take or +# endorse any position on legal or territorial claims. +# +#country- +#codes coordinates TZ comments +AD +4230+00131 Europe/Andorra +AE,OM,RE,SC,TF +2518+05518 Asia/Dubai Crozet +AF +3431+06912 Asia/Kabul +AL +4120+01950 Europe/Tirane +AM +4011+04430 Asia/Yerevan +AQ -6617+11031 Antarctica/Casey Casey +AQ -6835+07758 Antarctica/Davis Davis +AQ -6736+06253 Antarctica/Mawson Mawson +AQ -6448-06406 Antarctica/Palmer Palmer +AQ -6734-06808 Antarctica/Rothera Rothera +AQ -720041+0023206 Antarctica/Troll Troll +AQ -7824+10654 Antarctica/Vostok Vostok +AR -3436-05827 America/Argentina/Buenos_Aires Buenos Aires (BA, CF) +AR -3124-06411 America/Argentina/Cordoba most areas: CB, CC, CN, ER, FM, MN, SE, SF +AR -2447-06525 America/Argentina/Salta Salta (SA, LP, NQ, RN) +AR -2411-06518 America/Argentina/Jujuy Jujuy (JY) +AR -2649-06513 America/Argentina/Tucuman Tucumán (TM) +AR -2828-06547 America/Argentina/Catamarca Catamarca (CT), Chubut (CH) +AR -2926-06651 America/Argentina/La_Rioja La Rioja (LR) +AR -3132-06831 America/Argentina/San_Juan San Juan (SJ) +AR -3253-06849 America/Argentina/Mendoza Mendoza (MZ) +AR -3319-06621 America/Argentina/San_Luis San Luis (SL) +AR -5138-06913 America/Argentina/Rio_Gallegos Santa Cruz (SC) +AR -5448-06818 America/Argentina/Ushuaia Tierra del Fuego (TF) +AS,UM -1416-17042 Pacific/Pago_Pago Midway +AT +4813+01620 Europe/Vienna +AU -3133+15905 Australia/Lord_Howe Lord Howe Island +AU -5430+15857 Antarctica/Macquarie Macquarie Island +AU -4253+14719 Australia/Hobart Tasmania +AU -3749+14458 Australia/Melbourne Victoria +AU -3352+15113 Australia/Sydney New South Wales (most areas) +AU -3157+14127 Australia/Broken_Hill New South Wales (Yancowinna) +AU -2728+15302 Australia/Brisbane Queensland (most areas) +AU -2016+14900 Australia/Lindeman Queensland (Whitsunday Islands) +AU -3455+13835 Australia/Adelaide South Australia +AU -1228+13050 Australia/Darwin Northern Territory +AU -3157+11551 Australia/Perth Western Australia (most areas) +AU -3143+12852 Australia/Eucla Western Australia (Eucla) +AZ +4023+04951 Asia/Baku +BB +1306-05937 America/Barbados +BD +2343+09025 Asia/Dhaka +BE,LU,NL +5050+00420 Europe/Brussels +BG +4241+02319 Europe/Sofia +BM +3217-06446 Atlantic/Bermuda +BO -1630-06809 America/La_Paz +BR -0351-03225 America/Noronha Atlantic islands +BR -0127-04829 America/Belem Pará (east), Amapá +BR -0343-03830 America/Fortaleza Brazil (northeast: MA, PI, CE, RN, PB) +BR -0803-03454 America/Recife Pernambuco +BR -0712-04812 America/Araguaina Tocantins +BR -0940-03543 America/Maceio Alagoas, Sergipe +BR -1259-03831 America/Bahia Bahia +BR -2332-04637 America/Sao_Paulo Brazil (southeast: GO, DF, MG, ES, RJ, SP, PR, SC, RS) +BR -2027-05437 America/Campo_Grande Mato Grosso do Sul +BR -1535-05605 America/Cuiaba Mato Grosso +BR -0226-05452 America/Santarem Pará (west) +BR -0846-06354 America/Porto_Velho Rondônia +BR +0249-06040 America/Boa_Vista Roraima +BR -0308-06001 America/Manaus Amazonas (east) +BR -0640-06952 America/Eirunepe Amazonas (west) +BR -0958-06748 America/Rio_Branco Acre +BT +2728+08939 Asia/Thimphu +BY +5354+02734 Europe/Minsk +BZ +1730-08812 America/Belize +CA +4734-05243 America/St_Johns Newfoundland, Labrador (SE) +CA +4439-06336 America/Halifax Atlantic - NS (most areas), PE +CA +4612-05957 America/Glace_Bay Atlantic - NS (Cape Breton) +CA +4606-06447 America/Moncton Atlantic - New Brunswick +CA +5320-06025 America/Goose_Bay Atlantic - Labrador (most areas) +CA,BS +4339-07923 America/Toronto Eastern - ON & QC (most areas) +CA +6344-06828 America/Iqaluit Eastern - NU (most areas) +CA +4953-09709 America/Winnipeg Central - ON (west), Manitoba +CA +744144-0944945 America/Resolute Central - NU (Resolute) +CA +624900-0920459 America/Rankin_Inlet Central - NU (central) +CA +5024-10439 America/Regina CST - SK (most areas) +CA +5017-10750 America/Swift_Current CST - SK (midwest) +CA +5333-11328 America/Edmonton Mountain - AB, BC(E), NT(E), SK(W) +CA +690650-1050310 America/Cambridge_Bay Mountain - NU (west) +CA +682059-1334300 America/Inuvik Mountain - NT (west) +CA +5546-12014 America/Dawson_Creek MST - BC (Dawson Cr, Ft St John) +CA +5848-12242 America/Fort_Nelson MST - BC (Ft Nelson) +CA +6043-13503 America/Whitehorse MST - Yukon (east) +CA +6404-13925 America/Dawson MST - Yukon (west) +CA +4916-12307 America/Vancouver Pacific - BC (most areas) +CH,DE,LI +4723+00832 Europe/Zurich Büsingen +CI,BF,GH,GM,GN,IS,ML,MR,SH,SL,SN,TG +0519-00402 Africa/Abidjan +CK -2114-15946 Pacific/Rarotonga +CL -3327-07040 America/Santiago most of Chile +CL -4534-07204 America/Coyhaique Aysén Region +CL -5309-07055 America/Punta_Arenas Magallanes Region +CL -2709-10926 Pacific/Easter Easter Island +CN +3114+12128 Asia/Shanghai Beijing Time +CN +4348+08735 Asia/Urumqi Xinjiang Time +CO +0436-07405 America/Bogota +CR +0956-08405 America/Costa_Rica +CU +2308-08222 America/Havana +CV +1455-02331 Atlantic/Cape_Verde +CY +3510+03322 Asia/Nicosia most of Cyprus +CY +3507+03357 Asia/Famagusta Northern Cyprus +CZ,SK +5005+01426 Europe/Prague +DE,DK,NO,SE,SJ +5230+01322 Europe/Berlin most of Germany +DO +1828-06954 America/Santo_Domingo +DZ +3647+00303 Africa/Algiers +EC -0210-07950 America/Guayaquil Ecuador (mainland) +EC -0054-08936 Pacific/Galapagos Galápagos Islands +EE +5925+02445 Europe/Tallinn +EG +3003+03115 Africa/Cairo +EH +2709-01312 Africa/El_Aaiun +ES +4024-00341 Europe/Madrid Spain (mainland) +ES +3553-00519 Africa/Ceuta Ceuta, Melilla +ES +2806-01524 Atlantic/Canary Canary Islands +FI,AX +6010+02458 Europe/Helsinki +FJ -1808+17825 Pacific/Fiji +FK -5142-05751 Atlantic/Stanley +FM +0519+16259 Pacific/Kosrae Kosrae +FO +6201-00646 Atlantic/Faroe +FR,MC +4852+00220 Europe/Paris +GB,GG,IM,JE +513030-0000731 Europe/London +GE +4143+04449 Asia/Tbilisi +GF +0456-05220 America/Cayenne +GI +3608-00521 Europe/Gibraltar +GL +6411-05144 America/Nuuk most of Greenland +GL +7646-01840 America/Danmarkshavn National Park (east coast) +GL +7029-02158 America/Scoresbysund Scoresbysund/Ittoqqortoormiit +GL +7634-06847 America/Thule Thule/Pituffik +GR +3758+02343 Europe/Athens +GS -5416-03632 Atlantic/South_Georgia +GT +1438-09031 America/Guatemala +GU,MP +1328+14445 Pacific/Guam +GW +1151-01535 Africa/Bissau +GY +0648-05810 America/Guyana +HK +2217+11409 Asia/Hong_Kong +HN +1406-08713 America/Tegucigalpa +HT +1832-07220 America/Port-au-Prince +HU +4730+01905 Europe/Budapest +ID -0610+10648 Asia/Jakarta Java, Sumatra +ID -0002+10920 Asia/Pontianak Borneo (west, central) +ID -0507+11924 Asia/Makassar Borneo (east, south), Sulawesi/Celebes, Bali, Nusa Tengarra, Timor (west) +ID -0232+14042 Asia/Jayapura New Guinea (West Papua / Irian Jaya), Malukus/Moluccas +IE +5320-00615 Europe/Dublin +IL +314650+0351326 Asia/Jerusalem +IN +2232+08822 Asia/Kolkata +IO -0720+07225 Indian/Chagos +IQ +3321+04425 Asia/Baghdad +IR +3540+05126 Asia/Tehran +IT,SM,VA +4154+01229 Europe/Rome +JM +175805-0764736 America/Jamaica +JO +3157+03556 Asia/Amman +JP,AU +353916+1394441 Asia/Tokyo Eyre Bird Observatory +KE,DJ,ER,ET,KM,MG,SO,TZ,UG,YT -0117+03649 Africa/Nairobi +KG +4254+07436 Asia/Bishkek +KI,MH,TV,UM,WF +0125+17300 Pacific/Tarawa Gilberts, Marshalls, Wake +KI -0247-17143 Pacific/Kanton Phoenix Islands +KI +0152-15720 Pacific/Kiritimati Line Islands +KP +3901+12545 Asia/Pyongyang +KR +3733+12658 Asia/Seoul +KZ +4315+07657 Asia/Almaty most of Kazakhstan +KZ +4448+06528 Asia/Qyzylorda Qyzylorda/Kyzylorda/Kzyl-Orda +KZ +5312+06337 Asia/Qostanay Qostanay/Kostanay/Kustanay +KZ +5017+05710 Asia/Aqtobe Aqtöbe/Aktobe +KZ +4431+05016 Asia/Aqtau Mangghystaū/Mankistau +KZ +4707+05156 Asia/Atyrau Atyraū/Atirau/Gur'yev +KZ +5113+05121 Asia/Oral West Kazakhstan +LB +3353+03530 Asia/Beirut +LK +0656+07951 Asia/Colombo +LR +0618-01047 Africa/Monrovia +LT +5441+02519 Europe/Vilnius +LV +5657+02406 Europe/Riga +LY +3254+01311 Africa/Tripoli +MA +3339-00735 Africa/Casablanca +MD +4700+02850 Europe/Chisinau +MH +0905+16720 Pacific/Kwajalein Kwajalein +MM,CC +1647+09610 Asia/Yangon +MN +4755+10653 Asia/Ulaanbaatar most of Mongolia +MN +4801+09139 Asia/Hovd Bayan-Ölgii, Hovd, Uvs +MO +221150+1133230 Asia/Macau +MQ +1436-06105 America/Martinique +MT +3554+01431 Europe/Malta +MU -2010+05730 Indian/Mauritius +MV,TF +0410+07330 Indian/Maldives Kerguelen, St Paul I, Amsterdam I +MX +1924-09909 America/Mexico_City Central Mexico +MX +2105-08646 America/Cancun Quintana Roo +MX +2058-08937 America/Merida Campeche, Yucatán +MX +2540-10019 America/Monterrey Durango; Coahuila, Nuevo León, Tamaulipas (most areas) +MX +2550-09730 America/Matamoros Coahuila, Nuevo León, Tamaulipas (US border) +MX +2838-10605 America/Chihuahua Chihuahua (most areas) +MX +3144-10629 America/Ciudad_Juarez Chihuahua (US border - west) +MX +2934-10425 America/Ojinaga Chihuahua (US border - east) +MX +2313-10625 America/Mazatlan Baja California Sur, Nayarit (most areas), Sinaloa +MX +2048-10515 America/Bahia_Banderas Bahía de Banderas +MX +2904-11058 America/Hermosillo Sonora +MX +3232-11701 America/Tijuana Baja California +MY,BN +0133+11020 Asia/Kuching Sabah, Sarawak +MZ,BI,BW,CD,MW,RW,ZM,ZW -2558+03235 Africa/Maputo Central Africa Time +NA -2234+01706 Africa/Windhoek +NC -2216+16627 Pacific/Noumea +NF -2903+16758 Pacific/Norfolk +NG,AO,BJ,CD,CF,CG,CM,GA,GQ,NE +0627+00324 Africa/Lagos West Africa Time +NI +1209-08617 America/Managua +NP +2743+08519 Asia/Kathmandu +NR -0031+16655 Pacific/Nauru +NU -1901-16955 Pacific/Niue +NZ,AQ -3652+17446 Pacific/Auckland New Zealand time +NZ -4357-17633 Pacific/Chatham Chatham Islands +PA,CA,KY +0858-07932 America/Panama EST - ON (Atikokan), NU (Coral H) +PE -1203-07703 America/Lima +PF -1732-14934 Pacific/Tahiti Society Islands +PF -0900-13930 Pacific/Marquesas Marquesas Islands +PF -2308-13457 Pacific/Gambier Gambier Islands +PG,AQ,FM -0930+14710 Pacific/Port_Moresby Papua New Guinea (most areas), Chuuk, Yap, Dumont d'Urville +PG -0613+15534 Pacific/Bougainville Bougainville +PH +143512+1205804 Asia/Manila +PK +2452+06703 Asia/Karachi +PL +5215+02100 Europe/Warsaw +PM +4703-05620 America/Miquelon +PN -2504-13005 Pacific/Pitcairn +PR,AG,CA,AI,AW,BL,BQ,CW,DM,GD,GP,KN,LC,MF,MS,SX,TT,VC,VG,VI +182806-0660622 America/Puerto_Rico AST - QC (Lower North Shore) +PS +3130+03428 Asia/Gaza Gaza Strip +PS +313200+0350542 Asia/Hebron West Bank +PT +3843-00908 Europe/Lisbon Portugal (mainland) +PT +3238-01654 Atlantic/Madeira Madeira Islands +PT +3744-02540 Atlantic/Azores Azores +PW +0720+13429 Pacific/Palau +PY -2516-05740 America/Asuncion +QA,BH +2517+05132 Asia/Qatar +RO +4426+02606 Europe/Bucharest +RS,BA,HR,ME,MK,SI +4450+02030 Europe/Belgrade +RU +5443+02030 Europe/Kaliningrad MSK-01 - Kaliningrad +RU +554521+0373704 Europe/Moscow MSK+00 - Moscow area +# Mention RU and UA alphabetically. See "territorial claims" above. +RU,UA +4457+03406 Europe/Simferopol Crimea +RU +5836+04939 Europe/Kirov MSK+00 - Kirov +RU +4844+04425 Europe/Volgograd MSK+00 - Volgograd +RU +4621+04803 Europe/Astrakhan MSK+01 - Astrakhan +RU +5134+04602 Europe/Saratov MSK+01 - Saratov +RU +5420+04824 Europe/Ulyanovsk MSK+01 - Ulyanovsk +RU +5312+05009 Europe/Samara MSK+01 - Samara, Udmurtia +RU +5651+06036 Asia/Yekaterinburg MSK+02 - Urals +RU +5500+07324 Asia/Omsk MSK+03 - Omsk +RU +5502+08255 Asia/Novosibirsk MSK+04 - Novosibirsk +RU +5322+08345 Asia/Barnaul MSK+04 - Altai +RU +5630+08458 Asia/Tomsk MSK+04 - Tomsk +RU +5345+08707 Asia/Novokuznetsk MSK+04 - Kemerovo +RU +5601+09250 Asia/Krasnoyarsk MSK+04 - Krasnoyarsk area +RU +5216+10420 Asia/Irkutsk MSK+05 - Irkutsk, Buryatia +RU +5203+11328 Asia/Chita MSK+06 - Zabaykalsky +RU +6200+12940 Asia/Yakutsk MSK+06 - Lena River +RU +623923+1353314 Asia/Khandyga MSK+06 - Tomponsky, Ust-Maysky +RU +4310+13156 Asia/Vladivostok MSK+07 - Amur River +RU +643337+1431336 Asia/Ust-Nera MSK+07 - Oymyakonsky +RU +5934+15048 Asia/Magadan MSK+08 - Magadan +RU +4658+14242 Asia/Sakhalin MSK+08 - Sakhalin Island +RU +6728+15343 Asia/Srednekolymsk MSK+08 - Sakha (E), N Kuril Is +RU +5301+15839 Asia/Kamchatka MSK+09 - Kamchatka +RU +6445+17729 Asia/Anadyr MSK+09 - Bering Sea +SA,AQ,KW,YE +2438+04643 Asia/Riyadh Syowa +SB,FM -0932+16012 Pacific/Guadalcanal Pohnpei +SD +1536+03232 Africa/Khartoum +SG,AQ,MY +0117+10351 Asia/Singapore peninsular Malaysia, Concordia +SR +0550-05510 America/Paramaribo +SS +0451+03137 Africa/Juba +ST +0020+00644 Africa/Sao_Tome +SV +1342-08912 America/El_Salvador +SY +3330+03618 Asia/Damascus +TC +2128-07108 America/Grand_Turk +TD +1207+01503 Africa/Ndjamena +TH,CX,KH,LA,VN +1345+10031 Asia/Bangkok north Vietnam +TJ +3835+06848 Asia/Dushanbe +TK -0922-17114 Pacific/Fakaofo +TL -0833+12535 Asia/Dili +TM +3757+05823 Asia/Ashgabat +TN +3648+01011 Africa/Tunis +TO -210800-1751200 Pacific/Tongatapu +TR +4101+02858 Europe/Istanbul +TW +2503+12130 Asia/Taipei +UA +5026+03031 Europe/Kyiv most of Ukraine +US +404251-0740023 America/New_York Eastern (most areas) +US +421953-0830245 America/Detroit Eastern - MI (most areas) +US +381515-0854534 America/Kentucky/Louisville Eastern - KY (Louisville area) +US +364947-0845057 America/Kentucky/Monticello Eastern - KY (Wayne) +US +394606-0860929 America/Indiana/Indianapolis Eastern - IN (most areas) +US +384038-0873143 America/Indiana/Vincennes Eastern - IN (Da, Du, K, Mn) +US +410305-0863611 America/Indiana/Winamac Eastern - IN (Pulaski) +US +382232-0862041 America/Indiana/Marengo Eastern - IN (Crawford) +US +382931-0871643 America/Indiana/Petersburg Eastern - IN (Pike) +US +384452-0850402 America/Indiana/Vevay Eastern - IN (Switzerland) +US +415100-0873900 America/Chicago Central (most areas) +US +375711-0864541 America/Indiana/Tell_City Central - IN (Perry) +US +411745-0863730 America/Indiana/Knox Central - IN (Starke) +US +450628-0873651 America/Menominee Central - MI (Wisconsin border) +US +470659-1011757 America/North_Dakota/Center Central - ND (Oliver) +US +465042-1012439 America/North_Dakota/New_Salem Central - ND (Morton rural) +US +471551-1014640 America/North_Dakota/Beulah Central - ND (Mercer) +US +394421-1045903 America/Denver Mountain (most areas) +US +433649-1161209 America/Boise Mountain - ID (south), OR (east) +US,CA +332654-1120424 America/Phoenix MST - AZ (most areas), Creston BC +US +340308-1181434 America/Los_Angeles Pacific +US +611305-1495401 America/Anchorage Alaska (most areas) +US +581807-1342511 America/Juneau Alaska - Juneau area +US +571035-1351807 America/Sitka Alaska - Sitka area +US +550737-1313435 America/Metlakatla Alaska - Annette Island +US +593249-1394338 America/Yakutat Alaska - Yakutat +US +643004-1652423 America/Nome Alaska (west) +US +515248-1763929 America/Adak Alaska - western Aleutians +US +211825-1575130 Pacific/Honolulu Hawaii +UY -345433-0561245 America/Montevideo +UZ +3940+06648 Asia/Samarkand Uzbekistan (west) +UZ +4120+06918 Asia/Tashkent Uzbekistan (east) +VE +1030-06656 America/Caracas +VN +1045+10640 Asia/Ho_Chi_Minh south Vietnam +VU -1740+16825 Pacific/Efate +WS -1350-17144 Pacific/Apia +ZA,LS,SZ -2615+02800 Africa/Johannesburg +# +# The next section contains experimental tab-separated comments for +# use by user agents like tzselect that identify continents and oceans. +# +# For example, the comment "#@AQAntarctica/" means the country code +# AQ is in the continent Antarctica regardless of the Zone name, +# so Pacific/Auckland should be listed under Antarctica as well as +# under the Pacific because its line's country codes include AQ. +# +# If more than one country code is affected each is listed separated +# by commas, e.g., #@IS,SHAtlantic/". If a country code is in +# more than one continent or ocean, each is listed separated by +# commas, e.g., the second column of "#@CY,TRAsia/,Europe/". +# +# These experimental comments are present only for country codes where +# the continent or ocean is not already obvious from the Zone name. +# For example, there is no such comment for RU since it already +# corresponds to Zone names starting with both "Europe/" and "Asia/". +# +#@AQ Antarctica/ +#@IS,SH Atlantic/ +#@CY,TR Asia/,Europe/ +#@SJ Arctic/ +#@CC,CX,KM,MG,YT Indian/ diff --git a/lib/python3.12/site-packages/pytz/zoneinfo/zonenow.tab b/lib/python3.12/site-packages/pytz/zoneinfo/zonenow.tab new file mode 100644 index 0000000000000000000000000000000000000000..093f0a0cb7495b5d50751bc0cb5b22df4a67c763 --- /dev/null +++ b/lib/python3.12/site-packages/pytz/zoneinfo/zonenow.tab @@ -0,0 +1,296 @@ +# tzdb timezone descriptions, for users who do not care about old timestamps +# +# This file is in the public domain. +# +# From Paul Eggert (2023-12-18): +# This file contains a table where each row stands for a timezone +# where civil timestamps are predicted to agree from now on. +# This file is like zone1970.tab (see zone1970.tab's comments), +# but with the following changes: +# +# 1. Each timezone corresponds to a set of clocks that are planned +# to agree from now on. This is a larger set of clocks than in +# zone1970.tab, where each timezone's clocks must agree from 1970 on. +# 2. The first column is irrelevant and ignored. +# 3. The table is sorted in a different way: +# first by standard time UTC offset; +# then, if DST is used, by daylight saving UTC offset; +# then by time zone abbreviation. +# 4. Every timezone has a nonempty comments column, with wording +# distinguishing the timezone only from other timezones with the +# same UTC offset at some point during the year. +# +# The format of this table is experimental, and may change in future versions. +# +# This table is intended as an aid for users, to help them select timezones +# appropriate for their practical needs. It is not intended to take or +# endorse any position on legal or territorial claims. +# +#XX coordinates TZ comments +# +# -11 - SST +XX -1416-17042 Pacific/Pago_Pago Midway; Samoa ("SST") +# +# -11 +XX -1901-16955 Pacific/Niue Niue +# +# -10 - HST +XX +211825-1575130 Pacific/Honolulu Hawaii ("HST") +# +# -10 +XX -1732-14934 Pacific/Tahiti Tahiti; Cook Islands +# +# -10/-09 - HST / HDT (North America DST) +XX +515248-1763929 America/Adak western Aleutians in Alaska ("HST/HDT") +# +# -09:30 +XX -0900-13930 Pacific/Marquesas Marquesas +# +# -09 +XX -2308-13457 Pacific/Gambier Gambier +# +# -09/-08 - AKST/AKDT (North America DST) +XX +611305-1495401 America/Anchorage most of Alaska ("AKST/AKDT") +# +# -08 +XX -2504-13005 Pacific/Pitcairn Pitcairn +# +# -08/-07 - PST/PDT (North America DST) +XX +340308-1181434 America/Los_Angeles Pacific ("PST/PDT") - US & Canada; Mexico near US border +# +# -07 - MST +XX +332654-1120424 America/Phoenix Mountain Standard ("MST") - Arizona; western Mexico; Yukon +# +# -07/-06 - MST/MDT (North America DST) +XX +394421-1045903 America/Denver Mountain ("MST/MDT") - US & Canada; Mexico near US border +# +# -06 +XX -0054-08936 Pacific/Galapagos Galápagos +# +# -06 - CST +XX +1924-09909 America/Mexico_City Central Standard ("CST") - Saskatchewan; central Mexico; Central America +# +# -06/-05 (Chile DST) +XX -2709-10926 Pacific/Easter Easter Island +# +# -06/-05 - CST/CDT (North America DST) +XX +415100-0873900 America/Chicago Central ("CST/CDT") - US & Canada; Mexico near US border +# +# -05 +XX -1203-07703 America/Lima eastern South America +# +# -05 - EST +XX +175805-0764736 America/Jamaica Eastern Standard ("EST") - Caymans; Jamaica; eastern Mexico; Panama +# +# -05/-04 - CST/CDT (Cuba DST) +XX +2308-08222 America/Havana Cuba +# +# -05/-04 - EST/EDT (North America DST) +XX +404251-0740023 America/New_York Eastern ("EST/EDT") - US & Canada +# +# -04 +XX +1030-06656 America/Caracas western South America +# +# -04 - AST +XX +1828-06954 America/Santo_Domingo Atlantic Standard ("AST") - eastern Caribbean +# +# -04/-03 (Chile DST) +XX -3327-07040 America/Santiago most of Chile +# +# -04/-03 - AST/ADT (North America DST) +XX +4439-06336 America/Halifax Atlantic ("AST/ADT") - Canada; Bermuda +# +# -03:30/-02:30 - NST/NDT (North America DST) +XX +4734-05243 America/St_Johns Newfoundland ("NST/NDT") +# +# -03 +XX -2332-04637 America/Sao_Paulo eastern and southern South America +# +# -03/-02 (North America DST) +XX +4703-05620 America/Miquelon St Pierre & Miquelon +# +# -02 +XX -0351-03225 America/Noronha Fernando de Noronha; South Georgia +# +# -02/-01 (EU DST) +XX +6411-05144 America/Nuuk most of Greenland +# +# -01 +XX +1455-02331 Atlantic/Cape_Verde Cape Verde +# +# -01/+00 (EU DST) +XX +3744-02540 Atlantic/Azores Azores +# +# +00 - GMT +XX +0519-00402 Africa/Abidjan far western Africa; Iceland ("GMT") +# +# +00/+01 - GMT/BST (EU DST) +XX +513030-0000731 Europe/London United Kingdom ("GMT/BST") +# +# +00/+01 - WET/WEST (EU DST) +XX +3843-00908 Europe/Lisbon western Europe ("WET/WEST") +# +# +00/+02 - Troll DST +XX -720041+0023206 Antarctica/Troll Troll Station in Antarctica +# +# +01 - CET +XX +3647+00303 Africa/Algiers Algeria, Tunisia ("CET") +# +# +01 - WAT +XX +0627+00324 Africa/Lagos western Africa ("WAT") +# +# +01/+00 - IST/GMT (EU DST in reverse) +XX +5320-00615 Europe/Dublin Ireland ("IST/GMT") +# +# +01/+00 - (Morocco DST) +XX +3339-00735 Africa/Casablanca Morocco +# +# +01/+02 - CET/CEST (EU DST) +XX +4852+00220 Europe/Paris central Europe ("CET/CEST") +# +# +02 - CAT +XX -2558+03235 Africa/Maputo central Africa ("CAT") +# +# +02 - EET +XX +3254+01311 Africa/Tripoli Libya; Kaliningrad ("EET") +# +# +02 - SAST +XX -2615+02800 Africa/Johannesburg southern Africa ("SAST") +# +# +02/+03 - EET/EEST (EU DST) +XX +3758+02343 Europe/Athens eastern Europe ("EET/EEST") +# +# +02/+03 - EET/EEST (Egypt DST) +XX +3003+03115 Africa/Cairo Egypt +# +# +02/+03 - EET/EEST (Lebanon DST) +XX +3353+03530 Asia/Beirut Lebanon +# +# +02/+03 - EET/EEST (Moldova DST) +XX +4700+02850 Europe/Chisinau Moldova +# +# +02/+03 - EET/EEST (Palestine DST) +XX +3130+03428 Asia/Gaza Palestine +# +# +02/+03 - IST/IDT (Israel DST) +XX +314650+0351326 Asia/Jerusalem Israel +# +# +03 +XX +4101+02858 Europe/Istanbul Near East; Belarus +# +# +03 - EAT +XX -0117+03649 Africa/Nairobi eastern Africa ("EAT") +# +# +03 - MSK +XX +554521+0373704 Europe/Moscow Moscow ("MSK") +# +# +03:30 +XX +3540+05126 Asia/Tehran Iran +# +# +04 +XX +2518+05518 Asia/Dubai Russia; Caucasus; Persian Gulf; Seychelles; Réunion +# +# +04:30 +XX +3431+06912 Asia/Kabul Afghanistan +# +# +05 +XX +4120+06918 Asia/Tashkent Russia; Kazakhstan; Tajikistan; Turkmenistan; Uzbekistan; Maldives +# +# +05 - PKT +XX +2452+06703 Asia/Karachi Pakistan ("PKT") +# +# +05:30 +XX +0656+07951 Asia/Colombo Sri Lanka +# +# +05:30 - IST +XX +2232+08822 Asia/Kolkata India ("IST") +# +# +05:45 +XX +2743+08519 Asia/Kathmandu Nepal +# +# +06 +XX +2343+09025 Asia/Dhaka Russia; Kyrgyzstan; Bhutan; Bangladesh; Chagos +# +# +06:30 +XX +1647+09610 Asia/Yangon Myanmar; Cocos +# +# +07 +XX +1345+10031 Asia/Bangkok Russia; Indochina; Christmas Island +# +# +07 - WIB +XX -0610+10648 Asia/Jakarta Indonesia ("WIB") +# +# +08 +XX +0117+10351 Asia/Singapore Russia; Brunei; Malaysia; Singapore; Concordia +# +# +08 - AWST +XX -3157+11551 Australia/Perth Western Australia ("AWST") +# +# +08 - CST +XX +3114+12128 Asia/Shanghai China ("CST") +# +# +08 - HKT +XX +2217+11409 Asia/Hong_Kong Hong Kong ("HKT") +# +# +08 - PHT +XX +143512+1205804 Asia/Manila Philippines ("PHT") +# +# +08 - WITA +XX -0507+11924 Asia/Makassar Indonesia ("WITA") +# +# +08:45 +XX -3143+12852 Australia/Eucla Eucla +# +# +09 +XX +5203+11328 Asia/Chita Russia; Palau; East Timor +# +# +09 - JST +XX +353916+1394441 Asia/Tokyo Japan ("JST"); Eyre Bird Observatory +# +# +09 - KST +XX +3733+12658 Asia/Seoul Korea ("KST") +# +# +09 - WIT +XX -0232+14042 Asia/Jayapura Indonesia ("WIT") +# +# +09:30 - ACST +XX -1228+13050 Australia/Darwin Northern Territory ("ACST") +# +# +09:30/+10:30 - ACST/ACDT (Australia DST) +XX -3455+13835 Australia/Adelaide South Australia ("ACST/ACDT") +# +# +10 +XX +4310+13156 Asia/Vladivostok Russia; Yap; Chuuk; Papua New Guinea; Dumont d'Urville +# +# +10 - AEST +XX -2728+15302 Australia/Brisbane Queensland ("AEST") +# +# +10 - ChST +XX +1328+14445 Pacific/Guam Mariana Islands ("ChST") +# +# +10/+11 - AEST/AEDT (Australia DST) +XX -3352+15113 Australia/Sydney southeast Australia ("AEST/AEDT") +# +# +10:30/+11 +XX -3133+15905 Australia/Lord_Howe Lord Howe Island +# +# +11 +XX -0613+15534 Pacific/Bougainville Russia; Kosrae; Bougainville; Solomons +# +# +11/+12 (Australia DST) +XX -2903+16758 Pacific/Norfolk Norfolk Island +# +# +12 +XX +5301+15839 Asia/Kamchatka Russia; Tuvalu; Fiji; etc. +# +# +12/+13 (New Zealand DST) +XX -3652+17446 Pacific/Auckland New Zealand ("NZST/NZDT") +# +# +12:45/+13:45 (Chatham DST) +XX -4357-17633 Pacific/Chatham Chatham Islands +# +# +13 +XX -210800-1751200 Pacific/Tongatapu Kanton; Tokelau; Samoa (western); Tonga +# +# +14 +XX +0152-15720 Pacific/Kiritimati Kiritimati diff --git a/lib/python3.12/site-packages/trl/__init__.py b/lib/python3.12/site-packages/trl/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..19c0380c348c21d6c6a2d6d48a08fe1bb314ca50 --- /dev/null +++ b/lib/python3.12/site-packages/trl/__init__.py @@ -0,0 +1,222 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +__version__ = "0.19.0" + +from typing import TYPE_CHECKING + +from .import_utils import OptionalDependencyNotAvailable, _LazyModule, is_diffusers_available + + +_import_structure = { + "scripts": ["init_zero_verbose", "ScriptArguments", "TrlParser"], + "data_utils": [ + "apply_chat_template", + "extract_prompt", + "is_conversational", + "maybe_apply_chat_template", + "maybe_convert_to_chatml", + "maybe_extract_prompt", + "maybe_unpair_preference_dataset", + "pack_dataset", + "pack_examples", + "truncate_dataset", + "unpair_preference_dataset", + ], + "environment": ["TextEnvironment", "TextHistory"], + "extras": ["BestOfNSampler"], + "models": [ + "SUPPORTED_ARCHITECTURES", + "AutoModelForCausalLMWithValueHead", + "AutoModelForSeq2SeqLMWithValueHead", + "PreTrainedModelWrapper", + "clone_chat_template", + "create_reference_model", + "setup_chat_format", + ], + "trainer": [ + "AlignPropConfig", + "AlignPropTrainer", + "AllTrueJudge", + "BaseBinaryJudge", + "BaseJudge", + "BasePairwiseJudge", + "BaseRankJudge", + "BCOConfig", + "BCOTrainer", + "CPOConfig", + "CPOTrainer", + "DataCollatorForCompletionOnlyLM", + "DPOConfig", + "DPOTrainer", + "FDivergenceConstants", + "FDivergenceType", + "GKDConfig", + "GKDTrainer", + "GRPOConfig", + "GRPOTrainer", + "HfPairwiseJudge", + "IterativeSFTConfig", + "IterativeSFTTrainer", + "KTOConfig", + "KTOTrainer", + "LogCompletionsCallback", + "MergeModelCallback", + "ModelConfig", + "NashMDConfig", + "NashMDTrainer", + "OnlineDPOConfig", + "OnlineDPOTrainer", + "OpenAIPairwiseJudge", + "ORPOConfig", + "ORPOTrainer", + "PairRMJudge", + "PPOConfig", + "PPOTrainer", + "PRMConfig", + "PRMTrainer", + "RewardConfig", + "RewardTrainer", + "RLOOConfig", + "RLOOTrainer", + "SFTConfig", + "SFTTrainer", + "WinRateCallback", + "XPOConfig", + "XPOTrainer", + ], + "trainer.callbacks": ["MergeModelCallback", "RichProgressCallback", "SyncRefModelCallback"], + "trainer.utils": ["get_kbit_device_map", "get_peft_config", "get_quantization_config"], +} + +try: + if not is_diffusers_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["models"].extend( + [ + "DDPOPipelineOutput", + "DDPOSchedulerOutput", + "DDPOStableDiffusionPipeline", + "DefaultDDPOStableDiffusionPipeline", + ] + ) + _import_structure["trainer"].extend(["DDPOConfig", "DDPOTrainer"]) + +if TYPE_CHECKING: + from .data_utils import ( + apply_chat_template, + extract_prompt, + is_conversational, + maybe_apply_chat_template, + maybe_convert_to_chatml, + maybe_extract_prompt, + maybe_unpair_preference_dataset, + pack_dataset, + pack_examples, + truncate_dataset, + unpair_preference_dataset, + ) + from .environment import TextEnvironment, TextHistory + from .extras import BestOfNSampler + from .models import ( + SUPPORTED_ARCHITECTURES, + AutoModelForCausalLMWithValueHead, + AutoModelForSeq2SeqLMWithValueHead, + PreTrainedModelWrapper, + clone_chat_template, + create_reference_model, + setup_chat_format, + ) + from .scripts import ScriptArguments, TrlParser, init_zero_verbose + from .trainer import ( + AlignPropConfig, + AlignPropTrainer, + AllTrueJudge, + BaseBinaryJudge, + BaseJudge, + BasePairwiseJudge, + BaseRankJudge, + BCOConfig, + BCOTrainer, + CPOConfig, + CPOTrainer, + DataCollatorForCompletionOnlyLM, + DPOConfig, + DPOTrainer, + FDivergenceConstants, + FDivergenceType, + GKDConfig, + GKDTrainer, + GRPOConfig, + GRPOTrainer, + HfPairwiseJudge, + IterativeSFTConfig, + IterativeSFTTrainer, + KTOConfig, + KTOTrainer, + LogCompletionsCallback, + MergeModelCallback, + ModelConfig, + NashMDConfig, + NashMDTrainer, + OnlineDPOConfig, + OnlineDPOTrainer, + OpenAIPairwiseJudge, + ORPOConfig, + ORPOTrainer, + PairRMJudge, + PPOConfig, + PPOTrainer, + PRMConfig, + PRMTrainer, + RewardConfig, + RewardTrainer, + RLOOConfig, + RLOOTrainer, + SFTConfig, + SFTTrainer, + WinRateCallback, + XPOConfig, + XPOTrainer, + ) + from .trainer.callbacks import RichProgressCallback, SyncRefModelCallback + from .trainer.utils import get_kbit_device_map, get_peft_config, get_quantization_config + + try: + if not is_diffusers_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .models import ( + DDPOPipelineOutput, + DDPOSchedulerOutput, + DDPOStableDiffusionPipeline, + DefaultDDPOStableDiffusionPipeline, + ) + from .trainer import DDPOConfig, DDPOTrainer + +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + extra_objects={"__version__": __version__}, + ) diff --git a/lib/python3.12/site-packages/trl/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/trl/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..26c46633d936ecd77dddc4f2690844c555f23230 Binary files /dev/null and b/lib/python3.12/site-packages/trl/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/__pycache__/cli.cpython-312.pyc b/lib/python3.12/site-packages/trl/__pycache__/cli.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65c02580a222ab9365fde03ad59f8aaabf804389 Binary files /dev/null and b/lib/python3.12/site-packages/trl/__pycache__/cli.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/__pycache__/core.cpython-312.pyc b/lib/python3.12/site-packages/trl/__pycache__/core.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3837daf4454290c814069f14653b9151f546395e Binary files /dev/null and b/lib/python3.12/site-packages/trl/__pycache__/core.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/__pycache__/data_utils.cpython-312.pyc b/lib/python3.12/site-packages/trl/__pycache__/data_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47f27894b28db3aab92761d7e64bedddd62d2656 Binary files /dev/null and b/lib/python3.12/site-packages/trl/__pycache__/data_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/__pycache__/import_utils.cpython-312.pyc b/lib/python3.12/site-packages/trl/__pycache__/import_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..30924f0980b775ba2d42f5493a8abdf0ed8a3716 Binary files /dev/null and b/lib/python3.12/site-packages/trl/__pycache__/import_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/__pycache__/mergekit_utils.cpython-312.pyc b/lib/python3.12/site-packages/trl/__pycache__/mergekit_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aec75a8e560ee0bca35d2ce3683de3bb4ea814a9 Binary files /dev/null and b/lib/python3.12/site-packages/trl/__pycache__/mergekit_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/accelerate_configs/fsdp1.yaml b/lib/python3.12/site-packages/trl/accelerate_configs/fsdp1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c01b0b567bc93bf87ec136ea975b3793d273a45c --- /dev/null +++ b/lib/python3.12/site-packages/trl/accelerate_configs/fsdp1.yaml @@ -0,0 +1,28 @@ +compute_environment: LOCAL_MACHINE +debug: false +distributed_type: FSDP +downcast_bf16: 'no' +enable_cpu_affinity: false +fsdp_config: + fsdp_activation_checkpointing: false + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_backward_prefetch: BACKWARD_PRE + fsdp_cpu_ram_efficient_loading: true + fsdp_forward_prefetch: true + fsdp_offload_params: false + fsdp_reshard_after_forward: FULL_SHARD + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_sync_module_states: true + fsdp_use_orig_params: true + fsdp_version: 1 +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/lib/python3.12/site-packages/trl/accelerate_configs/fsdp2.yaml b/lib/python3.12/site-packages/trl/accelerate_configs/fsdp2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..af498f3eced9c2434b80113f2f22d40395e0ab8a --- /dev/null +++ b/lib/python3.12/site-packages/trl/accelerate_configs/fsdp2.yaml @@ -0,0 +1,25 @@ +# Requires accelerate 1.7.0 or higher +compute_environment: LOCAL_MACHINE +debug: false +distributed_type: FSDP +downcast_bf16: 'no' +enable_cpu_affinity: false +fsdp_config: + fsdp_activation_checkpointing: false + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_cpu_ram_efficient_loading: true + fsdp_offload_params: false + fsdp_reshard_after_forward: true + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_version: 2 +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/lib/python3.12/site-packages/trl/accelerate_configs/multi_gpu.yaml b/lib/python3.12/site-packages/trl/accelerate_configs/multi_gpu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..15dad9be3ba44f7c934e1ecab98a93cb83cbc79a --- /dev/null +++ b/lib/python3.12/site-packages/trl/accelerate_configs/multi_gpu.yaml @@ -0,0 +1,16 @@ +compute_environment: LOCAL_MACHINE +debug: false +distributed_type: MULTI_GPU +downcast_bf16: 'no' +gpu_ids: all +machine_rank: 0 +main_training_function: main +mixed_precision: 'bf16' +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/lib/python3.12/site-packages/trl/accelerate_configs/single_gpu.yaml b/lib/python3.12/site-packages/trl/accelerate_configs/single_gpu.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ebd00a067118e56f3d63ab0f24827cfea21b24b9 --- /dev/null +++ b/lib/python3.12/site-packages/trl/accelerate_configs/single_gpu.yaml @@ -0,0 +1,16 @@ +compute_environment: LOCAL_MACHINE +debug: false +distributed_type: "NO" +downcast_bf16: 'no' +gpu_ids: all +machine_rank: 0 +main_training_function: main +mixed_precision: 'bf16' +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/lib/python3.12/site-packages/trl/accelerate_configs/zero1.yaml b/lib/python3.12/site-packages/trl/accelerate_configs/zero1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d5b5f782fb30f9fcbcc8fc58262f09eaf2e10368 --- /dev/null +++ b/lib/python3.12/site-packages/trl/accelerate_configs/zero1.yaml @@ -0,0 +1,20 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + deepspeed_multinode_launcher: standard + gradient_accumulation_steps: 1 + zero3_init_flag: false + zero_stage: 1 +distributed_type: DEEPSPEED +downcast_bf16: 'no' +machine_rank: 0 +main_training_function: main +mixed_precision: 'bf16' +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/lib/python3.12/site-packages/trl/accelerate_configs/zero2.yaml b/lib/python3.12/site-packages/trl/accelerate_configs/zero2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..239b14ac3a9ae8de73122d1154bf0d71903dc15f --- /dev/null +++ b/lib/python3.12/site-packages/trl/accelerate_configs/zero2.yaml @@ -0,0 +1,21 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + deepspeed_multinode_launcher: standard + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: false + zero_stage: 2 +distributed_type: DEEPSPEED +downcast_bf16: 'no' +machine_rank: 0 +main_training_function: main +mixed_precision: 'bf16' +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/lib/python3.12/site-packages/trl/accelerate_configs/zero3.yaml b/lib/python3.12/site-packages/trl/accelerate_configs/zero3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b5a1201f8a2ee8706b63f0f80c664a1fc61a7d9d --- /dev/null +++ b/lib/python3.12/site-packages/trl/accelerate_configs/zero3.yaml @@ -0,0 +1,22 @@ +compute_environment: LOCAL_MACHINE +debug: false +deepspeed_config: + deepspeed_multinode_launcher: standard + offload_optimizer_device: none + offload_param_device: none + zero3_init_flag: true + zero3_save_16bit_model: true + zero_stage: 3 +distributed_type: DEEPSPEED +downcast_bf16: 'no' +machine_rank: 0 +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +tpu_env: [] +tpu_use_cluster: false +tpu_use_sudo: false +use_cpu: false diff --git a/lib/python3.12/site-packages/trl/cli.py b/lib/python3.12/site-packages/trl/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..9927687ce1ca3d775832d24efb1d2ac6fc22c909 --- /dev/null +++ b/lib/python3.12/site-packages/trl/cli.py @@ -0,0 +1,137 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.resources as resources +import os +import sys +import warnings + +from accelerate.commands.launch import launch_command, launch_command_parser + +from .scripts.dpo import make_parser as make_dpo_parser +from .scripts.env import print_env +from .scripts.grpo import make_parser as make_grpo_parser +from .scripts.kto import make_parser as make_kto_parser +from .scripts.sft import make_parser as make_sft_parser +from .scripts.utils import TrlParser +from .scripts.vllm_serve import main as vllm_serve_main +from .scripts.vllm_serve import make_parser as make_vllm_serve_parser + + +def main(): + parser = TrlParser(prog="TRL CLI", usage="trl", allow_abbrev=False) + + # Add the subparsers + subparsers = parser.add_subparsers(help="available commands", dest="command", parser_class=TrlParser) + + # Add the subparsers for every script + make_dpo_parser(subparsers) + subparsers.add_parser("env", help="Print the environment information") + make_grpo_parser(subparsers) + make_kto_parser(subparsers) + make_sft_parser(subparsers) + make_vllm_serve_parser(subparsers) + + # Parse the arguments; the remaining ones (`launch_args`) are passed to the 'accelerate launch' subparser. + # Duplicates may occur if the same argument is provided in both the config file and CLI. + # For example: launch_args = `["--num_processes", "4", "--num_processes", "8"]`. + # Deduplication and precedence (CLI over config) are handled later by launch_command_parser. + args, launch_args = parser.parse_args_and_config(return_remaining_strings=True) + + # Replace `--accelerate_config foo` with `--config_file trl/accelerate_configs/foo.yaml` if it is present in the + # launch_args. It allows the user to use predefined accelerate configs from the `trl` package. + if "--accelerate_config" in launch_args: + # Get the index of the '--accelerate_config' argument and the corresponding config name + config_index = launch_args.index("--accelerate_config") + config_name = launch_args[config_index + 1] + + # If the config_name correspond to a path in the filesystem, we don't want to override it + if os.path.isfile(config_name): + accelerate_config_path = config_name + elif resources.files("trl.accelerate_configs").joinpath(f"{config_name}.yaml").exists(): + # Get the predefined accelerate config path from the package resources + accelerate_config_path = resources.files("trl.accelerate_configs").joinpath(f"{config_name}.yaml") + else: + raise ValueError( + f"Accelerate config {config_name} is neither a file nor a valid config in the `trl` package. " + "Please provide a valid config name or a path to a config file." + ) + + # Remove '--accelerate_config' and its corresponding config name + launch_args.pop(config_index) + launch_args.pop(config_index) + + # Insert '--config_file' and the absolute path to the front of the list + launch_args = ["--config_file", str(accelerate_config_path)] + launch_args + + if args.command == "dpo": + # Get the default args for the launch command + dpo_training_script = resources.files("trl.scripts").joinpath("dpo.py") + args = launch_command_parser().parse_args([str(dpo_training_script)]) + + # Feed the args to the launch command + args.training_script_args = sys.argv[2:] # remove "trl" and "dpo" + launch_command(args) # launch training + + elif args.command == "env": + print_env() + + elif args.command == "grpo": + # Get the default args for the launch command + grpo_training_script = resources.files("trl.scripts").joinpath("grpo.py") + args = launch_command_parser().parse_args([str(grpo_training_script)]) + + # Feed the args to the launch command + args.training_script_args = sys.argv[2:] # remove "trl" and "grpo" + launch_command(args) # launch training + + elif args.command == "kto": + # Get the default args for the launch command + kto_training_script = resources.files("trl.scripts").joinpath("kto.py") + args = launch_command_parser().parse_args([str(kto_training_script)]) + + # Feed the args to the launch command + args.training_script_args = sys.argv[2:] # remove "trl" and "kto" + launch_command(args) # launch training + + elif args.command == "sft": + # Get the path to the training script + sft_training_script = resources.files("trl.scripts").joinpath("sft.py") + + # This simulates running: `accelerate launch sft.py `. + # Note that the training script args may include launch-related arguments (e.g., `--num_processes`), + # but we rely on the script to ignore any that don't apply to it. + training_script_args = sys.argv[2:] # Remove "trl" and "sft" + args = launch_command_parser().parse_args(launch_args + [str(sft_training_script)] + training_script_args) + launch_command(args) # launch training + + elif args.command == "vllm-serve": + (script_args,) = parser.parse_args_and_config() + + # Known issue: Using DeepSpeed with tensor_parallel_size=1 and data_parallel_size>1 may cause a crash when + # launched via the CLI. Suggest running the module directly. + # More information: https://github.com/vllm-project/vllm/issues/17079 + if script_args.tensor_parallel_size == 1 and script_args.data_parallel_size > 1: + warnings.warn( + "Detected configuration: tensor_parallel_size=1 and data_parallel_size>1. This setup is known to " + "cause a crash when using the `trl vllm-serve` CLI entry point. As a workaround, please run the " + "server using the module path instead: `python -m trl.scripts.vllm_serve`", + RuntimeWarning, + ) + + vllm_serve_main(script_args) + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/trl/core.py b/lib/python3.12/site-packages/trl/core.py new file mode 100644 index 0000000000000000000000000000000000000000..bd8733bedeb313fea48aecd0c67aa7d89d5132d4 --- /dev/null +++ b/lib/python3.12/site-packages/trl/core.py @@ -0,0 +1,159 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import warnings +from collections.abc import Mapping +from contextlib import contextmanager +from typing import Optional, Union + +import numpy as np +import torch +from transformers import is_torch_npu_available, is_torch_xpu_available + + +def flatten_dict(nested: dict, sep: str = "/") -> dict: + """Flatten dictionary and concatenate nested keys with separator.""" + + def recurse(nest: dict, prefix: str, into: dict) -> None: + for k, v in nest.items(): + if sep in k: + raise ValueError(f"separator '{sep}' not allowed to be in key '{k}'") + if isinstance(v, Mapping): + recurse(v, prefix + k + sep, into) + else: + into[prefix + k] = v + + flat = {} + recurse(nested, "", flat) + return flat + + +def masked_mean(values: torch.Tensor, mask: torch.Tensor, axis: Optional[bool] = None) -> torch.Tensor: + """Compute mean of tensor with a masked values.""" + if axis is not None: + return (values * mask).sum(axis=axis) / mask.sum(axis=axis) + else: + return (values * mask).sum() / mask.sum() + + +def masked_var(values: torch.Tensor, mask: torch.Tensor, unbiased: bool = True) -> torch.Tensor: + """Compute variance of tensor with masked values.""" + mean = masked_mean(values, mask) + centered_values = values - mean + variance = masked_mean(centered_values**2, mask) + if unbiased: + mask_sum = mask.sum() + if mask_sum == 0: + raise ValueError( + "The sum of the mask is zero, which can happen when `mini_batch_size=1`;" + "try increase the `mini_batch_size` or `gradient_accumulation_steps`" + ) + # note that if mask_sum == 1, then there is a division by zero issue + # to avoid it you just need to use a larger minibatch_size + bessel_correction = mask_sum / (mask_sum - 1) + variance = variance * bessel_correction + return variance + + +def masked_whiten(values: torch.Tensor, mask: torch.Tensor, shift_mean: bool = True) -> torch.Tensor: + """Whiten values with masked values.""" + mean, var = masked_mean(values, mask), masked_var(values, mask) + whitened = (values - mean) * torch.rsqrt(var + 1e-8) + if not shift_mean: + whitened += mean + return whitened + + +class LengthSampler: + """ + Samples a length + """ + + def __init__(self, min_value: int, max_value: int): + self.values = list(range(min_value, max_value)) + + def __call__(self) -> int: + return np.random.choice(self.values) + + +class PPODecorators: + optimize_device_cache = False + + @classmethod + @contextmanager + def empty_device_cache(cls): + yield + if cls.optimize_device_cache: + if is_torch_xpu_available(): + gc.collect() + torch.xpu.empty_cache() + gc.collect() + elif is_torch_npu_available(): + gc.collect() + torch.npu.empty_cache() + gc.collect() + elif torch.cuda.is_available(): + gc.collect() + torch.cuda.empty_cache() + gc.collect() + + +def randn_tensor( + shape: Union[tuple, list], + generator: Optional[Union[list[torch.Generator], torch.Generator]] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + layout: Optional[torch.layout] = None, +) -> torch.Tensor: + """A helper function to create random tensors on the desired `device` with the desired `dtype`. When + passing a list of generators, you can seed each batch size individually. If CPU generators are passed, the tensor + is always created on the CPU. + """ + # device on which tensor is created defaults to device + rand_device = device + batch_size = shape[0] + + layout = layout or torch.strided + device = device or torch.device("cpu") + + if generator is not None: + gen_device_type = generator.device.type if not isinstance(generator, list) else generator[0].device.type + if gen_device_type != device.type and gen_device_type == "cpu": + rand_device = "cpu" + if device != "mps": + warnings.warn( + f"The passed generator was created on 'cpu' even though a tensor on {device} was expected." + f" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably" + f" slighly speed up this function by passing a generator that was created on the {device} device.", + UserWarning, + ) + elif gen_device_type != device.type and gen_device_type == "cuda": + raise ValueError(f"Cannot generate a {device} tensor from a generator of type {gen_device_type}.") + + # make sure generator list of length 1 is treated like a non-list + if isinstance(generator, list) and len(generator) == 1: + generator = generator[0] + + if isinstance(generator, list): + shape = (1,) + shape[1:] + latents = [ + torch.randn(shape, generator=generator[i], device=rand_device, dtype=dtype, layout=layout) + for i in range(batch_size) + ] + latents = torch.cat(latents, dim=0).to(device) + else: + latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype, layout=layout).to(device) + + return latents diff --git a/lib/python3.12/site-packages/trl/data_utils.py b/lib/python3.12/site-packages/trl/data_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fbf9129ea5bccd5e6426ed6d0490ae5463116ed4 --- /dev/null +++ b/lib/python3.12/site-packages/trl/data_utils.py @@ -0,0 +1,814 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings +from collections import defaultdict, deque +from collections.abc import Sequence +from itertools import takewhile +from typing import Any, Callable, Optional, TypeVar, Union + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.types +from datasets import Dataset, DatasetDict +from transformers import PreTrainedTokenizerBase + + +DatasetType = TypeVar("DatasetType", Dataset, DatasetDict) + + +def is_conversational(example: dict[str, Any]) -> bool: + r""" + Check if the example is in a conversational format. + + Args: + example (`dict[str, Any]`): + A single data entry of a dataset. The example can have different keys depending on the dataset type. + + Returns: + `bool`: + `True` if the data is in a conversational format, `False` otherwise. + + Examples: + + ```python + >>> example = {"prompt": [{"role": "user", "content": "What color is the sky?"}]} + >>> is_conversational(example) + True + + >>> example = {"prompt": "The sky is"} + >>> is_conversational(example) + False + ``` + """ + supported_keys = ["prompt", "chosen", "rejected", "completion", "messages"] + example_keys = {key for key in example.keys() if key in supported_keys} + + # It must have one of the supported keys + if example_keys: + key = example_keys.pop() # take the first supported key + maybe_messages = example[key] + # It must be a list of messages + if isinstance(maybe_messages, list): + maybe_message = maybe_messages[0] + # Each message must a list of dictionaries with keys "role" and "content" + if isinstance(maybe_message, dict) and "role" in maybe_message and "content" in maybe_message: + return True + + return False + + +def apply_chat_template( + example: dict[str, list[dict[str, str]]], + tokenizer: PreTrainedTokenizerBase, + tools: Optional[list[Union[dict, Callable]]] = None, +) -> dict[str, str]: + r""" + Apply a chat template to a conversational example along with the schema for a list of functions in `tools`. + + For more details, see [`maybe_apply_chat_template`]. + """ + # Check that the example has the correct keys + supported_keys = ["prompt", "chosen", "rejected", "completion", "messages", "label"] + example_keys = {key for key in example.keys() if key in supported_keys} + if example_keys not in [ + {"messages"}, # language modeling + {"prompt"}, # prompt-only + {"prompt", "completion"}, # prompt-completion + {"prompt", "chosen", "rejected"}, # preference + {"chosen", "rejected"}, # preference with implicit prompt + {"prompt", "completion", "label"}, # unpaired preference + ]: + raise KeyError(f"Invalid keys in the example: {example_keys}") + + # Apply the chat template to the whole conversation + if "messages" in example: + messages = tokenizer.apply_chat_template(example["messages"], tools=tools, tokenize=False) + + # Apply the chat template to the prompt, adding the generation prompt + if "prompt" in example: + last_role = example["prompt"][-1]["role"] + if last_role == "user": + add_generation_prompt = True + continue_final_message = False + elif last_role == "assistant": + add_generation_prompt = False + continue_final_message = True + else: + raise ValueError(f"Invalid role in the last message: {last_role}") + prompt = tokenizer.apply_chat_template( + example["prompt"], + tools=tools, + continue_final_message=continue_final_message, + tokenize=False, + add_generation_prompt=add_generation_prompt, + ) + + # Apply the chat template to the entire prompt + completion + if "prompt" in example: # explicit prompt and prompt-completion case + if "chosen" in example: + prompt_chosen = tokenizer.apply_chat_template( + example["prompt"] + example["chosen"], tools=tools, tokenize=False + ) + # DeepSeek-R1 inserts a token when using `add_generation_prompt`, which can cause discrepancies + # between the prompt alone and the combined prompt+completion. To ensure consistency, we extract the + # common prefix between the two. In most cases, this is a no-op. + prompt = "".join(x for x, _ in takewhile(lambda x: x[0] == x[1], zip(prompt, prompt_chosen))) + + chosen = prompt_chosen[len(prompt) :] + if "rejected" in example and "prompt" in example: # explicit prompt + prompt_rejected = tokenizer.apply_chat_template( + example["prompt"] + example["rejected"], tools=tools, tokenize=False + ) + # Handle DeepSeek-R1 token, see the above comment for details + prompt = "".join(x for x, _ in takewhile(lambda x: x[0] == x[1], zip(prompt, prompt_rejected))) + rejected = prompt_rejected[len(prompt) :] + if "completion" in example: + prompt_completion = tokenizer.apply_chat_template( + example["prompt"] + example["completion"], tools=tools, tokenize=False + ) + # Handle DeepSeek-R1 token, see the above comment for details + prompt = "".join(x for x, _ in takewhile(lambda x: x[0] == x[1], zip(prompt, prompt_completion))) + completion = prompt_completion[len(prompt) :] + else: # implicit prompt case + if "chosen" in example: + chosen = tokenizer.apply_chat_template(example["chosen"], tools=tools, tokenize=False) + if "rejected" in example: + rejected = tokenizer.apply_chat_template(example["rejected"], tools=tools, tokenize=False) + + # Extract the completion by removing the prompt part from the prompt-completion string + output = {} + if "messages" in example: + output["text"] = messages + if "prompt" in example: + output["prompt"] = prompt + if "chosen" in example: + output["chosen"] = chosen + if "rejected" in example: + output["rejected"] = rejected + if "completion" in example: + output["completion"] = completion + if "label" in example: + output["label"] = example["label"] + + return output + + +def maybe_apply_chat_template( + example: dict[str, list[dict[str, str]]], + tokenizer: PreTrainedTokenizerBase, + tools: Optional[list[Union[dict, Callable]]] = None, +) -> dict[str, str]: + r""" + If the example is in a conversational format, apply a chat template to it. + + Args: + example (`dict[str, list[dict[str, str]]`): + Dictionary representing a single data entry of a conversational dataset. Each data entry can have different + keys depending on the dataset type. The supported dataset types are: + + - Language modeling dataset: `"messages"`. + - Prompt-only dataset: `"prompt"`. + - Prompt-completion dataset: `"prompt"` and `"completion"`. + - Preference dataset: `"prompt"`, `"chosen"`, and `"rejected"`. + - Preference dataset with implicit prompt: `"chosen"` and `"rejected"`. + - Unpaired preference dataset: `"prompt"`, `"completion"`, and `"label"`. + + For keys `"messages"`, `"prompt"`, `"chosen"`, `"rejected"`, and `"completion"`, the values are lists of + messages, where each message is a dictionary with keys `"role"` and `"content"`. + tokenizer (`PreTrainedTokenizerBase`): + Tokenizer to apply the chat template with. + tools (`list[Union[dict, Callable]]` or `None`, *optional*, defaults to `None`): + A list of tools (callable functions) that will be accessible to the model. If the template does not support + function calling, this argument will have no effect + + Returns: + `dict[str, str]`: + Formatted example with the chat template applied. + + Notes: + - This function does not alter the keys, except for Language modeling dataset, where `"messages"` is replaced + by `"text"`. + + - In case of prompt-only data, if the last role is `"user"`, the generation prompt is added to the prompt. + Else, if the last role is `"assistant"`, the final message is continued. + + Example: + + ```python + >>> from transformers import AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct") + >>> example = { + ... "prompt": [{"role": "user", "content": "What color is the sky?"}], + ... "completion": [{"role": "assistant", "content": "It is blue."}], + ... } + >>> apply_chat_template(example, tokenizer) + {'prompt': '<|user|>\nWhat color is the sky?<|end|>\n<|assistant|>\n', 'completion': 'It is blue.<|end|>\n<|endoftext|>'} + ``` + """ + if is_conversational(example): + return apply_chat_template(example, tokenizer, tools) + else: + return example + + +def _unpair_row(examples: list[dict[str, list[dict[str, str]]]]) -> list[dict[str, list[dict[str, str]]]]: + batch_size = len(examples["chosen"]) + new_rows = { + "completion": examples["chosen"] + examples["rejected"], + "label": [True] * batch_size + [False] * batch_size, + } + if "prompt" in examples: + new_rows["prompt"] = examples["prompt"] + examples["prompt"] + return new_rows + + +def unpair_preference_dataset( + dataset: DatasetType, num_proc: Optional[int] = None, desc: Optional[str] = None +) -> DatasetType: + r""" + Unpair a preference dataset. + + Args: + dataset (`Dataset` or `DatasetDict`): + Preference dataset to unpair. The dataset must have columns `"chosen"`, `"rejected"` and optionally + `"prompt"`. + num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + desc (`str` or `None`, *optional*, defaults to `None`): + Meaningful description to be displayed alongside with the progress bar while mapping examples. + + Returns: + `Dataset`: The unpaired preference dataset. + + Example: + + ```python + >>> from datasets import Dataset + + >>> dataset_dict = { + ... "prompt": ["The sky is", "The sun is"], + ... "chosen": [" blue.", "in the sky."], + ... "rejected": [" green.", " in the sea."], + ... } + >>> dataset = Dataset.from_dict(dataset_dict) + >>> dataset = unpair_preference_dataset(dataset) + >>> dataset + Dataset({ + features: ['prompt', 'completion', 'label'], + num_rows: 4 + }) + + >>> dataset[0] + {'prompt': 'The sky is', 'completion': ' blue.', 'label': True} + ``` + """ + return dataset.map(_unpair_row, batched=True, remove_columns=["chosen", "rejected"], num_proc=num_proc, desc=desc) + + +def maybe_unpair_preference_dataset( + dataset: DatasetType, num_proc: Optional[int] = None, desc: Optional[str] = None +) -> DatasetType: + r""" + Unpair a preference dataset if it is paired. + + Args: + dataset (`Dataset` or `DatasetDict`): + Preference dataset to unpair. The dataset must have columns `"chosen"`, `"rejected"` and optionally + `"prompt"`. + num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + desc (`str` or `None`, *optional*, defaults to `None`): + Meaningful description to be displayed alongside with the progress bar while mapping examples. + + Returns: + `Dataset` or `DatasetDict`: The unpaired preference dataset if it was paired, otherwise the original dataset. + + Example: + + ```python + >>> from datasets import Dataset + + >>> dataset_dict = { + ... "prompt": ["The sky is", "The sun is"], + ... "chosen": [" blue.", "in the sky."], + ... "rejected": [" green.", " in the sea."], + ... } + >>> dataset = Dataset.from_dict(dataset_dict) + >>> dataset = unpair_preference_dataset(dataset) + >>> dataset + Dataset({ + features: ['prompt', 'completion', 'label'], + num_rows: 4 + }) + + >>> dataset[0] + {'prompt': 'The sky is', 'completion': ' blue.', 'label': True} + ``` + """ + if isinstance(dataset, DatasetDict): + column_names = dataset[list(dataset.keys())[0]].column_names + else: + column_names = dataset.column_names + if "chosen" in column_names and "rejected" in column_names: + return unpair_preference_dataset(dataset, num_proc=num_proc, desc=desc) + else: + return dataset + + +def extract_prompt(example: dict[str, Sequence]) -> dict[str, Sequence]: + r""" + Extracts the shared prompt from a preference data example, where the prompt is implicit within both the chosen and + rejected completions. + + For more details, see [`maybe_extract_prompt`]. + """ + for idx in range(min(len(example["chosen"]), len(example["rejected"]))): + if example["chosen"][idx] != example["rejected"][idx]: + if example["chosen"][idx - 1] == " ": # remove space before the prompt + idx -= 1 + break + return { + "prompt": example["chosen"][:idx], + "chosen": example["chosen"][idx:], + "rejected": example["rejected"][idx:], + } + + +def maybe_extract_prompt(example: dict[str, list]) -> dict[str, list]: + r""" + Extracts the shared prompt from a preference data example, where the prompt is implicit within both the chosen and + rejected completions. + + If the example already contains a `"prompt"` key, the function returns the example as is. Else, the function + identifies the longest common sequence (prefix) of conversation turns between the "chosen" and "rejected" + completions and extracts this as the prompt. It then removes this prompt from the respective "chosen" and + "rejected" completions. + + Args: + example (`dict[str, list]`): + A dictionary representing a single data entry in the preference dataset. It must contain the keys + `"chosen"` and `"rejected"`, where each value is either conversational or standard (`str`). + + Returns: + `dict[str, list]`: A dictionary containing: + - `"prompt"`: The longest common prefix between the "chosen" and "rejected" completions. + - `"chosen"`: The remainder of the "chosen" completion, with the prompt removed. + - `"rejected"`: The remainder of the "rejected" completion, with the prompt removed. + + Examples: + + ```python + >>> example = { + ... "chosen": [ + ... {"role": "user", "content": "What color is the sky?"}, + ... {"role": "assistant", "content": "It is blue."}, + ... ], + ... "rejected": [ + ... {"role": "user", "content": "What color is the sky?"}, + ... {"role": "assistant", "content": "It is green."}, + ... ], + ... } + >>> extract_prompt(example) + {'prompt': [{'role': 'user', 'content': 'What color is the sky?'}], + 'chosen': [{'role': 'assistant', 'content': 'It is blue.'}], + 'rejected': [{'role': 'assistant', 'content': 'It is green.'}]} + ``` + + Or, with the `map` method of `datasets.Dataset`: + + ```python + >>> from trl import extract_prompt + >>> from datasets import Dataset + + >>> dataset_dict = { + ... "chosen": [ + ... [ + ... {"role": "user", "content": "What color is the sky?"}, + ... {"role": "assistant", "content": "It is blue."}, + ... ], + ... [ + ... {"role": "user", "content": "Where is the sun?"}, + ... {"role": "assistant", "content": "In the sky."}, + ... ], + ... ], + ... "rejected": [ + ... [ + ... {"role": "user", "content": "What color is the sky?"}, + ... {"role": "assistant", "content": "It is green."}, + ... ], + ... [ + ... {"role": "user", "content": "Where is the sun?"}, + ... {"role": "assistant", "content": "In the sea."}, + ... ], + ... ], + ... } + >>> dataset = Dataset.from_dict(dataset_dict) + >>> dataset = dataset.map(extract_prompt) + >>> dataset[0] + {'prompt': [{'role': 'user', 'content': 'What color is the sky?'}], + 'chosen': [{'role': 'assistant', 'content': 'It is blue.'}], + 'rejected': [{'role': 'assistant', 'content': 'It is green.'}]} + ``` + """ + # Some dataset add a `"prompt"` column, even though the prompt is implicit and included in the "chosen" and + # "rejected" completions. E.g.: + # {"prompt": "What color is the sky?", + # "chosen": [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is blue."}], + # "rejected": [{"role": "user", "content": "What color is the sky?"}, {"role": "assistant", "content": "It is green."}]} + # That's why we check if the prompt is also conversational before deciding not to extract it. + if "chosen" not in example or "rejected" not in example: # not a preference example + return example + if "prompt" in example: + # Both conversational or both non-conversational + chosen_conv = is_conversational({"chosen": example["chosen"]}) + prompt_conv = is_conversational({"prompt": example["prompt"]}) + if (chosen_conv and prompt_conv) or (not chosen_conv and not prompt_conv): + return example + return extract_prompt({"chosen": example["chosen"], "rejected": example["rejected"]}) + + +def pack_examples(examples: dict[str, list[list]], seq_length: int) -> dict[str, list[list]]: + """ + Pack examples into chunks of size `seq_length`. + + Args: + examples (`dict[str, list[list]]`): + Dictionary of examples with keys as strings and values as lists of lists. + seq_length (`int`): + Maximum sequence length. + + Returns: + `dict[str, list[list]]`: Dictionary of examples with keys as strings and values as lists of lists. + + Example: + + ```python + >>> from trl import pack_examples + + >>> examples = { + ... "input_ids": [[1, 2, 3], [4, 5, 6, 7], [8]], + ... "attention_mask": [[0, 1, 1], [0, 0, 1, 1], [1]], + ... } + >>> pack_examples(examples, seq_length=5) + {'input_ids': [[1, 2, 3, 4, 5], [6, 7, 8]], 'attention_mask': [[0, 1, 1, 0, 0], [1, 1, 1]]} + + >>> pack_examples(examples, seq_length=2) + {'input_ids': [[1, 2], [3, 4], [5, 6], [7, 8]], 'attention_mask': [[0, 1], [1, 0], [0, 1], [1, 1]]} + ``` + """ + warnings.warn( + "`pack_examples` is deprecated and will be removed in version 0.20.0. Use `pack_dataset` with a dataset " + "instead.", + DeprecationWarning, + ) + # Join all the values into a single list + examples = {k: sum(v, []) for k, v in examples.items()} + # Split the values into chunks of size seq_length + examples = {k: [v[i : i + seq_length] for i in range(0, len(v), seq_length)] for k, v in examples.items()} + return examples + + +class _SegmentTree: + """ + A segment tree data structure that, when initialized as `_SegmentTree(maxval)`, efficiently finds the next larger + value for a given input within the range [1, maxval]. + + See [Fewer Truncations Improve Language Modeling](https://arxiv.org/abs/2404.10830) for more details. + """ + + def __init__(self, maxval: int): + self.maxval = maxval + self.tree = [0] * (2 * maxval) + + def add(self, val): + assert 0 < val <= self.maxval + i = self.maxval + val - 1 + self.tree[i] = val + while i > 1: + i >>= 1 + left, right = self.tree[i << 1], self.tree[(i << 1) + 1] + # Compare the values using if-else otherwise repeated calls to `builtins.max` become the bottleneck + self.tree[i] = left if left >= right else right + + def remove(self, val): + assert 0 < val <= self.maxval + i = self.maxval + val - 1 + self.tree[i] = 0 + while i > 1: + i >>= 1 + left, right = self.tree[i << 1], self.tree[(i << 1) + 1] + # Compare the values using if-else otherwise repeated calls to `builtins.max` become the bottleneck + self.tree[i] = left if left >= right else right + + def search(self, val): + assert 0 < val <= self.maxval + i = 1 + while i < self.maxval: + if self.tree[i << 1] >= val: + i = i << 1 + else: + i = (i << 1) + 1 + return self.tree[i] + + +def _pack_ffd(examples: pa.Table, seq_length: int) -> pa.Table: + """Pack sequences in a pyarrow Table using First Fit Decreasing strategy.""" + # Add position_ids to the examples + input_ids = examples["input_ids"] + position_ids_python = [list(range(len(sequence))) for sequence in input_ids.to_pylist()] + position_ids_array = pa.array(position_ids_python, type=examples["input_ids"].type) + examples = examples.append_column("position_ids", position_ids_array) + + columns = [] + list_column_idx = None + for idx, column in enumerate(examples.columns): + if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list(column.type): + column = pc.list_slice(column, 0, seq_length) + if list_column_idx is None: + list_column_idx = idx + columns.append(column) + examples = pa.Table.from_arrays(columns, names=examples.column_names) + + ids = np.arange(len(examples)) + assert list_column_idx is not None + lengths = pc.make_struct(pc.list_value_length(examples[list_column_idx]).combine_chunks(), ids) + lengths = lengths.sort("descending", by=0) + + segment_tree = _SegmentTree(seq_length) + segment_tree.add(seq_length) # the max, `seq_length` bin is always available + space_to_bin = defaultdict(deque) + + # Bin is represented as a dict (of example ids and sum of their lengths) to allow in-place updates + bins: list[dict] = [] + for length, idx in zip(lengths.field(0).to_numpy(), lengths.field(1).to_numpy()): + space = segment_tree.search(length) + + if space < seq_length: + bin = space_to_bin[space].popleft() + else: + bin = {"ids": [], "length": 0} + bins.append(bin) + + bin["ids"].append(idx) + bin["length"] += length + if space < seq_length and not space_to_bin[space]: + segment_tree.remove(space) + + space = space - length + space_to_bin[space].append(bin) + if space > 0: + segment_tree.add(space) + + examples = pc.take(examples, [id_ for bin in bins for id_ in bin["ids"]]) + offsets = np.array([0] + [bin["length"] for bin in bins]) + offsets = np.cumsum(offsets) + + columns = [] + for column in examples.columns: + assert len(column.chunks) == 1 # `pc.take` returns a ChunkedArray with a single chunk + column = column.chunks[0] + if pa.types.is_list(column.type) or pa.types.is_large_list(column.type): + dtype = column.offsets.type.to_pandas_dtype() + column = type(column).from_arrays(offsets.astype(dtype), column.values) + columns.append(column) + return pa.Table.from_arrays(columns, names=examples.column_names) + + +def _pack_wrapped(examples: pa.Table, seq_length: int) -> pa.Table: + """Pack sequences in a pyarrow Table using a wrapped strategy.""" + columns = [] + for column in examples.columns: + if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list(column.type): + if isinstance(column, pa.ChunkedArray): + column = column.combine_chunks() + offsets, values = column.offsets, column.values + values = values[offsets[0].as_py() : offsets[-1].as_py()] + num_elements = len(values) + dtype = offsets.type.to_pandas_dtype() # np.int32 or np.int64 + offsets = np.arange(0, num_elements, seq_length, dtype=dtype) + offsets = np.concatenate((offsets, [num_elements])) + column = type(column).from_arrays(offsets, values) + columns.append(column) + return pa.Table.from_arrays(columns, names=examples.column_names) + + +def pack_dataset( + dataset: DatasetType, seq_length: int, strategy: str = "ffd", map_kwargs: Optional[dict[str, Any]] = None +) -> DatasetType: + r""" + Pack sequences in a dataset into chunks of size `seq_length`. + + Args: + dataset (`Dataset` or `DatasetDict`): + Dataset to pack + seq_length (`int`): + Target sequence length to pack to. + strategy (`str`, *optional*, defaults to `"ffd"`): + Packing strategy to use. Can be either: + + - `"ffd"` (First Fit Decreasing): Slower but preserves sequence boundaries. Sequences are never cut in the + middle. + - `"wrapped"`: Faster but more aggressive. Ignores sequence boundaries and will cut sequences in the middle + to completely fill each packed sequence with data. + map_kwargs (`dict` or `None`, *optional*, defaults to `None`): + Additional keyword arguments to pass to the dataset's map method when packing examples. + + Returns: + `Dataset` or `DatasetDict`: The dataset with packed sequences. The number of examples may decrease as sequences + are combined. + + Example: + ```python + >>> from datasets import Dataset + >>> from trl import pack_dataset + + >>> examples = { + ... "input_ids": [[1, 2, 3], [4, 5], [6, 7, 8], [9]], + ... "attention_mask": [[1, 1, 0], [1, 0], [1, 0, 0], [1]], + ... } + >>> dataset = Dataset.from_dict(examples) + >>> packed_dataset = pack_dataset(dataset, seq_length=4, strategy="ffd") + >>> packed_dataset[:] + {'input_ids': [[1, 2, 3, 9], [6, 7, 8, 4, 5]], + 'attention_mask': [[1, 1, 0, 1], [1, 0, 0, 1, 0]]} + ``` + """ + if map_kwargs is None: + map_kwargs = {} + # Fast packing with pyarrow + dataset = dataset.with_format("arrow") + if strategy == "ffd": + dataset = dataset.map(_pack_ffd, batched=True, fn_kwargs={"seq_length": seq_length}, **map_kwargs) + elif strategy == "wrapped": + dataset = dataset.map(_pack_wrapped, batched=True, fn_kwargs={"seq_length": seq_length}, **map_kwargs) + else: + raise ValueError(f"Invalid packing strategy: {strategy}. Use 'ffd' or 'wrapped'.") + dataset = dataset.with_format(None) + return dataset + + +def truncate_dataset( + dataset: DatasetType, max_length: int, map_kwargs: Optional[dict[str, Any]] = None +) -> DatasetType: + r""" + Truncate sequences in a dataset to a specifed `max_length`. + + Args: + dataset (`Dataset` or `DatasetDict`): + Dataset to truncate. + seq_length (`int`): + Maximum sequence length to truncate to. + map_kwargs (`dict` or `None`, *optional*, defaults to `None`): + Additional keyword arguments to pass to the dataset's map method when truncating examples. + + Returns: + `Dataset` or `DatasetDict`: The dataset with truncated sequences. + + Example: + ```python + >>> from datasets import Dataset + + >>> examples = { + ... "input_ids": [[1, 2, 3], [4, 5, 6, 7], [8]], + ... "attention_mask": [[0, 1, 1], [0, 0, 1, 1], [1]], + ... } + >>> dataset = Dataset.from_dict(examples) + >>> truncated_dataset = truncate_dataset(dataset, max_length=2) + >>> truncated_dataset[:] + {'input_ids': [[1, 2], [4, 5], [8]], + 'attention_mask': [[0, 1], [0, 0], [1]]} + ``` + """ + if map_kwargs is None: + map_kwargs = {} + if isinstance(dataset, Dataset): + # Fast truncation with pyarrow + def truncate(examples): + truncated_columns = [] + for column in examples.columns: + if pyarrow.types.is_list(column.type) or pyarrow.types.is_large_list(column.type): + column = pc.list_slice(column, 0, max_length) + truncated_columns.append(column) + return pa.Table.from_arrays(truncated_columns, names=examples.column_names) + + dataset = dataset.with_format("arrow") + dataset = dataset.map(truncate, batched=True, **map_kwargs) + dataset = dataset.with_format(None) + else: + + def truncate(examples): + truncated_examples = {} + for key, column in examples.items(): + if column and isinstance(column[0], list): + column = [val[:max_length] for val in column] + truncated_examples[key] = column + return truncated_examples + + dataset = dataset.map( + truncate, + batched=True, + **map_kwargs, + ) + return dataset + + +def is_conversational_from_value(example: dict[str, Any]) -> bool: + r""" + Check if the example is in a conversational format (from/value). Note that this format isn't recommended. Prefer + the ChatML format (role/content) + + Args: + example (`dict[str, Any]`): + A single data entry of a dataset. The example can have different keys depending on the dataset type. + + Returns: + `bool`: + `True` if the data is in a conversational Chatformat, `False` otherwise. + + Examples: + + ```python + >>> example = {"conversations": [{"from": "user", "value": "What color is the sky?"}]} + >>> is_conversational_from_value(example) + True + >>> example = {"conversations": [{"role": "user", "content": "What color is the sky?"}]} + >>> is_conversational_from_value(example) + False + >>> example = {"conversations": "The sky is"}) + >>> is_conversational_from_value(example) + False + ``` + """ + maybe_messages = example.get("conversations") + # It must be a list of messages + if isinstance(maybe_messages, list): + maybe_message = maybe_messages[0] + # Each message must a list of dictionaries with keys "from" and "value" + if isinstance(maybe_message, dict) and "from" in maybe_message and "value" in maybe_message: + return True + + return False + + +def maybe_convert_to_chatml(example: dict[str, list]) -> dict[str, list]: + """ + Convert a conversational dataset with fields `from` and `value` to ChatML format. + + This function modifies conversational data to align with OpenAI's ChatML format: + - Replaces the key `"from"` with `"role"` in message dictionaries. + - Replaces the key `"value"` with `"content"` in message dictionaries. + - Renames `"conversations"` to `"messages"` for consistency with ChatML. + + Args: + example (`dict[str, list]`): + A single data entry containing a list of messages. + + Returns: + `dict[str, list]`: + Example reformatted to ChatML style. + + Example: + ```python + >>> from trl import maybe_convert_to_chatml + + >>> example = { + ... "conversations": [ + ... {"from": "user", "value": "What color is the sky?"}, + ... {"from": "assistant", "value": "It is blue."}, + ... ] + ... } + >>> maybe_convert_to_chatml(example) + {'messages': [{'role': 'user', 'content': 'What color is the sky?'}, + {'role': 'assistant', 'content': 'It is blue.'}]} + ``` + """ + # List of possible keys containing message lists + for key in ["prompt", "completion", "chosen", "rejected", "messages", "conversations"]: + if key in example and isinstance(example[key], list): + messages = example[key] + for message in messages: + if isinstance(message, dict): + if "from" in message: + message["role"] = message.pop("from") + if "value" in message: + message["content"] = message.pop("value") + + # Rename "conversations" to "messages" + if "conversations" in example: + example["messages"] = example.pop("conversations") + + return example diff --git a/lib/python3.12/site-packages/trl/environment/__init__.py b/lib/python3.12/site-packages/trl/environment/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..70ab98f67c6db7e8b9cff63d60ad0ca1e84587c0 --- /dev/null +++ b/lib/python3.12/site-packages/trl/environment/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ..import_utils import _LazyModule + + +_import_structure = { + "base_environment": ["TextEnvironment", "TextHistory"], +} + +if TYPE_CHECKING: + from .base_environment import TextEnvironment, TextHistory +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/trl/environment/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/trl/environment/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e90fed8b70510f76208f70685712f57252255e5 Binary files /dev/null and b/lib/python3.12/site-packages/trl/environment/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/environment/__pycache__/base_environment.cpython-312.pyc b/lib/python3.12/site-packages/trl/environment/__pycache__/base_environment.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0fc5157a7415896f6f9d5b3bcc05212b575685c1 Binary files /dev/null and b/lib/python3.12/site-packages/trl/environment/__pycache__/base_environment.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/environment/base_environment.py b/lib/python3.12/site-packages/trl/environment/base_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..e6ea3c7623b4ce07a2316bc82fd46727cb8c9cd4 --- /dev/null +++ b/lib/python3.12/site-packages/trl/environment/base_environment.py @@ -0,0 +1,491 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +import warnings +from typing import Optional + +import torch +from accelerate.utils import extract_model_from_parallel +from transformers import StoppingCriteria, StoppingCriteriaList +from transformers.utils import is_rich_available + + +if is_rich_available(): + from rich import print + from rich.text import Text + + +class StringStoppingCriteria(StoppingCriteria): + """Custom `StoppingCriteria` which checks if all generations in the batch are completed.""" + + def __init__(self, stop_strings, tokenizer): + self.stop_strings = stop_strings + self.tokenizer = tokenizer + self.first_call = True + + def __call__(self, input_ids, scores, **kwargs): + """Returns true if all generated sequences contain any of the stop strings.""" + if self.first_call: + self.generated_tokens = [1 for _ in range(input_ids.shape[0])] + self.start_length = input_ids.shape[-1] - 1 + self.first_call = False + decoded_generations = self.tokenizer.batch_decode(input_ids[:, self.start_length :]) + done = [] + + for i, decoded_generation in enumerate(decoded_generations): + sequence_complete = any(stop_string in decoded_generation for stop_string in self.stop_strings) + done.append(sequence_complete) + if not sequence_complete: + self.generated_tokens[i] += 1 + + if all(done): + self.first_call = True + + return all(done) + + +class TextHistory: + """The TextHistory class keeps track of the history of an interaction between the language model and the environment.""" + + def __init__(self, text, tokens, system=True): + """ + Initialize TextHistory. + + Args: + text (`str`): The text of the first segment. + tokens (`torch.LongTensor`): The tokens of the first segment. + system (`bool`, *optional*): Whether the first segment is a system or user segment. + """ + self.system_spans = [] + self.text_spans = [] + self.token_spans = [] + self.token_masks = torch.tensor([], dtype=torch.long).to(tokens.device) + self.text = "" + self.tokens = torch.tensor([], dtype=torch.long).to(tokens.device) + self.completed = False + self.truncated = False + self.reward = 0.0 + + self.prompt_color = "black on grey85" + self.system_color = "black on cyan3" + self.model_color = "black on deep_sky_blue1" + self.reward_color = "black on plum1" + + self.append_segment(text, tokens, system=system) + + def append_segment(self, text, tokens, system=True): + """ + Append a new segment to the history. + + Args: + text (`str`): The text of the new segment. + tokens (`torch.LongTensor`): The tokens of the new segment. + system (`bool`, *optional*): Whether the new segment is a system or user segment. + """ + + if len(text) == 0 or len(tokens) == 0: + raise ValueError("Can't append empty text or token list to history.") + + original_text_length = len(self.text) + + self.text += text + self.text_spans.append((original_text_length, len(self.text))) + self.system_spans.append(system) + + original_token_length = len(self.tokens) + + self.tokens = torch.cat((self.tokens, tokens)) + if system: + self.token_masks = torch.cat((self.token_masks, torch.zeros_like(tokens))) + else: + self.token_masks = torch.cat((self.token_masks, torch.ones_like(tokens))) + self.token_spans.append((original_token_length, len(self.tokens))) + + def complete(self, truncated=False): + """ + Mark the history as completed. + """ + self.completed = True + self.truncated = truncated + + @property + def last_text_segment(self): + """ + Get the last text segment. + """ + start, end = self.text_spans[-1] + return self.text[start:end] + + def split_query_response_tokens(self): + """ + Split the tokens into query and response tokens. + """ + split_index = self.token_spans[0][1] + query = self.tokens[:split_index] + response = self.tokens[split_index:] + mask = self.token_masks[split_index:] + + return query, response, mask + + def show_text(self, show_legend=False): + """ + Print the text history. + """ + if not is_rich_available(): + raise ImportError( + "The `rich` library is required to display text with formatting. Install it using `pip install rich`." + ) + + text = Text(self.text) + text.stylize(self.prompt_color, self.text_spans[0][0], self.text_spans[1][0]) + for i, (start, end) in enumerate(self.text_spans[1:]): + if self.system_spans[i + 1]: + text.stylize(self.system_color, start, end) + else: + text.stylize(self.model_color, start, end) + + text.append(f"\n\nReward: {self.reward}", style=self.reward_color) + print(text) + + if show_legend: + self.show_colour_legend() + + def show_tokens(self, tokenizer, show_legend=False): + """ + Print the history tokens. + """ + if not is_rich_available(): + raise ImportError( + "The `rich` library is required to display tokens with formatting. " + "Install it using `pip install rich`." + ) + + text = Text() + prompt_end = self.token_spans[0][1] + for i, (token, mask) in enumerate(zip(self.tokens, self.token_masks)): + if i < prompt_end: + text.append(tokenizer.convert_ids_to_tokens(token.item()), style=self.prompt_color) + text.append(" ") + elif mask == 0: + text.append(tokenizer.convert_ids_to_tokens(token.item()), style=self.system_color) + text.append(" ") + else: + text.append(tokenizer.convert_ids_to_tokens(token.item()), style=self.model_color) + text.append(" ") + text.append(f"\n\nReward: {self.reward}", style=self.reward_color) + print(text) + if show_legend: + self.show_colour_legend() + + def show_colour_legend(self): + """ + Print the colour legend. + """ + if not is_rich_available(): + raise ImportError( + "The `rich` library is required to display colour legends with formatting. " + "Install it using `pip install rich`." + ) + text = Text("\n\n(Colour Legend: ") + text.append("Prompt", style=self.prompt_color) + text.append("|") + text.append("System", style=self.system_color) + text.append("|") + text.append("Model", style=self.model_color) + text.append("|") + text.append("Reward", style=self.reward_color) + text.append(")") + print(text) + + +class TextEnvironment: + """ + The TextEnvironment enables interaction of a LLM with an environment using tools. + """ + + def __init__( + self, + model=None, + tokenizer=None, + tools=None, + reward_fn=None, + prompt=None, + max_turns=4, + max_tool_response=100, + max_length=None, + generation_kwargs=None, + ): + """ + Initialize TextEnvironment. + + Args: + model (`PreTrainedModelWrapper`): + The model to use for generation. + tokenizer (`transformers.PreTrainedTokenizer`): + The tokenizer to use for generation. + tools (list): + A list of tools to use for interaction. + reward_fn (function): + A function that takes a string and returns a reward. + prompt (str): + The base prompt to use for generation. Is prepended to the tasks. + max_turns (Optional[int]): + The maximum number of turns to allow. + max_tool_response (Optional[int]): + The maximum number of characters to allow in a tool response. + max_length (Optional[int]): + The maximum number of tokens to allow in an episode. + generation_kwargs (Optional[dict]): + A dictionary of keyword arguments to pass to the model's generate method. + """ + warnings.warn( + "This class is deprecated and will be removed in version 0.21.0. To enable tool use with LLMs, check out smolagents (https://huggingface.co/docs/smolagents/index)", + DeprecationWarning, + ) + self.model = model + self.tokenizer = tokenizer + self.prompt = prompt + if isinstance(tools, dict): + self.tools = tools + else: + self.tools = {tool.__class__.__name__: tool for tool in tools} + self.reward_fn = reward_fn + self.max_length = max_length + self.request_token = "" + self.call_token = "" + self.response_token = "" + self.submit_token = "" + self.max_turns = max_turns + self.max_tool_response = max_tool_response + + if generation_kwargs is None: + self.generation_kwargs = dict() + else: + self.generation_kwargs = generation_kwargs + + self.is_encoder_decoder = hasattr(self.model, "is_encoder_decoder") + self.current_device = extract_model_from_parallel(self.model).pretrained_model.device + + def run(self, queries, **rewards_kwargs): + """ + Run the environment on a list of queries. + + Args: + queries (list[str]): A list of queries to run the model in the environment on. + """ + turns = 0 + + queries = [self.prompt + task for task in queries] + queries_tokens = [ + self.tokenizer(query, return_tensors="pt").input_ids[0].to(self.model.pretrained_model.device) + for query in queries + ] + + histories = [TextHistory(q, qt, system=True) for q, qt in zip(queries, queries_tokens)] + + while any(not history.completed for history in histories) and turns < self.max_turns: + histories = self.generate(histories) + histories = self.tasks_end_check(histories) + # TODO: make this parallel rather than for-loop + for i in range(len(histories)): + histories[i] = self.step(histories[i]) + histories = self.tasks_end_check(histories, model_turn=False) + turns += 1 + self.compute_reward(histories, **rewards_kwargs) + + # convert a list of (q, r, m) tuples to lists of all qs, rs, and ms respectively + queries, responses, masks = map(list, zip(*[history.split_query_response_tokens() for history in histories])) + + rewards = [history.reward for history in histories] + return queries, responses, masks, rewards, histories + + def step(self, history): + """ + Step the environment forward one turn. + + Args: + history (`TextHistory`): The history to step forward. + """ + truncated, ended = self.task_end_check(history) + if ended: + history.complete(truncated=truncated) + if history.completed: + return history + + tool, query = self.parse_tool_call(history.last_text_segment) + if tool is None or query is None: + response = f"Unknown tool call: {history.last_text_segment}" + else: + if tool not in self.tools: + response = f"Unknown tool {tool}." + try: + response = self.tools[tool](query) + except Exception as error: + response = f"Tool error: {str(error)}" + + if len(response) > self.max_tool_response: + response = response[: (self.max_tool_response - 3)] + "..." + + history.append_segment( + response + self.response_token, + self.tokenizer(response + self.response_token, return_tensors="pt") + .input_ids[0] + .to(self.model.pretrained_model.device), + system=True, + ) + + return history + + def parse_tool_call(self, text): + """ + Parse request string. Expected format: query + """ + result = re.search(f"(?<={self.request_token}).*?(?={self.call_token})", text, re.DOTALL) + + # if we can't find a / span we return none + if result is None: + return None, None + else: + extracted_text = result.group() + + result = re.search(r"<(.*?)>", extracted_text) + + # if we can't find a tool name we return none + if result is None: + return None, None + else: + tool = result.group(1) + + # split off the tool name + query = ">".join(extracted_text.split(">")[1:]) + + return tool, query + + def compute_reward(self, histories, **reward_kwargs): + """ + Compute the reward for a list of histories. + """ + rewards = self.reward_fn([history.last_text_segment for history in histories], **reward_kwargs) + for history, reward in zip(histories, rewards): + history.reward = reward + return histories + + def generate(self, histories): + """ + Generate responses for a list of histories. + """ + active_histories = [i for i, history in enumerate(histories) if not history.completed] + + query_tensors = [histories[i].tokens for i in active_histories] + response_tensors = self._generate_batched(query_tensors) + response_texts = self.tokenizer.batch_decode(response_tensors) + + for i, response_text, response_tensor in zip(active_histories, response_texts, response_tensors): + histories[i].append_segment(response_text, response_tensor, system=False) + + return histories + + def tasks_end_check(self, histories, model_turn=True): + """ + Check if the current generation sequences have finished. + """ + for history in histories: + if not history.completed: + truncated, ended = self.task_end_check(history, model_turn=model_turn) + if ended: + history.complete(truncated=truncated) + return histories + + def task_end_check(self, history, model_turn=True): + """ + Check if the current generation sequence has finished. + """ + truncated = False + ended = False + if history.completed: + return truncated, ended + if self.max_length is not None and len(self.tokenizer(history.text).input_ids[0]) > self.max_length: + truncated = True + ended = True + elif self.tokenizer.eos_token in history.text: + ended = True + elif model_turn and not ( + (self.request_token in history.last_text_segment and self.call_token in history.last_text_segment) + or self.submit_token in history.last_text_segment + ): + ended = True + elif self.submit_token in history.last_text_segment: + ended = True + return truncated, ended + + def _generate_batched( + self, + query_tensors, + batch_size: int = 16, + pad_to_multiple_of: Optional[int] = None, + ): + """ + Generate responses for a list of query tensors. + + Args: + query_tensors (list[torch.Tensor]): A list of query tensors to generate responses for. + batch_size (int): The batch size to use for generation. + pad_to_multiple_of (int): The padding length to use for generation. + """ + outputs = [] + padding_side_default = self.tokenizer.padding_side + if not self.is_encoder_decoder: + self.tokenizer.padding_side = "left" + + # in case we have fewer examples than bs + batch_size = min(len(query_tensors), batch_size) + + for i in range(0, len(query_tensors), batch_size): + # prevent overflow if query tensors are not even multiple of bs + end_index = min(len(query_tensors), i + batch_size) + + batch = query_tensors[i:end_index] + batch_mask = [torch.ones_like(element) for element in batch] + inputs = {"input_ids": batch, "attention_mask": batch_mask} + + padded_inputs = self.tokenizer.pad( + inputs, + padding=True, + max_length=None, + pad_to_multiple_of=pad_to_multiple_of, + return_tensors="pt", + ).to(self.current_device) + + stopping_criteria = StringStoppingCriteria([self.call_token, self.submit_token], self.tokenizer) + + self.generation_kwargs["stopping_criteria"] = StoppingCriteriaList([stopping_criteria]) + + generations = extract_model_from_parallel(self.model).generate(**padded_inputs, **self.generation_kwargs) + + for generation, mask, generated_tokens in zip( + generations, padded_inputs["attention_mask"], stopping_criteria.generated_tokens + ): + if not self.is_encoder_decoder: + output = generation[(1 - mask).sum() :] # remove padding + else: + output = generation + + if not self.is_encoder_decoder: + output = output[(mask).sum() :] # remove prompt + + # remove chunk generated after stopping criteria in batch mode + outputs.append(output[:generated_tokens]) + self.tokenizer.padding_side = padding_side_default + return outputs diff --git a/lib/python3.12/site-packages/trl/extras/__init__.py b/lib/python3.12/site-packages/trl/extras/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fddcdf1af1a6353717f446e555faeebfaf169b08 --- /dev/null +++ b/lib/python3.12/site-packages/trl/extras/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ..import_utils import _LazyModule + + +_import_structure = { + "best_of_n_sampler": ["BestOfNSampler"], +} + +if TYPE_CHECKING: + from .best_of_n_sampler import BestOfNSampler +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/trl/extras/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/trl/extras/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..470731929fabe1de56cc43b0acb5b56180d28d5a Binary files /dev/null and b/lib/python3.12/site-packages/trl/extras/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/extras/__pycache__/best_of_n_sampler.cpython-312.pyc b/lib/python3.12/site-packages/trl/extras/__pycache__/best_of_n_sampler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79725d9a8cdc18de39473d8d60ce27b07e2cc219 Binary files /dev/null and b/lib/python3.12/site-packages/trl/extras/__pycache__/best_of_n_sampler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/extras/__pycache__/dataset_formatting.cpython-312.pyc b/lib/python3.12/site-packages/trl/extras/__pycache__/dataset_formatting.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b7ff3addea3c6bf11c2dfeae293ac47c4362766 Binary files /dev/null and b/lib/python3.12/site-packages/trl/extras/__pycache__/dataset_formatting.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/extras/__pycache__/profiling.cpython-312.pyc b/lib/python3.12/site-packages/trl/extras/__pycache__/profiling.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6518eee7bff0939afcde8404ec1f4a7abae46bfc Binary files /dev/null and b/lib/python3.12/site-packages/trl/extras/__pycache__/profiling.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/extras/__pycache__/vllm_client.cpython-312.pyc b/lib/python3.12/site-packages/trl/extras/__pycache__/vllm_client.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4a8ebbbb5c451ea5e8689ac502f286defac0bb5 Binary files /dev/null and b/lib/python3.12/site-packages/trl/extras/__pycache__/vllm_client.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/extras/best_of_n_sampler.py b/lib/python3.12/site-packages/trl/extras/best_of_n_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..9656e412b9a70b11a6d4d10fe4eb1c138a28fccd --- /dev/null +++ b/lib/python3.12/site-packages/trl/extras/best_of_n_sampler.py @@ -0,0 +1,132 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Callable, Optional, Union + +import torch +from transformers import GenerationConfig, PreTrainedTokenizer, PreTrainedTokenizerFast, set_seed + +from ..models import SUPPORTED_ARCHITECTURES, PreTrainedModelWrapper + + +class BestOfNSampler: + def __init__( + self, + model: PreTrainedModelWrapper, + tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], + queries_to_scores: Callable[[list[str]], list[float]], + length_sampler: Any, + sample_size: int = 4, + seed: Optional[int] = None, + n_candidates: int = 1, + generation_config: Optional[GenerationConfig] = None, + ) -> None: + r""" + Initialize the sampler for best-of-n generation + + Args: + model (`PreTrainedModelWrapper`): + The pretrained model to use for generation + tokenizer (`PreTrainedTokenizer` or `PreTrainedTokenizerFast`): + Tokenizer associated with the pretrained model + queries_to_scores (`Callable[[list[str]], list[float]]`): + Callable that takes a list of generated texts and returns the associated reward scores + length_sampler (`Any`): + Sampler used to sample the length of the generated text + sample_size (`int`): + Number of samples to generate for each query + seed (`int`, *optional*): + Random seed used to control generation + n_candidates (`int`): + Number of candidates to return for each query + generation_config (`GenerationConfig`, *optional*): + Generation config passed to the underlying model's `generate` method. See `GenerationConfig` + (https://huggingface.co/docs/transformers/v4.29.1/en/main_classes/text_generation#transformers.GenerationConfig) + for more details + """ + if seed is not None: + set_seed(seed) + + if not isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)): + raise ValueError( + f"tokenizer must be a PreTrainedTokenizer or PreTrainedTokenizerFast, got {type(tokenizer)}" + ) + if not isinstance(model, (SUPPORTED_ARCHITECTURES)): + raise ValueError( + f"model must be a PreTrainedModelWrapper, got {type(model)} - supported architectures are: {SUPPORTED_ARCHITECTURES}" + ) + + self.model = model + self.tokenizer = tokenizer + + self.queries_to_scores = queries_to_scores + self.length_sampler = length_sampler + self.gen_config = generation_config + self.sample_size = sample_size + self.n_candidates = n_candidates + + def generate( + self, + tokenized_query: Union[list[int], torch.Tensor, list[torch.Tensor], list[list[int]]], + skip_special_tokens: bool = True, + device: Optional[Union[str, torch.device]] = None, + **generation_kwargs, + ) -> list[list[str]]: + r""" + Generate the best of n samples for input queries + + Args: + tokenized_query (`list[int]` or `torch.Tensor` or `list[torch.Tensor]` or `list[int]`): + represents either a single tokenized query (a single tensor or a list of integers) or a batch of + tokenized queries (a list of tensors or a list of lists of integers) + skip_special_tokens (`bool`): + Whether to remove the special tokens from the output + device (`str` or `torch.device`, *optional*): + The device on which the model will be loaded + **generation_kwargs (`dict`, *optional*): + Additional keyword arguments passed along to the underlying model's `generate` method. This is used to + override generation config + + Returns: + list[list[str]]: A list of lists of generated texts + """ + queries = None + + if isinstance(tokenized_query, torch.Tensor) and tokenized_query.ndim == 1: + queries = tokenized_query.unsqueeze(0) + elif isinstance(tokenized_query, list): + element_type = type(tokenized_query[0]) + if element_type is int: + queries = torch.tensor(tokenized_query).unsqueeze(0) + elif element_type is torch.Tensor: + queries = [tensor.reshape((1, -1)) for tensor in tokenized_query] + else: + queries = [torch.tensor(query).reshape((1, -1)) for query in tokenized_query] + + result = [] + + for query in queries: + queries = query.repeat((self.sample_size, 1)) + output = self.model.generate( + queries.to(device), + max_new_tokens=self.length_sampler(), + generation_config=self.gen_config, + **generation_kwargs, + ).squeeze() + output = self.tokenizer.batch_decode(output, skip_special_tokens=skip_special_tokens) + scores = torch.tensor(self.queries_to_scores(output)) + output = [output[i] for i in scores.topk(self.n_candidates).indices] + result.append(output) + + return result diff --git a/lib/python3.12/site-packages/trl/extras/dataset_formatting.py b/lib/python3.12/site-packages/trl/extras/dataset_formatting.py new file mode 100644 index 0000000000000000000000000000000000000000..d1614ba89c2e3df7a0e6455b49e593beeae50d2c --- /dev/null +++ b/lib/python3.12/site-packages/trl/extras/dataset_formatting.py @@ -0,0 +1,106 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from typing import Callable, Literal, Optional, Union + +from datasets import Dataset, Value +from transformers import AutoTokenizer + +from ..trainer.utils import ConstantLengthDataset + + +FORMAT_MAPPING = { + "chatml": [{"content": Value(dtype="string", id=None), "role": Value(dtype="string", id=None)}], + "instruction": {"completion": Value(dtype="string", id=None), "prompt": Value(dtype="string", id=None)}, +} + + +def conversations_formatting_function( + tokenizer: AutoTokenizer, messages_field: Literal["messages", "conversations"], tools: Optional[list] = None +): + r""" + return a callable function that takes in a "messages" dataset and returns a formatted dataset, based on the + tokenizer apply chat template to the dataset along with the schema of the list of functions in the tools list. + """ + + def format_dataset(examples): + if isinstance(examples[messages_field][0], list): + output_texts = [] + for i in range(len(examples[messages_field])): + output_texts.append( + tokenizer.apply_chat_template(examples[messages_field][i], tokenize=False, tools=tools) + ) + return output_texts + else: + return tokenizer.apply_chat_template(examples[messages_field], tokenize=False, tools=tools) + + return format_dataset + + +def instructions_formatting_function(tokenizer: AutoTokenizer): + r""" + return a callable function that takes in an "instructions" dataset and returns a formatted dataset, based on the + tokenizer apply chat template to the dataset + """ + + def format_dataset(examples): + if isinstance(examples["prompt"], list): + output_texts = [] + for i in range(len(examples["prompt"])): + converted_sample = [ + {"role": "user", "content": examples["prompt"][i]}, + {"role": "assistant", "content": examples["completion"][i]}, + ] + output_texts.append(tokenizer.apply_chat_template(converted_sample, tokenize=False)) + return output_texts + else: + converted_sample = [ + {"role": "user", "content": examples["prompt"]}, + {"role": "assistant", "content": examples["completion"]}, + ] + return tokenizer.apply_chat_template(converted_sample, tokenize=False) + + return format_dataset + + +def get_formatting_func_from_dataset( + dataset: Union[Dataset, ConstantLengthDataset], tokenizer: AutoTokenizer, tools: Optional[list] = None +) -> Optional[Callable]: + r""" + Finds the correct formatting function based on the dataset structure. Currently supported datasets are: + - `ChatML` with [{"role": str, "content": str}] + - `instruction` with [{"prompt": str, "completion": str}] + + Args: + dataset (Dataset): User dataset + tokenizer (AutoTokenizer): Tokenizer used for formatting + + Returns: + Callable: Formatting function if the dataset format is supported else None + """ + if isinstance(dataset, Dataset): + if "messages" in dataset.features: + if dataset.features["messages"] == FORMAT_MAPPING["chatml"]: + logging.info("Formatting dataset with chatml format") + return conversations_formatting_function(tokenizer, "messages", tools) + if "conversations" in dataset.features: + if dataset.features["conversations"] == FORMAT_MAPPING["chatml"]: + logging.info("Formatting dataset with chatml format") + return conversations_formatting_function(tokenizer, "conversations", tools) + elif dataset.features == FORMAT_MAPPING["instruction"]: + logging.info("Formatting dataset with instruction format") + return instructions_formatting_function(tokenizer) + + return None diff --git a/lib/python3.12/site-packages/trl/extras/profiling.py b/lib/python3.12/site-packages/trl/extras/profiling.py new file mode 100644 index 0000000000000000000000000000000000000000..8e4c9e9188ce4b46afd7b35debbc7e08e70aa986 --- /dev/null +++ b/lib/python3.12/site-packages/trl/extras/profiling.py @@ -0,0 +1,100 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextlib +import functools +import time +from collections.abc import Generator + +from transformers import Trainer +from transformers.integrations import is_mlflow_available, is_wandb_available + + +if is_wandb_available(): + import wandb + +if is_mlflow_available(): + import mlflow + + +@contextlib.contextmanager +def profiling_context(trainer: Trainer, name: str) -> Generator[None, None, None]: + """ + A context manager function for profiling a block of code. Results are logged to Weights & Biases or MLflow + depending on the trainer's configuration. + + Args: + trainer (`~transformers.Trainer`): + Trainer object. + name (`str`): + Name of the block to be profiled. Used as a key in the logged dictionary. + + Example: + ```python + from transformers import Trainer + from trl.extras.profiling import profiling_context + + + class MyTrainer(Trainer): + def some_method(self): + A = np.random.rand(1000, 1000) + B = np.random.rand(1000, 1000) + with profiling_context(self, "matrix_multiplication"): + # Code to profile: simulate a computationally expensive operation + result = A @ B # Matrix multiplication + ``` + """ + start_time = time.perf_counter() + yield + end_time = time.perf_counter() + duration = end_time - start_time + + profiling_metrics = {f"profiling/Time taken: {trainer.__class__.__name__}.{name}": duration} + if "wandb" in trainer.args.report_to and wandb.run is not None and trainer.accelerator.is_main_process: + wandb.log(profiling_metrics) + + if "mlflow" in trainer.args.report_to and mlflow.run is not None and trainer.accelerator.is_main_process: + mlflow.log_metrics(profiling_metrics, step=trainer.state.global_step) + + +def profiling_decorator(func: callable) -> callable: + """ + Decorator to profile a function and log execution time using [`extras.profiling.profiling_context`]. + + Args: + func (`callable`): + Function to be profiled. + + Example: + ```python + from transformers import Trainer + from trl.extras.profiling import profiling_decorator + + + class MyTrainer(Trainer): + @profiling_decorator + def some_method(self): + A = np.random.rand(1000, 1000) + B = np.random.rand(1000, 1000) + # Code to profile: simulate a computationally expensive operation + result = A @ B + ``` + """ + + @functools.wraps(func) + def wrapper(self, *args, **kwargs): + with profiling_context(self, func.__name__): + return func(self, *args, **kwargs) + + return wrapper diff --git a/lib/python3.12/site-packages/trl/extras/vllm_client.py b/lib/python3.12/site-packages/trl/extras/vllm_client.py new file mode 100644 index 0000000000000000000000000000000000000000..71341a99542c6411b19cf75fc6fbbcbfdd68627d --- /dev/null +++ b/lib/python3.12/site-packages/trl/extras/vllm_client.py @@ -0,0 +1,337 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import atexit +import logging +import socket +import time +from typing import Optional +from urllib.parse import urlparse + +import torch +from torch import nn + +from ..import_utils import is_requests_available, is_vllm_ascend_available, is_vllm_available + + +if is_requests_available(): + import requests + from requests import ConnectionError + + +if is_vllm_available(): + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup + + if is_vllm_ascend_available(): + from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator + + +logger = logging.getLogger(__name__) + + +class VLLMClient: + """ + A client class to interact with a vLLM server. + + This class provides methods to generate completions, initialize and manage weight update groups, and update model + weights in a distributed setting. Before using it, start the vLLM server with `trl vllm-serve`. + + Args: + base_url (`str` or `None`, *optional*, defaults to `None`): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `host` and `server_port` are + ignored. + host (`str`, *optional*, defaults to `"0.0.0.0"`): + IP address of the vLLM server. Ignored if `base_url` is provided. + server_port (`int`, *optional*, defaults to `8000`): + Port number of the vLLM server. Ignored if `base_url` is provided. + group_port (`int`, *optional*, defaults to `51216`): + Port number for the weight update group. + connection_timeout (`float`, *optional*, defaults to `0.0`): + Total timeout duration in seconds to wait for the server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + + Examples: + Run the vLLM server with the model `Qwen/Qwen2.5-7B`: + + ``` + $ trl vllm-serve --model Qwen/Qwen2.5-7B + ... + INFO: Application startup complete. + INFO: Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit) + ``` + + Use the client to generate completions and update model weights: + + ```python + >>> from trl.extras.vllm_client import VLLMClient + + >>> client = VLLMClient() + >>> client.generate(["Hello, AI!", "Tell me a joke"]) + [[2980, 498, 1492, 752, 448, 264, 13027, 8645, 30, 358, 2776, 4460, 311, 3270, 264, 2025], + [911, 7988, 1251, 382, 3838, 653, 498, 1618, 4325, 879, 2581, 20027, 264, 21428, 30, 362]] + + >>> from transformers import AutoModelForCausalLM + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B", device_map="cuda") + >>> client.init_communicator() + >>> client.update_model_params(model) + ``` + + There are several ways to initialize the client: + + ```python + VLLMClient(base_url="http://localhost:8000") + VLLMClient(base_url="http://192.168.1.100:8000") + VLLMClient(host="localhost", server_port=8000) + VLLMClient(host="192.168.1.100", server_port=8000) + ``` + """ + + def __init__( + self, + base_url: Optional[str] = None, + host: str = "0.0.0.0", + server_port: int = 8000, + group_port: int = 51216, + connection_timeout: float = 0.0, + ): + if not is_requests_available(): + raise ImportError("requests is not installed. Please install it with `pip install requests`.") + if not is_vllm_available(): + raise ImportError("vLLM is not installed. Please install it with `pip install vllm`.") + + self.session = requests.Session() + + if base_url is not None: + # Parse the base_url to extract host and port + parsed_url = urlparse(base_url) + self.host = socket.gethostbyname(parsed_url.hostname) + scheme = parsed_url.scheme or "http" + self.base_url = f"{scheme}://{parsed_url.netloc}{parsed_url.path}" + else: + self.host = host + self.server_port = server_port + self.base_url = f"http://{self.host}:{self.server_port}" + self.group_port = group_port + self.check_server(connection_timeout) # check server and fail after timeout + + def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0): + """ + Check server availability with retries on failure, within a total timeout duration. If the server is not up + after the total timeout duration, raise a `ConnectionError`. + + Args: + retry_interval (`float`, *optional*, defaults to `2.0`): + Interval in seconds between retries. + total_timeout (`float`, *optional*, defaults to `0.0`): + Total timeout duration in seconds. + """ + url = f"{self.base_url}/health/" + start_time = time.time() # Record the start time + + while True: + try: + response = requests.get(url) + except requests.exceptions.RequestException as exc: + # Check if the total timeout duration has passed + elapsed_time = time.time() - start_time + if elapsed_time >= total_timeout: + raise ConnectionError( + f"The vLLM server can't be reached at {self.base_url} after {total_timeout} seconds. Make " + "sure the server is running by running `trl vllm-serve`." + ) from exc + else: + if response.status_code == 200: + if "X-Forwarded-For" in response.headers: + self.host = response.headers["X-Forwarded-For"] + logger.info("Server is up!") + return None + + # Retry logic: wait before trying again + logger.info(f"Server is not up yet. Retrying in {retry_interval} seconds...") + time.sleep(retry_interval) + + def generate( + self, + prompts: list[str], + n: int = 1, + repetition_penalty: float = 1.0, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + max_tokens: int = 16, + guided_decoding_regex: Optional[str] = None, + generation_kwargs: Optional[dict] = None, + ) -> list[list[int]]: + """ + Generates model completions for the provided prompts. + + Args: + prompts (`list[str]`): + List of text prompts for which the model will generate completions. + n (`int`, *optional*, defaults to `1`): + Number of completions to generate for each prompt. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Parameter for repetition penalty. 1.0 means no penalty. + temperature (`float`, *optional*, defaults to `1.0`): + Temperature parameter for sampling. Higher values increase diversity. + top_p (`float`, *optional*, defaults to `1.0`): + Top-p sampling parameter.`1.0` means no truncation. + top_k (`int`, *optional*, defaults to `-1`): + Top-k sampling parameter. `-1` means no truncation. + min_p (`float`, *optional*, defaults to `0.0`): + Minimum probability for sampling. + max_tokens (`int`, *optional*, defaults to `16`): + Maximum number of tokens to generate for each prompt. + guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`): + Regular expression to guide the decoding process. + generation_kwargs (`dict` or `None`, *optional*, defaults to `None`): + Additional generation parameters to pass to the vLLM `SamplingParams`. This can include parameters like + `seed`, `frequency_penalty`, etc. If it contains keys that conflict with the other parameters, they + will override them. + + Returns: + `list[list[int]]`: + List of lists of token IDs representing the model-generated completions for each prompt. + """ + url = f"{self.base_url}/generate/" + response = self.session.post( + url, + json={ + "prompts": prompts, + "n": n, + "repetition_penalty": repetition_penalty, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_tokens": max_tokens, + "guided_decoding_regex": guided_decoding_regex, + "generation_kwargs": generation_kwargs or {}, + }, + ) + if response.status_code == 200: + return response.json()["completion_ids"] + else: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + def init_communicator(self): + """ + Initializes the weight update group in a distributed setup for model synchronization. + """ + # Get the world size from the server + url = f"{self.base_url}/get_world_size/" + response = requests.get(url) + if response.status_code == 200: + vllm_world_size = response.json()["world_size"] + else: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + world_size = vllm_world_size + 1 # add the client to the world + self.rank = vllm_world_size # the client's rank is the last process + + # Initialize weight update group + url = f"{self.base_url}/init_communicator/" + # In the server side, the host is set to 0.0.0.0 + response = self.session.post(url, json={"host": "0.0.0.0", "port": self.group_port, "world_size": world_size}) + if response.status_code != 200: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + # Brief delay to allow server initialization. While not strictly required (client socket will retry on + # connection failure), this prevents log warnings like: + # [W416 23:24:57.460001114 socket.cpp:204] [c10d] The hostname of the client socket cannot be retrieved. err=-3 + time.sleep(0.1) + + # Set up the communication group for weight broadcasting + pg = StatelessProcessGroup.create(host=self.host, port=self.group_port, rank=self.rank, world_size=world_size) + self.pynccl_comm = PyNcclCommunicator(pg, device=0) + + # When the client object is deleted, close the weight update group + atexit.register(self.close_communicator) + + def update_named_param(self, name: str, weights: torch.Tensor): + """ + Updates a specific named parameter in the model and broadcasts it to other processes. + + Args: + name (`str`): + Name of the layer whose weights are being updated. + weights (`torch.Tensor`): + Tensor containing the updated weights. + """ + dtype, shape = str(weights.dtype), tuple(weights.shape) + url = f"{self.base_url}/update_named_param/" + response = self.session.post(url, json={"name": name, "dtype": dtype, "shape": shape}) + if response.status_code != 200: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + # Broadcast the weights to the other processes + self.pynccl_comm.broadcast(weights, src=self.rank) + self.pynccl_comm.group.barrier() + + def update_model_params(self, model: nn.Module): + """ + Updates all parameters of the given model by calling `update_named_param` for each parameter in the model. + + Args: + model (`nn.Module`): + Model whose parameters (weights/biases) are to be updated. + """ + for name, param in model.named_parameters(): + # Update each parameter individually + self.update_named_param(name, param.data) + + def reset_prefix_cache(self): + """ + Resets the prefix cache for the model. + """ + url = f"{self.base_url}/reset_prefix_cache/" + response = self.session.post(url) + if response.status_code != 200: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + def close_communicator(self): + """ + Closes the weight update group and cleans up the communication group. + """ + url = f"{self.base_url}/close_communicator/" + + try: + response = self.session.post(url) + except ConnectionError: + # The server might be already down, so we don't need to close the communicator + pass + else: + if response.status_code != 200: + raise Exception(f"Request failed: {response.status_code}, {response.text}") + + +# Example usage +if __name__ == "__main__": + from vllm import SamplingParams + + client = VLLMClient() + client.init_communicator() + + # Generate completions + responses = client.generate(["Hello, AI!", "Tell me a joke"], n=4, max_tokens=32, sampling_params=SamplingParams()) + print("Responses:", responses) # noqa + + # Update model weights + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B").to("cuda") + client.update_model_params(model) diff --git a/lib/python3.12/site-packages/trl/import_utils.py b/lib/python3.12/site-packages/trl/import_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..dd594e36ef6b4da68e3ee8f7ab59125e8ea53645 --- /dev/null +++ b/lib/python3.12/site-packages/trl/import_utils.py @@ -0,0 +1,156 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import os +from itertools import chain +from types import ModuleType +from typing import Any + +from packaging import version +from transformers.utils.import_utils import _is_package_available + + +LIGER_KERNEL_MIN_VERSION = "0.5.8" + +# Use same as transformers.utils.import_utils +_deepspeed_available = _is_package_available("deepspeed") +_diffusers_available = _is_package_available("diffusers") +_fastapi_available = _is_package_available("fastapi") +_is_liger_kernel_available, _liger_kernel_version = _is_package_available("liger_kernel", return_version=True) +_llm_blender_available = _is_package_available("llm_blender") +_mergekit_available = _is_package_available("mergekit") +_pydantic_available = _is_package_available("pydantic") +_requests_available = _is_package_available("requests") +_unsloth_available = _is_package_available("unsloth") +_uvicorn_available = _is_package_available("uvicorn") +_vllm_available = _is_package_available("vllm") +_vllm_ascend_available = _is_package_available("vllm_ascend") +_joblib_available = _is_package_available("joblib") + + +def is_deepspeed_available() -> bool: + return _deepspeed_available + + +def is_diffusers_available() -> bool: + return _diffusers_available + + +def is_fastapi_available() -> bool: + return _fastapi_available + + +def is_liger_kernel_available(min_version: str = LIGER_KERNEL_MIN_VERSION) -> bool: + return _is_liger_kernel_available and version.parse(_liger_kernel_version) >= version.parse(min_version) + + +def is_llm_blender_available() -> bool: + return _llm_blender_available + + +def is_mergekit_available() -> bool: + return _mergekit_available + + +def is_pydantic_available() -> bool: + return _pydantic_available + + +def is_requests_available() -> bool: + return _requests_available + + +def is_unsloth_available() -> bool: + return _unsloth_available + + +def is_uvicorn_available() -> bool: + return _uvicorn_available + + +def is_vllm_available() -> bool: + return _vllm_available + + +def is_vllm_ascend_available() -> bool: + return _vllm_ascend_available + + +def is_joblib_available() -> bool: + return _joblib_available + + +class _LazyModule(ModuleType): + """ + Module class that surfaces all objects but only performs associated imports when the objects are requested. + """ + + # Very heavily inspired by optuna.integration._IntegrationModule + # https://github.com/optuna/optuna/blob/master/optuna/integration/__init__.py + def __init__(self, name, module_file, import_structure, module_spec=None, extra_objects=None): + super().__init__(name) + self._modules = set(import_structure.keys()) + self._class_to_module = {} + for key, values in import_structure.items(): + for value in values: + self._class_to_module[value] = key + # Needed for autocompletion in an IDE + self.__all__ = list(import_structure.keys()) + list(chain(*import_structure.values())) + self.__file__ = module_file + self.__spec__ = module_spec + self.__path__ = [os.path.dirname(module_file)] + self._objects = {} if extra_objects is None else extra_objects + self._name = name + self._import_structure = import_structure + + # Needed for autocompletion in an IDE + def __dir__(self): + result = super().__dir__() + # The elements of self.__all__ that are submodules may or may not be in the dir already, depending on whether + # they have been accessed or not. So we only add the elements of self.__all__ that are not already in the dir. + for attr in self.__all__: + if attr not in result: + result.append(attr) + return result + + def __getattr__(self, name: str) -> Any: + if name in self._objects: + return self._objects[name] + if name in self._modules: + value = self._get_module(name) + elif name in self._class_to_module.keys(): + module = self._get_module(self._class_to_module[name]) + value = getattr(module, name) + else: + raise AttributeError(f"module {self.__name__} has no attribute {name}") + + setattr(self, name, value) + return value + + def _get_module(self, module_name: str): + try: + return importlib.import_module("." + module_name, self.__name__) + except Exception as e: + raise RuntimeError( + f"Failed to import {self.__name__}.{module_name} because of the following error (look up to see its" + f" traceback):\n{e}" + ) from e + + def __reduce__(self): + return (self.__class__, (self._name, self.__file__, self._import_structure)) + + +class OptionalDependencyNotAvailable(BaseException): + """Internally used error class for signalling an optional dependency was not found.""" diff --git a/lib/python3.12/site-packages/trl/mergekit_utils.py b/lib/python3.12/site-packages/trl/mergekit_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a2260079a84f9c8dcfe69114a671709791cacf03 --- /dev/null +++ b/lib/python3.12/site-packages/trl/mergekit_utils.py @@ -0,0 +1,281 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +from huggingface_hub import HfApi + +from trl.import_utils import is_mergekit_available + + +if is_mergekit_available(): + from mergekit.config import MergeConfiguration + from mergekit.merge import MergeOptions, run_merge + + +def upload_model_to_hf(folder_path: str, repo_id: str): + api = HfApi() + # Create the repository if it doesn't exist + repo = api.create_repo(repo_id, repo_type="model") + + # Upload the folder to the specified repository + api.upload_folder( + folder_path=folder_path, + repo_id=repo.repo_id, + repo_type=repo.repo_type, + ) + + +class MergeConfig: + r""" + Configuration class for merging two models using `mergekit`. + + This class provides a structured way to configure and generate merge configurations for various merge methods, such + as `linear`, `ties`, `dare_ties`, and `slerp`. + + Args: + method (`str`, *optional*, defaults to `"linear"`): + Merge method to use. Supported methods include: + + - `"linear"`: Linearly combines two models with specified weights. + - `"ties"`: Combines two models using the TIES method with density parameters. + - `"dare_ties"`: A variant of TIES for domain adaptation. + - `"slerp"`: Combines models using spherical linear interpolation. + + Note: + + For more details about the merge methods and how they are implemented, see the [MergeKit GitHub + repository](https://github.com/arcee-ai/mergekit?tab=readme-ov-file#merge-methods). + + Attributes: + method (`str`): The merge method to use. + policy_model_path (`str` or `None`): Path to the policy model. + target_model_path (`str` or `None`): Path to the target model. + policy_model_weight (`float`): Weight for the policy model (for `linear` and `ties` methods). + target_model_weight (`float`): Weight for the target model (for `linear` and `ties` methods). + policy_model_density (`list[float]`): Density parameters for the policy model (for `ties` and `dare_ties`). + target_model_density (`list[float]`): Density parameters for the target model (for `ties` and `dare_ties`). + normalize (`float` or `None`): Normalization factor for the TIES method. + t_values (`float` or `None`): Interpolation factor for the SLERP method. + dtype (`str`): Data type to use for merging, e.g., `"float16"`. + """ + + def __init__(self, method: str = "linear"): + if not is_mergekit_available(): + raise ImportError("MergeConfig requires the `mergekit` extra. To install, run `pip install mergekit`.") + self.method = method + self.policy_model_path = None + self.target_model_path = None + + # Initialize relevant parameters based on the method + if method == "linear": + self.policy_model_weight = 0.5 + self.target_model_weight = 0.5 + self.dtype = "float16" + elif method == "ties": + self.policy_model_weight = 1.0 + self.policy_model_density = [1.0, 0.7, 0.1] + self.target_model_weight = 1.0 + self.target_model_density = [1.0] + self.normalize = 1.0 + self.dtype = "float16" + elif method == "dare_ties": + self.policy_model_weight = 1.0 + self.policy_model_density = [1.0, 0.7, 0.1] + self.target_model_weight = 1.0 + self.target_model_density = [1.0] + self.normalize = 1.0 + self.dtype = "float16" + elif method == "slerp": + self.t_values = 0.5 + self.dtype = "float16" + else: + raise ValueError(f"Unsupported merge method: {method}") + + def create_merge_config_linear(self) -> "MergeConfiguration": + """ + Creates a merge configuration for a linear merge of two models with specified weights. + """ + # Create the merge configuration dictionary + merge_config_dict = { + "dtype": self.dtype, + "merge_method": "linear", + "models": [ + {"model": self.policy_model_path, "parameters": {"weight": self.policy_model_weight}}, + {"model": self.target_model_path, "parameters": {"weight": self.target_model_weight}}, + ], + } + + # Create the MergeConfiguration from the dictionary + merge_config = MergeConfiguration.model_validate(merge_config_dict) + + return merge_config + + def create_merge_config_ties(self) -> "MergeConfiguration": + """ + Creates a merge configuration for a TIES merge of two models, with specified weights and densities. + """ + # Create the TIES merge configuration dictionary + merge_config_dict = { + "merge_method": "ties", + "slices": None, # Optional slices if needed + "models": [ + { + "model": { + "model": {"path": self.target_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": {"density": self.target_model_density, "weight": self.target_model_weight}, + }, + { + "model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": {"density": self.policy_model_density, "weight": self.policy_model_weight}, + }, + ], + "parameters": {"normalize": self.normalize}, + "base_model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "dtype": self.dtype, + "tokenizer_source": None, + "tokenizer": None, + "chat_template": None, + "out_dtype": None, + } + + # Create the MergeConfiguration from the dictionary + merge_config = MergeConfiguration.model_validate(merge_config_dict) + + return merge_config + + def create_merge_config_dare_ties(self) -> "MergeConfiguration": + """ + Creates a merge configuration for a DARE TIES merge of two models, with specified weights and densities. + """ + # Create the DARE TIES merge configuration dictionary + merge_config_dict = { + "merge_method": "dare_ties", + "slices": None, # Optional slices if needed + "models": [ + { + "model": { + "model": {"path": self.target_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": {"density": self.target_model_density, "weight": self.target_model_weight}, + }, + { + "model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": {"density": self.policy_model_density, "weight": self.policy_model_weight}, + }, + ], + "parameters": {"normalize": self.normalize}, + "base_model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "dtype": self.dtype, + "tokenizer_source": None, + "tokenizer": None, + "chat_template": None, + "out_dtype": None, + } + + # Create the MergeConfiguration from the dictionary + merge_config = MergeConfiguration.model_validate(merge_config_dict) + + return merge_config + + def create_merge_config_slerp(self) -> "MergeConfiguration": + """ + Creates a merge configuration for a SLERP merge of a model with a base model. + """ + + # Create the SLERP merge configuration dictionary + merge_config_dict = { + "merge_method": "slerp", + "slices": None, # Optional slices if needed + "models": [ + { + "model": { + "model": {"path": self.target_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "parameters": None, # No specific parameters for SLERP model + } + ], + "parameters": { + "t": self.t_values # Set the t values for SLERP + }, + "base_model": { + "model": {"path": self.policy_model_path, "revision": None}, + "lora": None, + "override_architecture": None, + }, + "dtype": self.dtype, + "tokenizer_source": None, + "tokenizer": None, + "chat_template": None, + "out_dtype": None, + } + + # Create the MergeConfiguration from the dictionary + merge_config = MergeConfiguration.model_validate(merge_config_dict) + + return merge_config + + def create(self) -> "MergeConfiguration": + if self.method == "linear": + return self.create_merge_config_linear() + elif self.method == "ties": + return self.create_merge_config_ties() + elif self.method == "dare_ties": + return self.create_merge_config_dare_ties() + elif self.method == "slerp": + return self.create_merge_config_slerp() + + +def merge_models(config: MergeConfig, out_path: str): + """ + Merge two models using mergekit + + Args: + config (`MergeConfig`): The merge configuration. + out_path (`str`): The output path for the merged model. + """ + if not is_mergekit_available(): + raise ImportError("merge_models requires the `mergekit` extra. To install, run `pip install mergekit`.") + run_merge( + config, + out_path=out_path, + options=MergeOptions( + cuda=torch.cuda.is_available(), + copy_tokenizer=True, + lazy_unpickle=False, + low_cpu_memory=False, + ), + ) diff --git a/lib/python3.12/site-packages/trl/models/__init__.py b/lib/python3.12/site-packages/trl/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5635e31dfdd2796588f3b20c1adb5428d4a8c2f5 --- /dev/null +++ b/lib/python3.12/site-packages/trl/models/__init__.py @@ -0,0 +1,75 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ..import_utils import OptionalDependencyNotAvailable, _LazyModule, is_diffusers_available + + +_import_structure = { + "activation_offloading": ["get_act_offloading_ctx_manager"], + "modeling_base": ["GeometricMixtureWrapper", "PreTrainedModelWrapper", "create_reference_model"], + "modeling_value_head": ["AutoModelForCausalLMWithValueHead", "AutoModelForSeq2SeqLMWithValueHead"], + "utils": [ + "SUPPORTED_ARCHITECTURES", + "clone_chat_template", + "prepare_deepspeed", + "prepare_fsdp", + "setup_chat_format", + "unwrap_model_for_generation", + ], +} + +try: + if not is_diffusers_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["modeling_sd_base"] = [ + "DDPOPipelineOutput", + "DDPOSchedulerOutput", + "DDPOStableDiffusionPipeline", + "DefaultDDPOStableDiffusionPipeline", + ] + +if TYPE_CHECKING: + from .activation_offloading import get_act_offloading_ctx_manager + from .modeling_base import GeometricMixtureWrapper, PreTrainedModelWrapper, create_reference_model + from .modeling_value_head import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead + from .utils import ( + SUPPORTED_ARCHITECTURES, + clone_chat_template, + prepare_deepspeed, + prepare_fsdp, + setup_chat_format, + unwrap_model_for_generation, + ) + + try: + if not is_diffusers_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .modeling_sd_base import ( + DDPOPipelineOutput, + DDPOSchedulerOutput, + DDPOStableDiffusionPipeline, + DefaultDDPOStableDiffusionPipeline, + ) +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/trl/models/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/trl/models/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..005aa54e99d7ccf0886108970799f32b3fe1ff79 Binary files /dev/null and b/lib/python3.12/site-packages/trl/models/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/models/__pycache__/activation_offloading.cpython-312.pyc b/lib/python3.12/site-packages/trl/models/__pycache__/activation_offloading.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5447afc03504197be4fb4051422684cc5c88b779 Binary files /dev/null and b/lib/python3.12/site-packages/trl/models/__pycache__/activation_offloading.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/models/__pycache__/auxiliary_modules.cpython-312.pyc b/lib/python3.12/site-packages/trl/models/__pycache__/auxiliary_modules.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1166837e8be09757aa496e046310e4910e80af19 Binary files /dev/null and b/lib/python3.12/site-packages/trl/models/__pycache__/auxiliary_modules.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/models/__pycache__/modeling_base.cpython-312.pyc b/lib/python3.12/site-packages/trl/models/__pycache__/modeling_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..522aec96291b875019e85f6c74ce582b42ea2813 Binary files /dev/null and b/lib/python3.12/site-packages/trl/models/__pycache__/modeling_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/models/__pycache__/modeling_sd_base.cpython-312.pyc b/lib/python3.12/site-packages/trl/models/__pycache__/modeling_sd_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..975c8e1f6690b2e1d03424f1dc208f4858f81d19 Binary files /dev/null and b/lib/python3.12/site-packages/trl/models/__pycache__/modeling_sd_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/models/__pycache__/modeling_value_head.cpython-312.pyc b/lib/python3.12/site-packages/trl/models/__pycache__/modeling_value_head.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f65f70a3505bde5b69b4b38c3f1cdca0819595c Binary files /dev/null and b/lib/python3.12/site-packages/trl/models/__pycache__/modeling_value_head.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/models/__pycache__/sd_utils.cpython-312.pyc b/lib/python3.12/site-packages/trl/models/__pycache__/sd_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b8d058a0a739fe375f2acfa0cdc93406716a10f Binary files /dev/null and b/lib/python3.12/site-packages/trl/models/__pycache__/sd_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/models/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/trl/models/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ad3058cafdfde50eab11d55a703bc08461f34d9 Binary files /dev/null and b/lib/python3.12/site-packages/trl/models/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/models/activation_offloading.py b/lib/python3.12/site-packages/trl/models/activation_offloading.py new file mode 100644 index 0000000000000000000000000000000000000000..048ab086c6d5cfa7b528a2e948353ff51313a1bc --- /dev/null +++ b/lib/python3.12/site-packages/trl/models/activation_offloading.py @@ -0,0 +1,464 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of https://github.com/pytorch/torchtune. + +import warnings + +import psutil +import torch +from torch import nn +from torch.autograd.graph import saved_tensors_hooks + + +class OffloadActivations(saved_tensors_hooks): + """ + Context manager under which activation tensors created in the forward pass will be offloaded. + + Enable the memory efficiency technique of activation offloading, where activations bigger than `min_offload_size` + bytes will be offloaded to CPU in the forward and brought back in the backward. This is in contrast to maintaining + the activation on GPU VRAM throughout the program. + + This manager contains the option of using one additional CUDA stream to handle the communication between CUDA and + CPU, which is intended to overlap with the default computation stream to improve runtime. We designed + synchronization with a few heuristics for optimizing the tradeoff between runtime vs memory usage. + + Args: + use_pin_memory (`bool`, *optional*, defaults to `True`): + Whether to offloaded Tensor will be placed in pinned memory on the CPU. Pinned memory allows the Tensor to + be moved back onto GPU more quickly but is a limited resource. + use_streams (`bool`, *optional*, defaults to `True`): + Whether to use streams for performance optimization where the communications get overlapped with the + computation. Requires a torch build after torch-2.5.0. + min_offload_size (`int`, *optional*, defaults to `1024`): + Minimum number of bytes a Tensor must be in order to qualify for offloading. If the tensor is too small, we + do not want to waste bandwidth and resources moving it to CPU and back. + max_fwd_stash_size (`int`, *optional*, defaults to `5`): + Maximum size of the forward stash, or the maximum number of consecutive activations to keep alive during + the forward pass. This number must be at least 1. Keeping alive more activations will potentially allow + more overlap between the communication and compute streams at the cost of increasing memory usage. Keeping + alive fewer activations will conserve memory, but may cause poor overlap between the streams, increasing + runtime. + + Raises: + ValueError: if `max_fwd_stash_size` is not at least `1`. + + Example: + ```python + >>> with OffloadActivations(): + ... outputs = model(inputs, labels=labels) + >>> loss = outputs.loss + >>> loss.backward() + ``` + """ + + def __init__( + self, + use_pin_memory: bool = True, + use_streams: bool = True, + min_offload_size: int = 1024, + max_fwd_stash_size: int = 5, + ) -> None: + self.use_streams = use_streams + + self.min_tensor_size_bytes = min_offload_size # we don't want to bother with small tensors + self.tracker = {} # tensor_id => (new_tensor, if_modified) ---> track what saved/offloaded tensors are where + self.tensor_id = 0 + self.is_first_forward_call = True + self.is_first_backward_call = True + self.is_first_forward_pass = True + + # Managing cpu memory + self.use_pin_memory = use_pin_memory + self.virtual_memory_safe_pct = 60 # we should not exceed this percentage of memory + + self.accelerator_type = ( + torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda" + ) + # NOTE: xpu doesn't have `default_stream` API, use `current_stream` instead + self.s0 = ( + torch.xpu.current_stream() if self.accelerator_type == "xpu" else torch.cuda.default_stream() + ) # comp stream + + # For streaming + if self.use_streams: + self.s1 = torch.Stream() if self.accelerator_type == "xpu" else torch.cuda.Stream() # comms stream + self.fwd_stash = {} # tensor_id => (activation, ev1) + if max_fwd_stash_size < 1: + raise ValueError(f"max_fwd_stash_size should be at least 1 but is {max_fwd_stash_size}") + self.max_fwd_stash_size = max_fwd_stash_size + self.bwd_tensor_stash = {} # tensor_id => activation + self.bwd_ev_stash = {} # tensor_id => ev0 + self.curr_graph_id = None + self.curr_autograd_node = None + + # -------- platform util functions -------- # + def verify_sufficient_virtual_memory(): + curr_pct = get_cpu_ram_pct() + if curr_pct > self.virtual_memory_safe_pct: + warnings.warn(f"{curr_pct=}% > {self.virtual_memory_safe_pct=}% of virtual memory used") + + def get_cpu_ram_pct() -> float: + # get the percentage of memory used by the system + return psutil.virtual_memory().percent + + def get_tensor_id() -> int: + # create a unique id for each tensor we are managing + self.tensor_id += 1 + return self.tensor_id + + def get_num_bytes_tensor(x: torch.Tensor) -> int: + # get the number of bytes in a tensor, for memory management purposes + return x.element_size() * x.nelement() # x.element_size() * x._base_storage().nbytes() + + # -------- core pack / unpack work -------- # + def pack_tensor(activation: torch.Tensor) -> int: + # activations are passed in during forward pass - from here we take over and return a unique id + if self.is_first_forward_call: + if len(self.tracker) != 0: + raise ValueError("Backward pass should have cleared tracker of all tensors") + + # set training phase trackers + self.is_first_forward_call = False + self.is_first_backward_call = True + + # query for basic tensor info + num_bytes = get_num_bytes_tensor(activation) + tensor_id = get_tensor_id() + + # only offload hefty bois if they're activations on CUDA (our heuristic + # for that is to check if they're not params or buffers)! + if ( + activation.device.type in ["cuda", "xpu"] + and num_bytes >= self.min_tensor_size_bytes + and ( + not isinstance(activation, torch.nn.Parameter) + and not (hasattr(torch.nn, "Buffer") and isinstance(activation, torch.nn.Buffer)) + ) + ): + if self.use_streams: + # First, sync back and dereference previously offloaded tensors + # as the offloading should be done sufficiently long ago. + for id in list(self.fwd_stash.keys()): + if id <= tensor_id - self.max_fwd_stash_size: + _, ev = self.fwd_stash[id] + self.s0.wait_event(ev) + del self.fwd_stash[id] + else: + break + + # Sync in, offload, and add an event to sync back later + self.s1.wait_stream(self.s0) + + stream = self.s1 if self.use_streams else self.s0 + with stream if self.accelerator_type == "xpu" else torch.cuda.stream(stream): + cpu_tensor = torch.empty_like(activation, pin_memory=self.use_pin_memory, device="cpu") + cpu_tensor.copy_(activation, non_blocking=True) + self.tracker[tensor_id] = ( + cpu_tensor, + True, # True = (in future) modified + ) + + if self.use_streams: + event = self.s1.record_event() + + # Stash to keep activation alive til s1 is done + self.fwd_stash[tensor_id] = (activation, event) + else: + self.tracker[tensor_id] = ( + activation, + False, + ) # False = not modified, tensor is as is + + return tensor_id + + def unpack_tensor_single_stream(unpack_tensor_id: int) -> torch.Tensor: + # backward pass - we are called with the tensor_id, which + # we will use to retrieve the saved/offloaded tensor + if self.is_first_backward_call: + if self.is_first_forward_pass: + self.is_first_forward_pass = False + if self.use_pin_memory: + verify_sufficient_virtual_memory() + + self.is_first_backward_call = False + self.is_first_forward_call = True + + if unpack_tensor_id not in self.tracker: + raise ValueError(f"Untracked tensor with id {unpack_tensor_id}") + + maybe_accelerator_tensor, modified = self.tracker[unpack_tensor_id] + if modified: + accelerator_tensor = maybe_accelerator_tensor.to(self.accelerator_type, non_blocking=True) + maybe_accelerator_tensor = accelerator_tensor + + # clear tensor from tracking + del self.tracker[unpack_tensor_id] + return maybe_accelerator_tensor + + def unpack_tensor_with_streams(unpack_tensor_id: int) -> torch.Tensor: + # backward pass - we are called with the tensor_id, which + # we will use to retrieve the saved/offloaded tensor + if self.is_first_backward_call: + self.curr_graph_id = torch._C._current_graph_task_id() + + def wait_and_del_remaining_references() -> None: + for id in list(self.bwd_tensor_stash.keys()): + event = self.bwd_ev_stash[id] + self.s1.wait_event(event) + del self.bwd_tensor_stash[id] + + # Register a callback to the end of autograd to clean everything up + torch.autograd.variable.Variable._execution_engine.queue_callback(wait_and_del_remaining_references) + + if self.is_first_forward_pass: + self.is_first_forward_pass = False + if self.use_pin_memory: + verify_sufficient_virtual_memory() + + self.is_first_backward_call = False + self.is_first_forward_call = True + + if unpack_tensor_id not in self.tracker: + raise ValueError(f"untracked tensor with id {unpack_tensor_id}") + + maybe_accelerator_tensor, modified = self.tracker[unpack_tensor_id] + if modified: + # Get data on the current autograd node + graph_id = torch._C._current_graph_task_id() + node = torch._C._current_autograd_node() + prev_node_ids = [] + + # If we're on a new node, mark prev node's tensors to be freed later + if graph_id == self.curr_graph_id and self.curr_autograd_node != node: + self.curr_autograd_node = node + prev_node_ids = list(self.bwd_tensor_stash.keys()) + + brought_back_from_cpu = True + if unpack_tensor_id in self.fwd_stash: + maybe_accelerator_tensor = self.fwd_stash[unpack_tensor_id][0] + brought_back_from_cpu = False + else: + # Kick off the process to bring tensors back + with self.s1 if self.accelerator_type == "xpu" else torch.cuda.stream(self.s1): + accelerator_tensor = maybe_accelerator_tensor.to(self.accelerator_type, non_blocking=True) + maybe_accelerator_tensor = accelerator_tensor + + # Tell comp stream to wait for the info to be loaded before executing + self.s0.wait_stream(self.s1) + + # Stash the tensor to keep memory alive until compute stream is complete + self.bwd_tensor_stash[unpack_tensor_id] = maybe_accelerator_tensor + + # Note: [Track views of the unpacked] + # Why do we get the use count of the unpacked tensor here? We want an + # initial count to compare to later, during the post-hook of the + # backward node, when we need to decide whether we're allowed to free + # the tensor yet. In what obscure cases must we delay freeing the + # tensor (and thus call record_stream)? + # 1. Any of the outputs of the backward node is a view of the unpacked + # tensor. + # 2. In the case that this unpacked tensor will be used in a + # checkpointed region, if one of the recomputed saved tensors ends + # up as a view of the unpacked tensor. + # 3. The user abuses the system somehow and manually relies on the + # unpacked tensor to exist after the backward node has executed. + storage_refcount = torch._C._storage_Use_Count(maybe_accelerator_tensor.untyped_storage()._cdata) + + def hook(outputs, inputs): + # create events for the current node inputs/outputs if they were streamed in + if brought_back_from_cpu: + # See Note: [Track views of the unpacked] + # IF any of the outputs is a view of the tensor, OR if a view of + # the tensor has been saved as a part of checkpoint's recompute + # process, OR the user has abusedly incurred a reference on the + # unpacked tensor, THEN the tensor might be used later and we + # cannot presume to delete it after only the current node is + # done! So we use our frenemy, record_stream, to ensure the + # Tensor stays unmessed with until it's done getting used in the + # compute stream (s0 here). Note that the con here is we introduce + # non-deterministic (thus higher) memory usage, but this case + # should not happen often. + unpacked_tensor = self.bwd_tensor_stash[unpack_tensor_id] + if torch._C._storage_Use_Count(unpacked_tensor.untyped_storage()._cdata) > storage_refcount: + unpacked_tensor.record_stream(self.s0) + del self.bwd_tensor_stash[unpack_tensor_id] + else: + event = self.s0.record_event() + self.bwd_ev_stash[unpack_tensor_id] = event + + # if there are still things in the fwd_stash, get rid of them as we're in bwd now + for id in list(self.fwd_stash.keys()): + _, ev = self.fwd_stash[id] + self.s0.wait_event(ev) + del self.fwd_stash[id] + + # wait on prev node's events and del those + for id in prev_node_ids: + event = self.bwd_ev_stash[id] + self.s1.wait_event(event) + del self.bwd_tensor_stash[id] + + return outputs + + node.register_hook(hook) + + # clear tensor from tracking + del self.tracker[unpack_tensor_id] + return maybe_accelerator_tensor + + unpack_tensor = unpack_tensor_with_streams if self.use_streams else unpack_tensor_single_stream + super().__init__(pack_tensor, unpack_tensor) + + +class NoOpManager(saved_tensors_hooks): + """ + A `saved_tensors_hook` manager used to disable any other `saved_tensors_hook` manager applied before. This relies + on the behavior that only the most recently registered `saved_tensors_hook` will run. + + One example usage is to opt a local region of code out of activations offloading, which is usually applied globally + to best track state. + """ + + def __init__(self) -> None: + def noop(tensor): + return tensor + + super().__init__(noop, noop) + + +def get_act_offloading_ctx_manager( + model: nn.Module, + use_pin_memory: bool = True, + use_streams: bool = True, + min_offload_size: int = 1024, + max_fwd_stash_size: int = 5, + warn_if_no_head: bool = True, +) -> OffloadActivations: + """ + Returns the activation offloading context manager for the model. All but the last output Linear in every step will + be offloaded. + + If activation offloading is enabled, we return the OffloadActivations context manager. If activation offloading is + disabled, we return a NoOpManager context manager. + + Args: + model (`nn.Module`): + Model to wrap with the activation offloading context manager. + use_pin_memory (`bool`, *optional*, defaults to `True`): + Whether to offloaded Tensor will be placed in pinned memory on the CPU. Pinned memory allows the Tensor to + be moved back onto GPU more quickly but is a limited resource. + use_streams (`bool`, *optional*, defaults to `True`): + Whether to use streams for performance optimization where the communications get overlapped with the + computation. Requires a torch build after torch-2.5.0. + min_offload_size (`int`, *optional*, defaults to `1024`): + Minimum number of bytes a Tensor must be in order to qualify for offloading. If the tensor is too small, we + do not want to waste bandwidth and resources moving it to CPU and back. + max_fwd_stash_size (`int`, *optional*, defaults to `5`): + Maximum size of the forward stash, or the maximum number of consecutive activations to keep alive during + the forward pass. This number must be at least 1. Keeping alive more activations will potentially allow + more overlap between the communication and compute streams at the cost of increasing memory usage. Keeping + alive fewer activations will conserve memory, but may cause poor overlap between the streams, increasing + runtime. + warn_if_no_head (`bool`, *optional*, defaults to `True`): + Whether to warn if no output head is detected. If set to `False`, no warning will be raised if no output + head is detected. + + Returns: + `contextlib.ContextDecorator`: + Activation offloading context manager for the model. + """ + activations_handling_ctx = OffloadActivations( + use_pin_memory=use_pin_memory, + use_streams=use_streams, + min_offload_size=min_offload_size, + max_fwd_stash_size=max_fwd_stash_size, + ) + + # Below is our hack to disable offloading the last output Linear in every + # step, as the cost for offloading the activation and then soon after bringing + # it back is expensive. + output_head_detected = False + noop_ctx = NoOpManager() + + # Try to get the actual model if it's wrapped + unwrapped_model = model + if hasattr(unwrapped_model, "module"): + unwrapped_model = unwrapped_model.module + # check for PEFT models + if hasattr(unwrapped_model, "base_model") and hasattr(unwrapped_model, "peft_config"): + unwrapped_model = unwrapped_model.base_model + + # Check for different types of output heads + if hasattr(unwrapped_model, "output"): + if isinstance(unwrapped_model.output, nn.Module): + unwrapped_model.output.register_forward_pre_hook(lambda *args: noop_ctx.__enter__()) + unwrapped_model.output.register_forward_hook(lambda *args: noop_ctx.__exit__(), always_call=True) + output_head_detected = True + elif hasattr(unwrapped_model.output, "linear") and isinstance(unwrapped_model.output.linear, nn.Module): + unwrapped_model.output.linear.register_forward_pre_hook(lambda *args: noop_ctx.__enter__()) + unwrapped_model.output.linear.register_forward_hook(lambda *args: noop_ctx.__exit__(), always_call=True) + output_head_detected = True + + # Check for HuggingFace model output heads + elif hasattr(unwrapped_model, "lm_head"): + unwrapped_model.lm_head.register_forward_pre_hook(lambda *args: noop_ctx.__enter__()) + unwrapped_model.lm_head.register_forward_hook(lambda *args: noop_ctx.__exit__(), always_call=True) + output_head_detected = True + + # Check for decoder-based models + elif hasattr(unwrapped_model, "decoder"): + decoder = unwrapped_model.decoder + if hasattr(decoder, "output"): + decoder.output.register_forward_pre_hook(lambda *args: noop_ctx.__enter__()) + decoder.output.register_forward_hook(lambda *args: noop_ctx.__exit__(), always_call=True) + output_head_detected = True + # Some models have lm_head in the decoder + elif hasattr(decoder, "lm_head"): + decoder.lm_head.register_forward_pre_hook(lambda *args: noop_ctx.__enter__()) + decoder.lm_head.register_forward_hook(lambda *args: noop_ctx.__exit__(), always_call=True) + output_head_detected = True + + # Check for transformer models with final layer norm + elif hasattr(unwrapped_model, "final_layer_norm") or hasattr(unwrapped_model, "ln_f"): + final_norm = getattr(unwrapped_model, "final_layer_norm", None) or unwrapped_model.ln_f + final_norm.register_forward_pre_hook(lambda *args: noop_ctx.__enter__()) + final_norm.register_forward_hook(lambda *args: noop_ctx.__exit__(), always_call=True) + output_head_detected = True + + # Check for models with head module + elif hasattr(unwrapped_model, "head") and isinstance(unwrapped_model.head, nn.Module): + unwrapped_model.head.register_forward_pre_hook(lambda *args: noop_ctx.__enter__()) + unwrapped_model.head.register_forward_hook(lambda *args: noop_ctx.__exit__(), always_call=True) + output_head_detected = True + + if not output_head_detected and warn_if_no_head: + warnings.warn( + "During activation offloading, no output head was detected. If your model has an output head, it will be " + "offloaded. This usually greatly slows training, given the large vocabulary size. To change this " + "behavior, set your output head as model.output and make it an nn.Module. You can disable this warning by " + "passing `warn_if_no_head=False`." + ) + + # Disable offloading for any Liger modules + for name, module in unwrapped_model.named_modules(): + if "liger" in name.lower(): + module.register_forward_pre_hook(lambda *args: noop_ctx.__enter__()) + module.register_forward_hook(lambda *args: noop_ctx.__exit__(), always_call=True) + + return activations_handling_ctx diff --git a/lib/python3.12/site-packages/trl/models/auxiliary_modules.py b/lib/python3.12/site-packages/trl/models/auxiliary_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..fd4f9e71394d5a746a0c456afba03bf9e6fa5d61 --- /dev/null +++ b/lib/python3.12/site-packages/trl/models/auxiliary_modules.py @@ -0,0 +1,96 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +import torch +import torch.nn as nn +import torchvision +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import EntryNotFoundError +from transformers import CLIPModel, is_torch_npu_available, is_torch_xpu_available + + +class MLP(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.Sequential( + nn.Linear(768, 1024), + nn.Dropout(0.2), + nn.Linear(1024, 128), + nn.Dropout(0.2), + nn.Linear(128, 64), + nn.Dropout(0.1), + nn.Linear(64, 16), + nn.Linear(16, 1), + ) + + def forward(self, embed): + return self.layers(embed) + + +class AestheticScorer(torch.nn.Module): + """ + This model attempts to predict the aesthetic score of an image. The aesthetic score is a numerical approximation of + how much a specific image is liked by humans on average. This is from + https://github.com/christophschuhmann/improved-aesthetic-predictor + """ + + def __init__(self, *, dtype, model_id, model_filename): + super().__init__() + self.clip = CLIPModel.from_pretrained("openai/clip-vit-large-patch14") + self.normalize = torchvision.transforms.Normalize( + mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711] + ) + self.target_size = 224 + self.mlp = MLP() + try: + cached_path = hf_hub_download(model_id, model_filename) + except EntryNotFoundError: + cached_path = os.path.join(model_id, model_filename) + state_dict = torch.load(cached_path, map_location=torch.device("cpu"), weights_only=True) + self.mlp.load_state_dict(state_dict) + self.dtype = dtype + self.eval() + + def __call__(self, images): + device = next(self.parameters()).device + images = torchvision.transforms.Resize(self.target_size)(images) + images = self.normalize(images).to(self.dtype).to(device) + embed = self.clip.get_image_features(pixel_values=images) + # normalize embedding + embed = embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True) + reward = self.mlp(embed).squeeze(1) + return reward + + +def aesthetic_scorer(hub_model_id, model_filename): + scorer = AestheticScorer( + model_id=hub_model_id, + model_filename=model_filename, + dtype=torch.float32, + ) + if is_torch_npu_available(): + scorer = scorer.npu() + elif is_torch_xpu_available(): + scorer = scorer.xpu() + else: + scorer = scorer.cuda() + + def _fn(images, prompts, metadata): + images = (images).clamp(0, 1) + scores = scorer(images) + return scores, {} + + return _fn diff --git a/lib/python3.12/site-packages/trl/models/modeling_base.py b/lib/python3.12/site-packages/trl/models/modeling_base.py new file mode 100644 index 0000000000000000000000000000000000000000..e2036f496b6f65bdaf2c926bc15edf67e7129161 --- /dev/null +++ b/lib/python3.12/site-packages/trl/models/modeling_base.py @@ -0,0 +1,719 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import logging +import os +from copy import deepcopy +from typing import Optional + +import torch +import torch.nn as nn +from accelerate import PartialState +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import ( + EntryNotFoundError, + HFValidationError, + LocalEntryNotFoundError, + RepositoryNotFoundError, +) +from safetensors.torch import load_file as safe_load_file +from transformers import GenerationMixin, PreTrainedModel, is_torch_npu_available, is_torch_xpu_available +from transformers.utils import is_peft_available + + +if is_peft_available(): + from peft import ( + PeftConfig, + PeftModel, + PeftModelForCausalLM, + PeftModelForSeq2SeqLM, + PromptLearningConfig, + get_peft_model, + prepare_model_for_kbit_training, + ) + + +from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled + + +LAYER_PATTERNS = [ + "transformer.h.{layer}", + "model.decoder.layers.{layer}", + "gpt_neox.layers.{layer}", + "model.layers.{layer}", +] + + +class PreTrainedModelWrapper(nn.Module): + r""" + A wrapper class around a (`transformers.PreTrainedModel`) to be compatible with the (`~transformers.PreTrained`) + class in order to keep some attributes and methods of the (`~transformers.PreTrainedModel`) class. + + Attributes: + pretrained_model (`transformers.PreTrainedModel`): + The model to be wrapped. + parent_class (`transformers.PreTrainedModel`): + The parent class of the model to be wrapped. + supported_args (`list`): + The list of arguments that are supported by the wrapper class. + """ + + transformers_parent_class = None + supported_args = None + supported_modules = ("v_head",) + supported_rm_modules = ("score",) + supported_pretrained_model_architectures = ( + (PreTrainedModel) + if not is_peft_available() + else (PreTrainedModel, PeftModelForCausalLM, PeftModelForSeq2SeqLM) + ) + + def __init__( + self, pretrained_model=None, score_module=None, supports_rm_adapter=False, rm_adapter_name=None, **kwargs + ): + super().__init__() + self.pretrained_model = pretrained_model + + self.config = pretrained_model.config + self.prepare_inputs_for_generation = pretrained_model.prepare_inputs_for_generation + self.is_loaded_in_8bit = getattr(pretrained_model, "is_loaded_in_8bit", False) + self.is_loaded_in_4bit = getattr(pretrained_model, "is_loaded_in_4bit", False) + self.is_sequential_parallel = False + + if hasattr(pretrained_model, "gradient_checkpointing_disable"): + self.gradient_checkpointing_disable = pretrained_model.gradient_checkpointing_disable + + if hasattr(pretrained_model, "gradient_checkpointing_enable"): + self.gradient_checkpointing_enable = pretrained_model.gradient_checkpointing_enable + + if hasattr(pretrained_model, "enable_input_require_grads"): + self.enable_input_require_grads = pretrained_model.enable_input_require_grads + + self.supports_rm_adapter = supports_rm_adapter + self.rm_adapter_name = rm_adapter_name + self.policy_adapter_name = "default" + if score_module is not None: + self.score = score_module + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + r""" + Instantiates a new model from a pretrained model from `transformers`. The pretrained model is loaded using the + `from_pretrained` method of the `transformers.PreTrainedModel` class. The arguments that are specific to the + `transformers.PreTrainedModel` class are passed along this method and filtered out from the `kwargs` argument. + + Args: + pretrained_model_name_or_path (`str` or `transformers.PreTrainedModel`): + The path to the pretrained model or its name. + *model_args (`list`, *optional*)): + Additional positional arguments passed along to the underlying model's `from_pretrained` method. + **kwargs (`dict`, *optional*): + Additional keyword arguments passed along to the underlying model's `from_pretrained` method. We also + pre-process the kwargs to extract the arguments that are specific to the `transformers.PreTrainedModel` + class and the arguments that are specific to trl models. The kwargs also support + `prepare_model_for_kbit_training` arguments from `peft` library. + """ + if kwargs is not None: + peft_config = kwargs.pop("peft_config", None) + reward_adapter = kwargs.pop("reward_adapter", None) + reward_adapter_name = kwargs.pop("reward_adapter_name", "reward_adapter") + is_trainable = kwargs.pop("is_trainable", False) + trl_model_args, pretrained_kwargs, peft_quantization_kwargs = cls._split_kwargs(kwargs) + token = pretrained_kwargs.get("token", None) + else: + peft_config = None + is_trainable = False + trl_model_args = {} + pretrained_kwargs = {} + peft_quantization_kwargs = {} + token = None + + if reward_adapter is not None and not isinstance(reward_adapter, str): + raise ValueError( + "The `reward_adapter` argument should be a string representing the name of local path or the Hub id to the Reward Modeling adapter." + ) + + is_peft_model = False + + current_device = cls._get_current_device() + if isinstance(pretrained_model_name_or_path, str): + is_loaded_in_8bit = pretrained_kwargs["load_in_8bit"] if "load_in_8bit" in pretrained_kwargs else False + is_loaded_in_4bit = pretrained_kwargs["load_in_4bit"] if "load_in_4bit" in pretrained_kwargs else False + else: + is_loaded_in_8bit = getattr(pretrained_model_name_or_path, "is_loaded_in_8bit", False) + is_loaded_in_4bit = getattr(pretrained_model_name_or_path, "is_loaded_in_4bit", False) + + if (is_loaded_in_8bit or is_loaded_in_4bit) and "device_map" not in pretrained_kwargs: + # warn users + logging.warning( + "The `device_map` argument is not provided. We will override the device_map argument." + " to set the entire" + " model on the current device. If you want to set the model on multiple devices, please provide" + " a custom `device_map` argument." + ) + pretrained_kwargs["device_map"] = {"": current_device} + + if is_peft_available() and peft_config is not None and not isinstance(peft_config, PeftConfig): + raise ValueError("The `peft_config` argument should be an instance of `peft.PeftConfig` class.") + + # First, load the pre-trained model using the parent-class + # either `AutoModelForCausalLM` or `AutoModelForSeq2SeqLM` + if isinstance(pretrained_model_name_or_path, str): + if is_peft_available(): + try: + # If there is a trained peft adapter in the hub, load its config. + remote_adapter_config = hf_hub_download( + pretrained_model_name_or_path, + "adapter_config.json", + token=token, + ) + except (EntryNotFoundError, LocalEntryNotFoundError, HFValidationError, RepositoryNotFoundError): + remote_adapter_config = None + else: + remote_adapter_config = None + + local_adapter_present = os.path.exists(os.path.join(pretrained_model_name_or_path, "adapter_config.json")) + + if (local_adapter_present or remote_adapter_config is not None) and is_peft_available(): + if peft_config is not None: + logging.warning( + "`peft_config` argument ignored since a peft config file was found in " + f"{pretrained_model_name_or_path}" + ) + + # Load the trained peft adapter config + if local_adapter_present: + trained_adapter_config = PeftConfig.from_pretrained(pretrained_model_name_or_path) + else: + remote_adapter_dir = os.path.dirname(remote_adapter_config) + trained_adapter_config = PeftConfig.from_pretrained(remote_adapter_dir) + + # Load the pretrained base model + pretrained_model = cls.transformers_parent_class.from_pretrained( + trained_adapter_config.base_model_name_or_path, *model_args, **pretrained_kwargs + ) + + # Wrap the pretrained model with the trained peft adapter + pretrained_model = PeftModel.from_pretrained( + pretrained_model, pretrained_model_name_or_path, is_trainable=is_trainable, token=token + ) + logging.info("Trained peft adapter loaded") + else: + pretrained_model = cls.transformers_parent_class.from_pretrained( + pretrained_model_name_or_path, *model_args, **pretrained_kwargs + ) + + if peft_config is not None: + # Initialize a new peft adapter with the given config + if is_loaded_in_8bit or is_loaded_in_4bit: + pretrained_model = prepare_model_for_kbit_training( + pretrained_model, + **peft_quantization_kwargs, + ) + pretrained_model = get_peft_model(pretrained_model, peft_config) + logging.info("peft adapter initialised") + + elif isinstance(pretrained_model_name_or_path, cls.supported_pretrained_model_architectures): + pretrained_model = pretrained_model_name_or_path + + if peft_config is not None and isinstance(pretrained_model, PreTrainedModel): + # Initialize a new peft adapter with the given config + if is_loaded_in_8bit or is_loaded_in_4bit: + pretrained_model = prepare_model_for_kbit_training( + pretrained_model, + **peft_quantization_kwargs, + ) + pretrained_model = get_peft_model(pretrained_model, peft_config) + logging.info("peft adapter initialised") + else: + raise ValueError( + "pretrained_model_name_or_path should be a string or a PreTrainedModel, " + f"but is {type(pretrained_model_name_or_path)}" + ) + + if is_peft_available(): + if isinstance(pretrained_model, PeftModel): + is_peft_model = True + # for backward compatibility + if hasattr(pretrained_model, "active_peft_config") and isinstance( + pretrained_model.active_peft_config, PromptLearningConfig + ): + raise ValueError("PromptLearningConfig is not supported for PPO training.") + + # Add reward modeling adapter if specified + if not is_peft_model and reward_adapter is not None: + raise ValueError("reward_adapter can only be used with a PeftModel. ") + elif is_peft_model and reward_adapter is not None: + score_module = cls.add_and_load_reward_modeling_adapter( + pretrained_model, reward_adapter, reward_adapter_name, token=token + ) + multi_adapter_args = { + "score_module": score_module, + "supports_rm_adapter": True, + "rm_adapter_name": reward_adapter_name, + } + else: + multi_adapter_args = {"supports_rm_adapter": False} + + # Then, create the full model by instantiating the wrapper class + model = cls(pretrained_model, **multi_adapter_args, **trl_model_args) + + # if resume_training, load the state_dict again - this is ok since the + # state_dict is removed from the model after loading it. + is_resuming_training = True + if isinstance(pretrained_model_name_or_path, str): + safe_filename = os.path.join(pretrained_model_name_or_path, "model.safetensors") + filename = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin") + + sharded_index_filename = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin.index.json") + safe_sharded_index_filename = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + is_sharded = False + use_safe = os.path.exists(safe_filename) + + if not (os.path.exists(filename) or os.path.exists(safe_filename)): + # Try with `pytorch_model.bin` + filename, files_to_download, is_sharded, is_resuming_training = cls._get_checkpoint_from_hub( + pretrained_model, + pretrained_model_name_or_path, + sharded_index_filename, + token=token, + ) + # Try with safetensors + if filename is None and files_to_download is None: + safe_filename, files_to_download, is_sharded, is_resuming_training = cls._get_checkpoint_from_hub( + pretrained_model, + pretrained_model_name_or_path, + safe_sharded_index_filename, + token=token, + model_name="model.safetensors", + model_index_name="model.safetensors.index.json", + ) + use_safe = True + else: + use_safe = False + + loading_func = safe_load_file if use_safe else torch.load + load_kwargs = {} if use_safe else {"map_location": "cpu", "weights_only": True} + + if is_resuming_training: + if is_sharded: + # download each file and add it to the state_dict + state_dict = {} + + for shard_file in files_to_download: + filename = hf_hub_download( + pretrained_model_name_or_path, + shard_file, + token=token, + ) + state_dict.update(loading_func(filename, **load_kwargs)) + else: + state_dict = loading_func(filename if not use_safe else safe_filename, **load_kwargs) + + else: + state_dict = pretrained_model_name_or_path.state_dict() + + model.is_peft_model = is_peft_model + model.current_device = current_device + + if is_resuming_training: + model.post_init(state_dict=state_dict) + + return model + + @classmethod + def _get_checkpoint_from_hub( + cls, + pretrained_model, + pretrained_model_name_or_path, + index_filename, + token=None, + model_name="pytorch_model.bin", + model_index_name="pytorch_model.bin.index.json", + ): + files_to_download = None + filename = None + is_resuming_training = True + is_sharded = False + + try: + filename = hf_hub_download( + pretrained_model_name_or_path, + model_name, + token=token, + ) + # sharded + except (EntryNotFoundError, LocalEntryNotFoundError, HFValidationError, RepositoryNotFoundError): + if os.path.exists(index_filename): + index_file_name = index_filename + else: + try: + index_file_name = hf_hub_download( + pretrained_model_name_or_path, + model_index_name, + token=token, + ) + except (EntryNotFoundError, LocalEntryNotFoundError, HFValidationError, RepositoryNotFoundError): + # not continue training, do not have v_head weight + is_resuming_training = False + logging.warning( + f"A {type(pretrained_model)} model is loaded from '{pretrained_model_name_or_path}', " + f"and no v_head weight is found. This IS expected if you are not resuming PPO training." + ) + # load json + if is_resuming_training: + with open(index_file_name) as f: + index = json.load(f) + # check filename with `v_head` or any known extra module: + files_to_download = set() + for k, v in index["weight_map"].items(): + if any(module in k for module in cls.supported_modules): + files_to_download.add(v) + is_sharded = True + + return filename, files_to_download, is_sharded, is_resuming_training + + @classmethod + def _get_current_device(cls): + r""" + Get the current device. For GPU & XPU, we return the local process index using the `accelerate.PartialState` + object to handle corner cases when running scripts in distributed environments. + + Returns: + current_device (`Union[int, str]`): + The current device. + """ + state = PartialState() + if torch.cuda.is_available() or is_torch_xpu_available(): + return state.local_process_index + elif is_torch_npu_available(): + return f"npu:{state.local_process_index}" + else: + return "cpu" + + @classmethod + def _split_kwargs(cls, kwargs): + """ + Separate the kwargs from the arguments that we support inside `supported_args` and the ones that we don't. + """ + check_peft_kwargs = False + + if is_peft_available(): + from peft import prepare_model_for_kbit_training + + check_peft_kwargs = True + + supported_kwargs = {} + unsupported_kwargs = {} + peft_kwargs = {} + + for key, value in kwargs.items(): + if key in cls.supported_args: + supported_kwargs[key] = value + else: + unsupported_kwargs[key] = value + + if check_peft_kwargs: + if key in prepare_model_for_kbit_training.__code__.co_varnames: + peft_kwargs[key] = value + if key in unsupported_kwargs: + unsupported_kwargs.pop(key) + + return supported_kwargs, unsupported_kwargs, peft_kwargs + + @classmethod + def add_and_load_reward_modeling_adapter( + cls, pretrained_model, adapter_model_id, adapter_name="reward_model_adapter", token=None + ): + r""" + Add and load a reward modeling adapter. This method can only be used if the model is a `PeftModel` and if you + have initialized the model with the `reward_modeling_adapter_id` argument, pointing to the id of the reward + modeling adapter. The latest needs also to contain the score head in order to produce the reward. + """ + pretrained_model.load_adapter(adapter_model_id, adapter_name, is_trainable=False) + pretrained_model.train() + + filename = os.path.join(adapter_model_id, "adapter_model.bin") + safe_loading = False + if not os.path.exists(filename): + try: + local_filename = hf_hub_download( + adapter_model_id, + "adapter_model.bin", + token=token, + ) + except Exception: + filename = os.path.join(adapter_model_id, "adapter_model.safetensors") + safe_loading = True + if not os.path.exists(filename): + try: + local_filename = hf_hub_download( + adapter_model_id, + "adapter_model.safetensors", + token=token, + ) + except Exception as exc: + raise ValueError( + "Could not find adapter model in the Hub, make sure you have the correct adapter model id." + ) from exc + else: + local_filename = filename + else: + local_filename = filename + + loading_func = safe_load_file if safe_loading else torch.load + load_kwargs = {} if safe_loading else {"map_location": "cpu", "weights_only": True} + + adapter_state_dict = loading_func(local_filename, **load_kwargs) + + for score_name_candidate in cls.supported_rm_modules: + if any(score_name_candidate in name for name in adapter_state_dict.keys()): + score_name = score_name_candidate + # we have found the correct head name and can break + break + + score_dict = {} + + for name, param in adapter_state_dict.items(): + if score_name in name: + key_name = ".".join(name.split(".")[-1:]) + score_dict[key_name] = param.to(cls._get_current_device()) + + num_labels, hidden_dim = score_dict["weight"].shape + has_bias = any("bias" in name for name in adapter_state_dict.keys()) + + score = nn.Linear(hidden_dim, num_labels, bias=has_bias).to( + device=cls._get_current_device(), + dtype=pretrained_model.dtype, + ) + score.load_state_dict(score_dict) + for param in score.parameters(): + param.requires_grad = False + + return score + + def push_to_hub(self, *args, **kwargs): + r""" + Push the pretrained model to the hub. This method is a wrapper around + `transformers.PreTrainedModel.push_to_hub`. Please refer to the documentation of + `transformers.PreTrainedModel.push_to_hub` for more information. + + Args: + *args (`list`, *optional*): + Positional arguments passed along to the underlying model's `push_to_hub` method. + **kwargs (`dict`, *optional*): + Keyword arguments passed along to the underlying model's `push_to_hub` method. + """ + raise NotImplementedError + + def save_pretrained(self, *args, **kwargs): + r""" + Save the pretrained model to a directory. This method is a wrapper around + `transformers.PreTrainedModel.save_pretrained`. Please refer to the documentation of + `transformers.PreTrainedModel.save_pretrained` for more information. + + Args: + *args (`list`, *optional*): + Positional arguments passed along to the underlying model's `save_pretrained` method. + **kwargs (`dict`, *optional*): + Keyword arguments passed along to the underlying model's `save_pretrained` method. + """ + state_dict = kwargs.get("state_dict") + if state_dict is None: + state_dict = self.state_dict() + kwargs["state_dict"] = state_dict + + # if it is a peft model only save the `v_head` state_dict and + # pop the `state_dict` from the kwargs to avoid slient bugs with `peft` + if self.is_peft_model: + save_path = args[0] + save_path = os.path.join(save_path, "pytorch_model.bin") + torch.save(state_dict, save_path) + _ = kwargs.pop("state_dict", None) + + return self.pretrained_model.save_pretrained(*args, **kwargs) + + def state_dict(self, *args, **kwargs): + r""" + Return the state_dict of the pretrained model. + """ + raise NotImplementedError + + def post_init(self, *args, **kwargs): + r""" + Post initialization method. This method is called after the model is instantiated and loaded from a checkpoint. + It can be used to perform additional operations such as loading the state_dict. + """ + raise NotImplementedError + + def compute_reward_score(self, input_ids, attention_mask=None, **kwargs): + r""" + Computes the reward score for a given input. The method has first to enable the adapter and then compute the + reward score. After that the model disables the reward modeling adapter and enables the default ppo adapter + again. + """ + if not self.supports_rm_adapter: + raise ValueError("This model does not support reward modeling adapter.") + + # enable rm adapter + self.pretrained_model.set_adapter(self.rm_adapter_name) + self.pretrained_model.eval() + + with torch.no_grad(): + base_model_output = self.pretrained_model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + return_dict=True, + **kwargs, + ) + + last_hidden_states = base_model_output.hidden_states[-1] + scores = self.score(last_hidden_states) + + self.pretrained_model.set_adapter(self.policy_adapter_name) + self.pretrained_model.eval() + + return scores + + +def create_reference_model( + model: PreTrainedModelWrapper, num_shared_layers: Optional[int] = None, pattern: Optional[str] = None +) -> PreTrainedModelWrapper: + """ + Creates a static reference copy of a model. Note that model will be in `.eval()` mode. + + Args: + model (`PreTrainedModelWrapper`): The model to be copied. + num_shared_layers (`int`, *optional*): + The number of initial layers that are shared between both models and kept frozen. + pattern (`str`, *optional*): The shared layers are selected with a string pattern + (e.g. "transformer.h.{layer}" for GPT2) and if a custom pattern is necessary it can be passed here. + + Returns: + `PreTrainedModelWrapper` + """ + if is_deepspeed_zero3_enabled(): + raise ValueError( + "DeepSpeed ZeRO-3 is enabled and is not compatible with `create_reference_model()`. Please instantiate your reference model directly with `AutoModelForCausalLM.from_pretrained()`." + ) + + parameter_names = [n for n, _ in model.named_parameters()] + ref_model = deepcopy(model) + + # if no layers are shared, return copy of model + if num_shared_layers is None: + for param_name in parameter_names: + param = ref_model.get_parameter(param_name) + param.requires_grad = False + return ref_model.eval() + + # identify layer name pattern + if pattern is not None: + pattern = pattern.format(layer=num_shared_layers) + else: + for pattern_candidate in LAYER_PATTERNS: + pattern_candidate = pattern_candidate.format(layer=num_shared_layers) + if any(pattern_candidate in name for name in parameter_names): + pattern = pattern_candidate + break + + if pattern is None: + raise ValueError("Layer pattern could not be matched.") + + # divide parameters in shared and unshared parameter lists + shared_param_list = [] + unshared_param_list = [] + + shared_parameter = True + for name, _param in model.named_parameters(): + if pattern in name: + shared_parameter = False + if shared_parameter: + shared_param_list.append(name) + else: + unshared_param_list.append(name) + + # create reference of the original parameter if they are shared + for param_name in shared_param_list: + param = model.get_parameter(param_name) + param.requires_grad = False + + _ref_param = ref_model.get_parameter(param_name) + + # for all other parameters just make sure they don't use gradients + for param_name in unshared_param_list: + param = ref_model.get_parameter(param_name) + param.requires_grad = False + + if pattern is not None and len(unshared_param_list) == 0: + logging.warning("Pattern passed or found, but no layers matched in the model. Check for a typo.") + + return ref_model.eval() + + +class GeometricMixtureWrapper(GenerationMixin): + r""" + Geometric Mixture generation wrapper that samples from the logits of two model's geometric mixture. + + Args: + model (`PreTrainedModel`): The model to be wrapped. + ref_model (`PreTrainedModel`): The reference model. + generation_config (`GenerationConfig`): The generation config. + mixture_coef (`float`, *optional* - default: 0.5): The mixture coefficient. + """ + + main_input_name = "input_ids" + _supports_cache_class = False + _supports_static_cache = False + + def __init__(self, model, ref_model, generation_config, mixture_coef=0.5, device=None): + super().__init__() + + self.model = model + self.config = model.config + self.ref_model = ref_model + self.generation_config = generation_config + self.mixture_coef = mixture_coef + self.device = device + + def __call__(self, *args, **kwargs): + return self.forward(*args, **kwargs) + + @torch.inference_mode() + def forward(self, *args, **kwargs): + model_outputs = self.model(*args, **kwargs) + model_logits = model_outputs.logits + ref_model_logits = self.ref_model(*args, **kwargs).logits + + model_outputs.logits = torch.nn.functional.log_softmax( + self.mixture_coef * ref_model_logits + (1 - self.mixture_coef) * model_logits, dim=-1 + ) + + return model_outputs + + def prepare_inputs_for_generation(self, *args, **kwargs): + # turn off cache in the generation config + kwargs["use_cache"] = False + model_inputs = self.model.prepare_inputs_for_generation(*args, **kwargs) + _ = self.ref_model.prepare_inputs_for_generation(*args, **kwargs) + + return model_inputs + + def _validate_model_class(self): + self.model._validate_model_class() + + def _validate_model_kwargs(self, model_kwargs): + return self.model._validate_model_kwargs(model_kwargs) diff --git a/lib/python3.12/site-packages/trl/models/modeling_sd_base.py b/lib/python3.12/site-packages/trl/models/modeling_sd_base.py new file mode 100644 index 0000000000000000000000000000000000000000..8e61e4add246687d331188ed29f557dd6f41a6ee --- /dev/null +++ b/lib/python3.12/site-packages/trl/models/modeling_sd_base.py @@ -0,0 +1,918 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextlib +import os +import random +import warnings +from dataclasses import dataclass +from typing import Any, Callable, Optional, Union + +import numpy as np +import torch +import torch.utils.checkpoint as checkpoint +from diffusers import DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg +from transformers.utils import is_peft_available + +from ..core import randn_tensor +from .sd_utils import convert_state_dict_to_diffusers + + +if is_peft_available(): + from peft import LoraConfig + from peft.utils import get_peft_model_state_dict + + +@dataclass +class DDPOPipelineOutput: + """ + Output class for the diffusers pipeline to be finetuned with the DDPO trainer + + Args: + images (`torch.Tensor`): + The generated images. + latents (`list[torch.Tensor]`): + The latents used to generate the images. + log_probs (`list[torch.Tensor]`): + The log probabilities of the latents. + + """ + + images: torch.Tensor + latents: torch.Tensor + log_probs: torch.Tensor + + +@dataclass +class DDPOSchedulerOutput: + """ + Output class for the diffusers scheduler to be finetuned with the DDPO trainer + + Args: + latents (`torch.Tensor`): + Predicted sample at the previous timestep. Shape: `(batch_size, num_channels, height, width)` + log_probs (`torch.Tensor`): + Log probability of the above mentioned sample. Shape: `(batch_size)` + """ + + latents: torch.Tensor + log_probs: torch.Tensor + + +class DDPOStableDiffusionPipeline: + """ + Main class for the diffusers pipeline to be finetuned with the DDPO trainer + """ + + def __call__(self, *args, **kwargs) -> DDPOPipelineOutput: + raise NotImplementedError + + def scheduler_step(self, *args, **kwargs) -> DDPOSchedulerOutput: + raise NotImplementedError + + @property + def unet(self): + """ + Returns the 2d U-Net model used for diffusion. + """ + raise NotImplementedError + + @property + def vae(self): + """ + Returns the Variational Autoencoder model used from mapping images to and from the latent space + """ + raise NotImplementedError + + @property + def tokenizer(self): + """ + Returns the tokenizer used for tokenizing text inputs + """ + raise NotImplementedError + + @property + def scheduler(self): + """ + Returns the scheduler associated with the pipeline used for the diffusion process + """ + raise NotImplementedError + + @property + def text_encoder(self): + """ + Returns the text encoder used for encoding text inputs + """ + raise NotImplementedError + + @property + def autocast(self): + """ + Returns the autocast context manager + """ + raise NotImplementedError + + def set_progress_bar_config(self, *args, **kwargs): + """ + Sets the progress bar config for the pipeline + """ + raise NotImplementedError + + def save_pretrained(self, *args, **kwargs): + """ + Saves all of the model weights + """ + raise NotImplementedError + + def get_trainable_layers(self, *args, **kwargs): + """ + Returns the trainable parameters of the pipeline + """ + raise NotImplementedError + + def save_checkpoint(self, *args, **kwargs): + """ + Light wrapper around accelerate's register_save_state_pre_hook which is run before saving state + """ + raise NotImplementedError + + def load_checkpoint(self, *args, **kwargs): + """ + Light wrapper around accelerate's register_lad_state_pre_hook which is run before loading state + """ + raise NotImplementedError + + +def _left_broadcast(input_tensor, shape): + """ + As opposed to the default direction of broadcasting (right to left), this function broadcasts from left to right + + Args: + input_tensor (`torch.FloatTensor`): is the tensor to broadcast + shape (`tuple[int]`): is the shape to broadcast to + """ + input_ndim = input_tensor.ndim + if input_ndim > len(shape): + raise ValueError( + "The number of dimensions of the tensor to broadcast cannot be greater than the length of the shape to broadcast to" + ) + return input_tensor.reshape(input_tensor.shape + (1,) * (len(shape) - input_ndim)).broadcast_to(shape) + + +def _get_variance(self, timestep, prev_timestep): + alpha_prod_t = torch.gather(self.alphas_cumprod, 0, timestep.cpu()).to(timestep.device) + alpha_prod_t_prev = torch.where( + prev_timestep.cpu() >= 0, + self.alphas_cumprod.gather(0, prev_timestep.cpu()), + self.final_alpha_cumprod, + ).to(timestep.device) + beta_prod_t = 1 - alpha_prod_t + beta_prod_t_prev = 1 - alpha_prod_t_prev + + variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) + + return variance + + +def scheduler_step( + self, + model_output: torch.FloatTensor, + timestep: int, + sample: torch.FloatTensor, + eta: float = 0.0, + use_clipped_model_output: bool = False, + generator=None, + prev_sample: Optional[torch.FloatTensor] = None, +) -> DDPOSchedulerOutput: + """ + + Predict the sample at the previous timestep by reversing the SDE. Core function to propagate the diffusion process + from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): direct output from learned diffusion model. + timestep (`int`): current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + current instance of sample being created by diffusion process. + eta (`float`): weight of noise for added noise in diffusion step. + use_clipped_model_output (`bool`): if `True`, compute "corrected" `model_output` from the clipped + predicted original sample. Necessary because predicted original sample is clipped to [-1, 1] when + `self.config.clip_sample` is `True`. If no clipping has happened, "corrected" `model_output` would coincide + with the one provided as input and `use_clipped_model_output` will have not effect. + generator: random number generator. + variance_noise (`torch.FloatTensor`): instead of generating noise for the variance using `generator`, we + can directly provide the noise for the variance itself. This is useful for methods such as CycleDiffusion. + (https://huggingface.co/papers/2210.05559) + + Returns: + `DDPOSchedulerOutput`: the predicted sample at the previous timestep and the log probability of the sample + """ + + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + # See formulas (12) and (16) of DDIM paper https://huggingface.co/papers/2010.02502 + # Ideally, read DDIM paper in-detail understanding + + # Notation ( -> + # - pred_noise_t -> e_theta(x_t, t) + # - pred_original_sample -> f_theta(x_t, t) or x_0 + # - std_dev_t -> sigma_t + # - eta -> η + # - pred_sample_direction -> "direction pointing to x_t" + # - pred_prev_sample -> "x_t-1" + # 1. get previous step value (=t-1) + prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps + # to prevent OOB on gather + prev_timestep = torch.clamp(prev_timestep, 0, self.config.num_train_timesteps - 1) + + # 2. compute alphas, betas + alpha_prod_t = self.alphas_cumprod.gather(0, timestep.cpu()) + alpha_prod_t_prev = torch.where( + prev_timestep.cpu() >= 0, + self.alphas_cumprod.gather(0, prev_timestep.cpu()), + self.final_alpha_cumprod, + ) + alpha_prod_t = _left_broadcast(alpha_prod_t, sample.shape).to(sample.device) + alpha_prod_t_prev = _left_broadcast(alpha_prod_t_prev, sample.shape).to(sample.device) + + beta_prod_t = 1 - alpha_prod_t + + # 3. compute predicted original sample from predicted noise also called + # "predicted x_0" of formula (12) from https://huggingface.co/papers/2010.02502 + if self.config.prediction_type == "epsilon": + pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) + pred_epsilon = model_output + elif self.config.prediction_type == "sample": + pred_original_sample = model_output + pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) + elif self.config.prediction_type == "v_prediction": + pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output + pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" + " `v_prediction`" + ) + + # 4. Clip or threshold "predicted x_0" + if self.config.thresholding: + pred_original_sample = self._threshold_sample(pred_original_sample) + elif self.config.clip_sample: + pred_original_sample = pred_original_sample.clamp( + -self.config.clip_sample_range, self.config.clip_sample_range + ) + + # 5. compute variance: "sigma_t(η)" -> see formula (16) + # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) + variance = _get_variance(self, timestep, prev_timestep) + std_dev_t = eta * variance ** (0.5) + std_dev_t = _left_broadcast(std_dev_t, sample.shape).to(sample.device) + + if use_clipped_model_output: + # the pred_epsilon is always re-derived from the clipped x_0 in Glide + pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) + + # 6. compute "direction pointing to x_t" of formula (12) from https://huggingface.co/papers/2010.02502 + pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon + + # 7. compute x_t without "random noise" of formula (12) from https://huggingface.co/papers/2010.02502 + prev_sample_mean = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction + + if prev_sample is not None and generator is not None: + raise ValueError( + "Cannot pass both generator and prev_sample. Please make sure that either `generator` or" + " `prev_sample` stays `None`." + ) + + if prev_sample is None: + variance_noise = randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=model_output.dtype, + ) + prev_sample = prev_sample_mean + std_dev_t * variance_noise + + # log prob of prev_sample given prev_sample_mean and std_dev_t + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * (std_dev_t**2)) + - torch.log(std_dev_t) + - torch.log(torch.sqrt(2 * torch.as_tensor(np.pi))) + ) + # mean along all but batch dimension + log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim))) + + return DDPOSchedulerOutput(prev_sample.type(sample.dtype), log_prob) + + +# 1. The output type for call is different as the logprobs are now returned +# 2. An extra method called `scheduler_step` is added which is used to constraint the scheduler output +@torch.no_grad() +def pipeline_step( + self, + prompt: Optional[Union[str, list[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + negative_prompt: Optional[Union[str, list[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: int = 1, + cross_attention_kwargs: Optional[dict[str, Any]] = None, + guidance_rescale: float = 0.0, +): + r""" + Function invoked when calling the pipeline for generation. Args: prompt (`str` or `list[str]`, *optional*): The + prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height + (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): The height in pixels of the + generated image. + width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to 50): + 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*, defaults to 7.5): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. of + [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the + text `prompt`, usually at the expense of lower image quality. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies + to [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to + make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between [PIL](https://pillow.readthedocs.io/en/stable/): + `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple. + callback (`Callable`, *optional*): + A function that will be called every `callback_steps` steps during inference. The function will be called + with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function will be called. If not specified, the callback will be + called at every step. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py). + guidance_rescale (`float`, *optional*, defaults to 0.7): + Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of + [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). + Guidance rescale factor should fix overexposure when using zero terminal SNR. + + Examples: + + Returns: + `DDPOPipelineOutput`: The generated image, the predicted latents used to generate the image and the associated + log probabilities + """ + # 0. Default height and width to unet + height = height or self.unet.config.sample_size * self.vae_scale_factor + width = width or self.unet.config.sample_size * self.vae_scale_factor + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + height, + width, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + text_encoder_lora_scale = cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None + prompt_embeds = self._encode_prompt( + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=text_encoder_lora_scale, + ) + + # 4. Prepare timesteps + self.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + # 5. Prepare latent variables + num_channels_latents = self.unet.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + + # 6. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + all_latents = [latents] + all_log_probs = [] + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = self.scheduler.scale_model_input(latent_model_input, t) + + # predict the noise residual + noise_pred = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + if do_classifier_free_guidance and guidance_rescale > 0.0: + # Based on 3.4. in https://huggingface.co/papers/2305.08891 + noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) + + # compute the previous noisy sample x_t -> x_t-1 + scheduler_output = scheduler_step(self.scheduler, noise_pred, t, latents, eta) + latents = scheduler_output.latents + log_prob = scheduler_output.log_probs + + all_latents.append(latents) + all_log_probs.append(log_prob) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + callback(i, t, latents) + + if not output_type == "latent": + image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] + image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) + else: + image = latents + has_nsfw_concept = None + + if has_nsfw_concept is None: + do_denormalize = [True] * image.shape[0] + else: + do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] + + image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + + # Offload last model to CPU + if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: + self.final_offload_hook.offload() + + return DDPOPipelineOutput(image, all_latents, all_log_probs) + + +def pipeline_step_with_grad( + pipeline, + prompt: Optional[Union[str, list[str]]] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: int = 50, + guidance_scale: float = 7.5, + truncated_backprop: bool = True, + truncated_backprop_rand: bool = True, + gradient_checkpoint: bool = True, + truncated_backprop_timestep: int = 49, + truncated_rand_backprop_minmax: tuple = (0, 50), + negative_prompt: Optional[Union[str, list[str]]] = None, + num_images_per_prompt: Optional[int] = 1, + eta: float = 0.0, + generator: Optional[Union[torch.Generator, list[torch.Generator]]] = None, + latents: Optional[torch.FloatTensor] = None, + prompt_embeds: Optional[torch.FloatTensor] = None, + negative_prompt_embeds: Optional[torch.FloatTensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, + callback_steps: int = 1, + cross_attention_kwargs: Optional[dict[str, Any]] = None, + guidance_rescale: float = 0.0, +): + r""" + Function to get RGB image with gradients attached to the model weights. + + Args: + prompt (`str` or `list[str]`, *optional*, defaults to `None`): + The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds` + instead. + height (`int`, *optional*, defaults to `pipeline.unet.config.sample_size * pipeline.vae_scale_factor`): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to `pipeline.unet.config.sample_size * pipeline.vae_scale_factor`): + The width in pixels of the generated image. + num_inference_steps (`int`, *optional*, defaults to `50`): + 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*, defaults to `7.5`): + Guidance scale as defined in [Classifier-Free Diffusion + Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. of + [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting + `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the + text `prompt`, usually at the expense of lower image quality. + truncated_backprop (`bool`, *optional*, defaults to True): + Truncated Backpropation to fixed timesteps, helps prevent collapse during diffusion reward training as + shown in AlignProp (https://huggingface.co/papers/2310.03739). + truncated_backprop_rand (`bool`, *optional*, defaults to True): + Truncated Randomized Backpropation randomizes truncation to different diffusion timesteps, this helps + prevent collapse during diffusion reward training as shown in AlignProp + (https://huggingface.co/papers/2310.03739). Enabling truncated_backprop_rand allows adapting earlier + timesteps in diffusion while not resulting in a collapse. + gradient_checkpoint (`bool`, *optional*, defaults to True): + Adds gradient checkpointing to Unet forward pass. Reduces GPU memory consumption while slightly increasing + the training time. + truncated_backprop_timestep (`int`, *optional*, defaults to 49): + Absolute timestep to which the gradients are being backpropagated. Higher number reduces the memory usage + and reduces the chances of collapse. While a lower value, allows more semantic changes in the diffusion + generations, as the earlier diffusion timesteps are getting updated. However it also increases the chances + of collapse. + truncated_rand_backprop_minmax (`Tuple`, *optional*, defaults to (0,50)): + Range for randomized backprop. Here the value at 0 index indicates the earlier diffusion timestep to update + (closer to noise), while the value at index 1 indicates the later diffusion timestep to update. + negative_prompt (`str` or `list[str]`, *optional*): + The prompt or prompts not to guide the image generation. If not defined, one has to pass + `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is + less than `1`). + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + eta (`float`, *optional*, defaults to 0.0): + Corresponds to parameter eta (η) in the DDIM paper: https://huggingface.co/papers/2010.02502. Only applies + to [`schedulers.DDIMScheduler`], will be ignored for others. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to + make generation deterministic. + latents (`torch.FloatTensor`, *optional*): + Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image + generation. Can be used to tweak the same generation with different prompts. If not provided, a latents + tensor will ge generated by sampling using the supplied random `generator`. + prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + negative_prompt_embeds (`torch.FloatTensor`, *optional*): + Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. + If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. + output_type (`str`, *optional*, defaults to `"pil"`): + The output format of the generate image. Choose between [PIL](https://pillow.readthedocs.io/en/stable/): + `PIL.Image.Image` or `np.array`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a plain tuple. + callback (`Callable`, *optional*): + A function that will be called every `callback_steps` steps during inference. The function will be called + with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. + callback_steps (`int`, *optional*, defaults to 1): + The frequency at which the `callback` function will be called. If not specified, the callback will be + called at every step. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `pipeline.processor` in + [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py). + guidance_rescale (`float`, *optional*, defaults to 0.7): + Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are + Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of + [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). + Guidance rescale factor should fix overexposure when using zero terminal SNR. + + Examples: + + Returns: + `DDPOPipelineOutput`: The generated image, the predicted latents used to generate the image and the associated + log probabilities + """ + # 0. Default height and width to unet + height = height or pipeline.unet.config.sample_size * pipeline.vae_scale_factor + width = width or pipeline.unet.config.sample_size * pipeline.vae_scale_factor + + with torch.no_grad(): + # 1. Check inputs. Raise error if not correct + pipeline.check_inputs( + prompt, + height, + width, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + ) + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = pipeline._execution_device + # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) + # of the Imagen paper: https://huggingface.co/papers/2205.11487 . `guidance_scale = 1` + # corresponds to doing no classifier free guidance. + do_classifier_free_guidance = guidance_scale > 1.0 + + # 3. Encode input prompt + text_encoder_lora_scale = ( + cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None + ) + prompt_embeds = pipeline._encode_prompt( + prompt, + device, + num_images_per_prompt, + do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=text_encoder_lora_scale, + ) + + # 4. Prepare timesteps + pipeline.scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = pipeline.scheduler.timesteps + + # 5. Prepare latent variables + num_channels_latents = pipeline.unet.config.in_channels + latents = pipeline.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + # 6. Denoising loop + num_warmup_steps = len(timesteps) - num_inference_steps * pipeline.scheduler.order + all_latents = [latents] + all_log_probs = [] + with pipeline.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + # expand the latents if we are doing classifier free guidance + latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents + latent_model_input = pipeline.scheduler.scale_model_input(latent_model_input, t) + + # predict the noise residual + if gradient_checkpoint: + noise_pred = checkpoint.checkpoint( + pipeline.unet, + latent_model_input, + t, + prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + use_reentrant=False, + )[0] + else: + noise_pred = pipeline.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + cross_attention_kwargs=cross_attention_kwargs, + return_dict=False, + )[0] + + # truncating backpropagation is critical for preventing overoptimization (https://huggingface.co/papers/2304.05977). + if truncated_backprop: + # Randomized truncation randomizes the truncation process (https://huggingface.co/papers/2310.03739) + # the range of truncation is defined by truncated_rand_backprop_minmax + # Setting truncated_rand_backprop_minmax[0] to be low will allow the model to update earlier timesteps in the diffusion chain, while setitng it high will reduce the memory usage. + if truncated_backprop_rand: + rand_timestep = random.randint( + truncated_rand_backprop_minmax[0], truncated_rand_backprop_minmax[1] + ) + if i < rand_timestep: + noise_pred = noise_pred.detach() + else: + # fixed truncation process + if i < truncated_backprop_timestep: + noise_pred = noise_pred.detach() + + # perform guidance + if do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) + + if do_classifier_free_guidance and guidance_rescale > 0.0: + # Based on 3.4. in https://huggingface.co/papers/2305.08891 + noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale) + + # compute the previous noisy sample x_t -> x_t-1 + scheduler_output = scheduler_step(pipeline.scheduler, noise_pred, t, latents, eta) + latents = scheduler_output.latents + log_prob = scheduler_output.log_probs + + all_latents.append(latents) + all_log_probs.append(log_prob) + + # call the callback, if provided + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % pipeline.scheduler.order == 0): + progress_bar.update() + if callback is not None and i % callback_steps == 0: + callback(i, t, latents) + + if not output_type == "latent": + image = pipeline.vae.decode(latents / pipeline.vae.config.scaling_factor, return_dict=False)[0] + image, has_nsfw_concept = pipeline.run_safety_checker(image, device, prompt_embeds.dtype) + else: + image = latents + has_nsfw_concept = None + + if has_nsfw_concept is None: + do_denormalize = [True] * image.shape[0] + else: + do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept] + + image = pipeline.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize) + + # Offload last model to CPU + if hasattr(pipeline, "final_offload_hook") and pipeline.final_offload_hook is not None: + pipeline.final_offload_hook.offload() + + return DDPOPipelineOutput(image, all_latents, all_log_probs) + + +class DefaultDDPOStableDiffusionPipeline(DDPOStableDiffusionPipeline): + def __init__(self, pretrained_model_name: str, *, pretrained_model_revision: str = "main", use_lora: bool = True): + self.sd_pipeline = StableDiffusionPipeline.from_pretrained( + pretrained_model_name, revision=pretrained_model_revision + ) + + self.use_lora = use_lora + self.pretrained_model = pretrained_model_name + self.pretrained_revision = pretrained_model_revision + + try: + self.sd_pipeline.load_lora_weights( + pretrained_model_name, + weight_name="pytorch_lora_weights.safetensors", + revision=pretrained_model_revision, + ) + self.use_lora = True + except OSError: + if use_lora: + warnings.warn( + "Trying to load LoRA weights but no LoRA weights found. Set `use_lora=False` or check that " + "`pytorch_lora_weights.safetensors` exists in the model folder.", + UserWarning, + ) + + self.sd_pipeline.scheduler = DDIMScheduler.from_config(self.sd_pipeline.scheduler.config) + self.sd_pipeline.safety_checker = None + + # memory optimization + self.sd_pipeline.vae.requires_grad_(False) + self.sd_pipeline.text_encoder.requires_grad_(False) + self.sd_pipeline.unet.requires_grad_(not self.use_lora) + + def __call__(self, *args, **kwargs) -> DDPOPipelineOutput: + return pipeline_step(self.sd_pipeline, *args, **kwargs) + + def rgb_with_grad(self, *args, **kwargs) -> DDPOPipelineOutput: + return pipeline_step_with_grad(self.sd_pipeline, *args, **kwargs) + + def scheduler_step(self, *args, **kwargs) -> DDPOSchedulerOutput: + return scheduler_step(self.sd_pipeline.scheduler, *args, **kwargs) + + @property + def unet(self): + return self.sd_pipeline.unet + + @property + def vae(self): + return self.sd_pipeline.vae + + @property + def tokenizer(self): + return self.sd_pipeline.tokenizer + + @property + def scheduler(self): + return self.sd_pipeline.scheduler + + @property + def text_encoder(self): + return self.sd_pipeline.text_encoder + + @property + def autocast(self): + return contextlib.nullcontext if self.use_lora else None + + def save_pretrained(self, output_dir): + if self.use_lora: + state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(self.sd_pipeline.unet)) + self.sd_pipeline.save_lora_weights(save_directory=output_dir, unet_lora_layers=state_dict) + self.sd_pipeline.save_pretrained(output_dir) + + def set_progress_bar_config(self, *args, **kwargs): + self.sd_pipeline.set_progress_bar_config(*args, **kwargs) + + def get_trainable_layers(self): + if self.use_lora: + lora_config = LoraConfig( + r=4, + lora_alpha=4, + init_lora_weights="gaussian", + target_modules=["to_k", "to_q", "to_v", "to_out.0"], + ) + self.sd_pipeline.unet.add_adapter(lora_config) + + # To avoid accelerate unscaling problems in FP16. + for param in self.sd_pipeline.unet.parameters(): + # only upcast trainable parameters (LoRA) into fp32 + if param.requires_grad: + param.data = param.to(torch.float32) + return self.sd_pipeline.unet + else: + return self.sd_pipeline.unet + + def save_checkpoint(self, models, weights, output_dir): + if len(models) != 1: + raise ValueError("Given how the trainable params were set, this should be of length 1") + if self.use_lora and hasattr(models[0], "peft_config") and getattr(models[0], "peft_config", None) is not None: + state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(models[0])) + self.sd_pipeline.save_lora_weights(save_directory=output_dir, unet_lora_layers=state_dict) + elif not self.use_lora and isinstance(models[0], UNet2DConditionModel): + models[0].save_pretrained(os.path.join(output_dir, "unet")) + else: + raise ValueError(f"Unknown model type {type(models[0])}") + + def load_checkpoint(self, models, input_dir): + if len(models) != 1: + raise ValueError("Given how the trainable params were set, this should be of length 1") + if self.use_lora: + lora_state_dict, network_alphas = self.sd_pipeline.lora_state_dict( + input_dir, weight_name="pytorch_lora_weights.safetensors" + ) + self.sd_pipeline.load_lora_into_unet(lora_state_dict, network_alphas=network_alphas, unet=models[0]) + + elif not self.use_lora and isinstance(models[0], UNet2DConditionModel): + load_model = UNet2DConditionModel.from_pretrained(input_dir, subfolder="unet") + models[0].register_to_config(**load_model.config) + models[0].load_state_dict(load_model.state_dict()) + del load_model + else: + raise ValueError(f"Unknown model type {type(models[0])}") diff --git a/lib/python3.12/site-packages/trl/models/modeling_value_head.py b/lib/python3.12/site-packages/trl/models/modeling_value_head.py new file mode 100644 index 0000000000000000000000000000000000000000..f32953be412494b2be6657f632586e15e01121c7 --- /dev/null +++ b/lib/python3.12/site-packages/trl/models/modeling_value_head.py @@ -0,0 +1,434 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import torch.nn as nn +from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, is_torch_npu_available, is_torch_xpu_available + +from .modeling_base import PreTrainedModelWrapper + + +class ValueHead(nn.Module): + r""" + The ValueHead class implements a head for GPT2 that returns a scalar for each output token. + """ + + def __init__(self, config, **kwargs): + super().__init__() + if not hasattr(config, "summary_dropout_prob"): + summary_dropout_prob = kwargs.pop("summary_dropout_prob", 0.1) + else: + summary_dropout_prob = config.summary_dropout_prob + + self.dropout = nn.Dropout(summary_dropout_prob) if summary_dropout_prob else nn.Identity() + + # some models such as OPT have a projection layer before the word embeddings - e.g. OPT-350m + if hasattr(config, "hidden_size"): + hidden_size = config.hidden_size + if hasattr(config, "word_embed_proj_dim"): + hidden_size = config.word_embed_proj_dim + elif hasattr(config, "is_encoder_decoder"): + if config.is_encoder_decoder and hasattr(config, "decoder"): + if hasattr(config.decoder, "hidden_size"): + hidden_size = config.decoder.hidden_size + + self.summary = nn.Linear(hidden_size, 1) + + self.flatten = nn.Flatten() + + def forward(self, hidden_states): + output = self.dropout(hidden_states) + + # For now force upcast in fp32 if needed. Let's keep the + # output in fp32 for numerical stability. + if output.dtype != self.summary.weight.dtype: + output = output.to(self.summary.weight.dtype) + + output = self.summary(output) + return output + + +class AutoModelForCausalLMWithValueHead(PreTrainedModelWrapper): + r""" + An autoregressive model with a value head in addition to the language model head. This class inherits from + `~trl.PreTrainedModelWrapper` and wraps a `transformers.PreTrainedModel` class. The wrapper class supports classic + functions such as `from_pretrained`, `push_to_hub` and `generate`. To call a method of the wrapped model, simply + manipulate the `pretrained_model` attribute of this class. + + Class attributes: + - **transformers_parent_class** (`transformers.PreTrainedModel`) -- The parent class of the wrapped model. This + should be set to `transformers.AutoModelForCausalLM` for this class. + - **supported_args** (`tuple`) -- A tuple of strings that are used to identify the arguments that are supported + by the `ValueHead` class. Currently, the supported args are: + - **summary_dropout_prob** (`float`, `optional`, defaults to `None`) -- The dropout probability for the + `ValueHead` class. + - **v_head_initializer_range** (`float`, `optional`, defaults to `0.2`) -- The initializer range for the + `ValueHead` if a specific initialization strategy is selected. + - **v_head_init_strategy** (`str`, `optional`, defaults to `None`) -- The initialization strategy for the + `ValueHead`. Currently, the supported strategies are: + - **`None`** -- Initializes the weights of the `ValueHead` with a random distribution. This is the + default strategy. + - **"normal"** -- Initializes the weights of the `ValueHead` with a normal distribution. + """ + + transformers_parent_class = AutoModelForCausalLM + supported_args = ( + "summary_dropout_prob", + "v_head_initializer_range", + "v_head_init_strategy", + ) + + def __init__(self, pretrained_model, **kwargs): + r""" + Initializes the model. + + Args: + pretrained_model (`transformers.PreTrainedModel`): + The model to wrap. It should be a causal language model such as GPT2. or any model mapped inside the + `AutoModelForCausalLM` class. + kwargs (`dict`, `optional`): + Additional keyword arguments, that are passed to the `ValueHead` class. + """ + super().__init__(pretrained_model, **kwargs) + v_head_kwargs, _, _ = self._split_kwargs(kwargs) + self.v_head = ValueHead(self.pretrained_model.config, **v_head_kwargs) + self._init_weights(**v_head_kwargs) + + def _init_weights(self, **kwargs): + r""" + Initializes the weights of the value head. The default initialization strategy is random. Users can pass a + different initialization strategy by passing the `v_head_init_strategy` argument when calling + `.from_pretrained`. Supported strategies are: + - `normal`: initializes the weights with a normal distribution. + + Args: + **kwargs (`dict`, `optional`): + Additional keyword arguments, that are passed to the `ValueHead` class. These arguments can contain the + `v_head_init_strategy` argument as well as the `v_head_initializer_range` argument. + """ + initializer_range = kwargs.pop("v_head_initializer_range", 0.2) + # random init by default + init_strategy = kwargs.pop("v_head_init_strategy", None) + if init_strategy is None: + # do nothing + pass + elif init_strategy == "normal": + self.v_head.summary.weight.data.normal_(mean=0.0, std=initializer_range) + self.v_head.summary.bias.data.zero_() + + def forward( + self, + input_ids=None, + past_key_values=None, + attention_mask=None, + return_past_key_values=False, + **kwargs, + ): + r""" + Applies a forward pass to the wrapped model and returns the logits of the value head. + + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. + past_key_values (`tuple(tuple(torch.FloatTensor))`, `optional`): + Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model + (see `past_key_values` input) to speed up sequential decoding. + attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, `optional`): + Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + return_past_key_values (bool): A flag indicating if the computed hidden-states should be returned. + kwargs (`dict`, `optional`): + Additional keyword arguments, that are passed to the wrapped model. + """ + kwargs["output_hidden_states"] = True # this had already been set in the LORA / PEFT examples + kwargs["past_key_values"] = past_key_values + + if self.is_peft_model and self.pretrained_model.active_peft_config.peft_type == "PREFIX_TUNING": + kwargs.pop("past_key_values") + + base_model_output = self.pretrained_model( + input_ids=input_ids, + attention_mask=attention_mask, + **kwargs, + ) + + last_hidden_state = base_model_output.hidden_states[-1] + lm_logits = base_model_output.logits + loss = base_model_output.loss + + if last_hidden_state.device != self.v_head.summary.weight.device: + last_hidden_state = last_hidden_state.to(self.v_head.summary.weight.device) + + value = self.v_head(last_hidden_state).squeeze(-1) + + # force upcast in fp32 if logits are in half-precision + if lm_logits.dtype != torch.float32: + lm_logits = lm_logits.float() + + if return_past_key_values: + return (lm_logits, loss, value, base_model_output.past_key_values) + else: + return (lm_logits, loss, value) + + def generate(self, *args, **kwargs): + r""" + A simple wrapper around the `generate` method of the wrapped model. Please refer to the + [`generate`](https://huggingface.co/docs/transformers/internal/generation_utils) method of the wrapped model + for more information about the supported arguments. + + Args: + *args (`list`, *optional*): + Positional arguments passed to the `generate` method of the wrapped model. + **kwargs (`dict`, *optional*): + Keyword arguments passed to the `generate` method of the wrapped model. + """ + return self.pretrained_model.generate(*args, **kwargs) + + def state_dict(self, *args, **kwargs): + r""" + Returns the state dictionary of the model. We add the state dictionary of the value head to the state + dictionary of the wrapped model by prepending the key with `v_head.`. + """ + if not self.is_peft_model: + pretrained_model_state_dict = self.pretrained_model.state_dict(*args, **kwargs) + else: + # if it is a peft model, only save the v_head + pretrained_model_state_dict = {} + + v_head_state_dict = self.v_head.state_dict(*args, **kwargs) + for k, v in v_head_state_dict.items(): + pretrained_model_state_dict[f"v_head.{k}"] = v + return pretrained_model_state_dict + + def push_to_hub(self, *args, **kwargs): + self.pretrained_model.v_head = self.v_head + + return self.pretrained_model.push_to_hub(*args, **kwargs) + + def post_init(self, state_dict): + r""" + We add the state dictionary of the value head to the state dictionary of the wrapped model by prepending the + key with `v_head.`. This function removes the `v_head.` prefix from the keys of the value head state + dictionary. + """ + for k in list(state_dict.keys()): + if "v_head." in k: + state_dict[k.replace("v_head.", "")] = state_dict.pop(k) + self.v_head.load_state_dict(state_dict, strict=False) + del state_dict + + if hasattr(self.pretrained_model, "hf_device_map"): + if ( + "cpu" in self.pretrained_model.hf_device_map.values() + or "disk" in self.pretrained_model.hf_device_map.values() + ): + raise ValueError( + "The model is offloaded on CPU or disk - CPU & disk offloading is not supported for ValueHead models." + ) + + first_device = list(set(self.pretrained_model.hf_device_map.values()))[0] + if isinstance(first_device, int): + if is_torch_npu_available(): + first_device = f"npu:{first_device}" + elif is_torch_xpu_available(): + first_device = f"xpu:{first_device}" + else: + first_device = f"cuda:{first_device}" + self.v_head = self.v_head.to(first_device) + + def set_device_hook(module, input, outputs): + new_output = () + for output in outputs: + if isinstance(output, torch.Tensor): + new_output += (output.to(first_device),) + else: + new_output += (output,) + return new_output + + self.register_forward_hook(set_device_hook) + + self.is_sequential_parallel = True + + +class AutoModelForSeq2SeqLMWithValueHead(PreTrainedModelWrapper): + r""" + A seq2seq model with a value head in addition to the language model head. This class inherits from + `~trl.PreTrainedModelWrapper` and wraps a `transformers.PreTrainedModel` class. The wrapper class supports classic + functions such as `from_pretrained` and `push_to_hub` and also provides some additional functionalities such as + `generate`. + + Args: + pretrained_model (`transformers.PreTrainedModel`): + The model to wrap. It should be a causal language model such as GPT2. or any model mapped inside the + `AutoModelForSeq2SeqLM` class. + kwargs: + Additional keyword arguments passed along to the `ValueHead` class. + """ + + transformers_parent_class = AutoModelForSeq2SeqLM + lm_head_namings = ["lm_head", "embed_out", "output_projection"] + supported_args = ( + "summary_dropout_prob", + "v_head_initializer_range", + "v_head_init_strategy", + ) + + def __init__(self, pretrained_model, **kwargs): + super().__init__(pretrained_model, **kwargs) + v_head_kwargs, _, _ = self._split_kwargs(kwargs) + self.is_encoder_decoder = True + + if not self._has_lm_head(): + raise ValueError("The model does not have a language model head, please use a model that has one.") + + self.v_head = ValueHead(self.pretrained_model.config, **v_head_kwargs) + + self._init_weights(**v_head_kwargs) + + def _has_lm_head(self): + # check module names of all modules inside `pretrained_model` to find the language model head + for name, _module in self.pretrained_model.named_modules(): + if any(attribute in name for attribute in self.lm_head_namings): + return True + return False + + def post_init(self, state_dict): + r""" + We add the state dictionary of the value head to the state dictionary of the wrapped model by prepending the + key with `v_head.`. This function removes the `v_head.` prefix from the keys of the value head state + dictionary. + """ + for k in list(state_dict.keys()): + if "v_head." in k: + state_dict[k.replace("v_head.", "")] = state_dict.pop(k) + self.v_head.load_state_dict(state_dict, strict=False) + del state_dict + + if hasattr(self.pretrained_model, "hf_device_map"): + if ( + "cpu" in self.pretrained_model.hf_device_map.values() + or "disk" in self.pretrained_model.hf_device_map.values() + ): + raise ValueError( + "The model is offloaded on CPU or disk - CPU & disk offloading is not supported for ValueHead models." + ) + + # get the lm_head device + for name, module in self.pretrained_model.named_modules(): + if any(attribute in name for attribute in self.lm_head_namings): + lm_head_device = module.weight.device + break + + # put v_head on the same device as the lm_head to avoid issues + self.v_head = self.v_head.to(lm_head_device) + + def set_device_hook(module, input, outputs): + r""" + A hook that sets the device of the output of the model to the device of the first parameter of the + model. + + Args: + module (`nn.Module`): + The module to which the hook is attached. + input (`tuple`): + The input to the module. + outputs (`tuple`): + The output of the module. + """ + new_output = () + for output in outputs: + if isinstance(output, torch.Tensor): + new_output += (output.to(lm_head_device),) + else: + new_output += (output,) + return new_output + + self.register_forward_hook(set_device_hook) + self.is_sequential_parallel = True + + def state_dict(self, *args, **kwargs): + r""" + Returns the state dictionary of the model. We add the state dictionary of the value head to the state + dictionary of the wrapped model by prepending the key with `v_head.`. + """ + if not self.is_peft_model: + pretrained_model_state_dict = self.pretrained_model.state_dict(*args, **kwargs) + else: + # if it is a peft model, only save the v_head + pretrained_model_state_dict = {} + + v_head_state_dict = self.v_head.state_dict(*args, **kwargs) + for k, v in v_head_state_dict.items(): + pretrained_model_state_dict[f"v_head.{k}"] = v + return pretrained_model_state_dict + + def push_to_hub(self, *args, **kwargs): + self.pretrained_model.v_head = self.v_head + + return self.pretrained_model.push_to_hub(*args, **kwargs) + + def _init_weights(self, **kwargs): + r""" + We initialize the weights of the value head. + """ + initializer_range = kwargs.pop("v_head_initializer_range", 0.2) + # random init by default + init_strategy = kwargs.pop("v_head_init_strategy", None) + if init_strategy is None: + # do nothing + pass + elif init_strategy == "normal": + self.v_head.summary.weight.data.normal_(mean=0.0, std=initializer_range) + self.v_head.summary.bias.data.zero_() + + def forward( + self, + input_ids=None, + past_key_values=None, + attention_mask=None, + return_past_key_values=False, + **kwargs, + ): + kwargs["past_key_values"] = past_key_values + if self.is_peft_model and self.pretrained_model.active_peft_config.peft_type == "PREFIX_TUNING": + kwargs.pop("past_key_values") + + base_model_output = self.pretrained_model( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, # We force the model to output hidden states + **kwargs, + ) + + last_hidden_state = base_model_output.decoder_hidden_states[-1] + lm_logits = base_model_output.logits + loss = base_model_output.loss + + value = self.v_head(last_hidden_state).squeeze(-1) + + # force upcast in fp32 if logits are in half-precision + if lm_logits.dtype != torch.float32: + lm_logits = lm_logits.float() + + if return_past_key_values: + return (lm_logits, loss, value, base_model_output.past_key_values) + else: + return (lm_logits, loss, value) + + def generate(self, *args, **kwargs): + r""" + We call `generate` on the wrapped model. + """ + return self.pretrained_model.generate(*args, **kwargs) diff --git a/lib/python3.12/site-packages/trl/models/sd_utils.py b/lib/python3.12/site-packages/trl/models/sd_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ef40b7a803b0b87cd25ff6cfb6d755362acc41fe --- /dev/null +++ b/lib/python3.12/site-packages/trl/models/sd_utils.py @@ -0,0 +1,149 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +State dict utilities: utility methods for converting state dicts easily File copied from diffusers to avoid import +issues and make TRL compatible with most of diffusers versions. +""" + +import enum + + +class StateDictType(enum.Enum): + """ + The mode to use when converting state dicts. + """ + + DIFFUSERS_OLD = "diffusers_old" + PEFT = "peft" + + +PEFT_TO_DIFFUSERS = { + ".q_proj.lora_B": ".q_proj.lora_linear_layer.up", + ".q_proj.lora_A": ".q_proj.lora_linear_layer.down", + ".k_proj.lora_B": ".k_proj.lora_linear_layer.up", + ".k_proj.lora_A": ".k_proj.lora_linear_layer.down", + ".v_proj.lora_B": ".v_proj.lora_linear_layer.up", + ".v_proj.lora_A": ".v_proj.lora_linear_layer.down", + ".out_proj.lora_B": ".out_proj.lora_linear_layer.up", + ".out_proj.lora_A": ".out_proj.lora_linear_layer.down", + "to_k.lora_A": "to_k.lora.down", + "to_k.lora_B": "to_k.lora.up", + "to_q.lora_A": "to_q.lora.down", + "to_q.lora_B": "to_q.lora.up", + "to_v.lora_A": "to_v.lora.down", + "to_v.lora_B": "to_v.lora.up", + "to_out.0.lora_A": "to_out.0.lora.down", + "to_out.0.lora_B": "to_out.0.lora.up", +} + +DIFFUSERS_OLD_TO_DIFFUSERS = { + ".to_q_lora.up": ".q_proj.lora_linear_layer.up", + ".to_q_lora.down": ".q_proj.lora_linear_layer.down", + ".to_k_lora.up": ".k_proj.lora_linear_layer.up", + ".to_k_lora.down": ".k_proj.lora_linear_layer.down", + ".to_v_lora.up": ".v_proj.lora_linear_layer.up", + ".to_v_lora.down": ".v_proj.lora_linear_layer.down", + ".to_out_lora.up": ".out_proj.lora_linear_layer.up", + ".to_out_lora.down": ".out_proj.lora_linear_layer.down", +} + +DIFFUSERS_STATE_DICT_MAPPINGS = { + StateDictType.DIFFUSERS_OLD: DIFFUSERS_OLD_TO_DIFFUSERS, + StateDictType.PEFT: PEFT_TO_DIFFUSERS, +} + +KEYS_TO_ALWAYS_REPLACE = { + ".processor.": ".", +} + + +def convert_state_dict(state_dict, mapping): + r""" + Simply iterates over the state dict and replaces the patterns in `mapping` with the corresponding values. + + Args: + state_dict (`dict[str, torch.Tensor]`): + The state dict to convert. + mapping (`dict[str, str]`): + The mapping to use for conversion, the mapping should be a dictionary with the following structure: + - key: the pattern to replace + - value: the pattern to replace with + + Returns: + converted_state_dict (`dict`) + The converted state dict. + """ + converted_state_dict = {} + for k, v in state_dict.items(): + # First, filter out the keys that we always want to replace + for pattern in KEYS_TO_ALWAYS_REPLACE.keys(): + if pattern in k: + new_pattern = KEYS_TO_ALWAYS_REPLACE[pattern] + k = k.replace(pattern, new_pattern) + + for pattern in mapping.keys(): + if pattern in k: + new_pattern = mapping[pattern] + k = k.replace(pattern, new_pattern) + break + converted_state_dict[k] = v + return converted_state_dict + + +def convert_state_dict_to_diffusers(state_dict, original_type=None, **kwargs): + r""" + Converts a state dict to new diffusers format. The state dict can be from previous diffusers format + (`OLD_DIFFUSERS`), or PEFT format (`PEFT`) or new diffusers format (`DIFFUSERS`). In the last case the method will + return the state dict as is. + + The method only supports the conversion from diffusers old, PEFT to diffusers new for now. + + Args: + state_dict (`dict[str, torch.Tensor]`): + The state dict to convert. + original_type (`StateDictType`, *optional*): + The original type of the state dict, if not provided, the method will try to infer it automatically. + kwargs (`dict`, *args*): + Additional arguments to pass to the method. + + - **adapter_name**: For example, in case of PEFT, some keys will be pre-pended + with the adapter name, therefore needs a special handling. By default PEFT also takes care of that in + `get_peft_model_state_dict` method: + https://github.com/huggingface/peft/blob/ba0477f2985b1ba311b83459d29895c809404e99/src/peft/utils/save_and_load.py#L92 + but we add it here in case we don't want to rely on that method. + """ + peft_adapter_name = kwargs.pop("adapter_name", None) + if peft_adapter_name is not None: + peft_adapter_name = "." + peft_adapter_name + else: + peft_adapter_name = "" + + if original_type is None: + # Old diffusers to PEFT + if any("to_out_lora" in k for k in state_dict.keys()): + original_type = StateDictType.DIFFUSERS_OLD + elif any(f".lora_A{peft_adapter_name}.weight" in k for k in state_dict.keys()): + original_type = StateDictType.PEFT + elif any("lora_linear_layer" in k for k in state_dict.keys()): + # nothing to do + return state_dict + else: + raise ValueError("Could not automatically infer state dict type") + + if original_type not in DIFFUSERS_STATE_DICT_MAPPINGS.keys(): + raise ValueError(f"Original type {original_type} is not supported") + + mapping = DIFFUSERS_STATE_DICT_MAPPINGS[original_type] + return convert_state_dict(state_dict, mapping) diff --git a/lib/python3.12/site-packages/trl/models/utils.py b/lib/python3.12/site-packages/trl/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e3066511c17d9da240c233c466996344443c0f72 --- /dev/null +++ b/lib/python3.12/site-packages/trl/models/utils.py @@ -0,0 +1,424 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools +from contextlib import contextmanager +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal, Optional, Union + +import torch.nn as nn +from packaging import version +from transformers import AutoTokenizer, PreTrainedModel, PreTrainedTokenizer + +from .modeling_value_head import AutoModelForCausalLMWithValueHead, AutoModelForSeq2SeqLMWithValueHead + + +SUPPORTED_ARCHITECTURES = ( + AutoModelForCausalLMWithValueHead, + AutoModelForSeq2SeqLMWithValueHead, +) + +if TYPE_CHECKING: + from accelerate import Accelerator + from deepspeed.runtime.engine import DeepSpeedEngine + from torch.nn import Module + from torch.nn.parallel.distributed import DistributedDataParallel + + +# TODO: Add Abstract Base Class if more formats are added +@dataclass +class ChatMlSpecialTokens: + """Dataclass for special tokens used in ChatML, including system, user, assistant, bos, eos, and pad tokens.""" + + bos_token: str = "<|im_start|>" + eos_token: str = "<|im_end|>" + pad_token: str = "<|im_end|>" + + @property + def system(self): + return f"{self.bos_token}system" + + @property + def user(self): + return f"{self.bos_token}user" + + @property + def assistant(self): + return f"{self.bos_token}assistant" + + @property + def chat_template(self): + return ( + "{% for message in messages %}" + f"{{{{'{self.bos_token}' + message['role'] + '\n' + message['content'] + '{self.eos_token}' + '\n'}}}}" + "{% endfor %}" + "{% if add_generation_prompt %}" + f"{{{{ '{self.assistant}\n' }}}}" + "{% endif %}" + ) + + +FORMAT_MAPPING = {"chatml": ChatMlSpecialTokens} + + +def setup_chat_format( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + format: Optional[Literal["chatml"]] = "chatml", + resize_to_multiple_of: Optional[int] = None, +) -> tuple[PreTrainedModel, PreTrainedTokenizer]: + """ + Setup chat format by adding special tokens to the tokenizer, setting the correct format, and extending the + embedding layer of the model based on the new special tokens. + + We recommend using [`clone_chat_template`] instead of this function. + + + + If the model already has a chat template, this will throw an error. If you want to overwrite it, please set + `tokenizer.chat_template` to `None`. + + Args: + model (`~transformers.PreTrainedModel`): The model to be modified. + tokenizer (`~transformers.PreTrainedTokenizer`): The tokenizer to be modified. + format (`Optional[Literal["chatml"]]`): The format to be set. Defaults to "chatml". + resize_to_multiple_of (`int` or `None`): Number to resize the embedding layer to. Defaults to None. + + Returns: + model (`~transformers.PreTrainedModel`): + The modified model. + tokenizer (`~transformers.PreTrainedTokenizer`): + The modified tokenizer. + """ + # check if model already had a chat template + if tokenizer.chat_template is not None: + raise ValueError( + "Chat template is already added to the tokenizer. If you want to overwrite it, please set it to None" + ) + + # check if format available and retrieve + if format not in FORMAT_MAPPING: + raise ValueError(f"Format {format} not available. Please use one of {FORMAT_MAPPING.keys()}") + + chat_format = FORMAT_MAPPING[format]() + + # set special tokens and them + tokenizer.eos_token = chat_format.eos_token + tokenizer.pad_token = chat_format.pad_token + tokenizer.bos_token = chat_format.bos_token + tokenizer.add_special_tokens({"additional_special_tokens": [chat_format.bos_token, chat_format.eos_token]}) + # set chat format for tokenizer + tokenizer.chat_template = chat_format.chat_template + + # resize embedding layer to a multiple of 64, https://x.com/karpathy/status/1621578354024677377 + model.resize_token_embeddings( + # After studying many tokenizers, we found that len(tokenizer.vocab) is the most reliable way to get the vocab + # size. Avoid using tokenizer.vocab_size or tokenizer.vocab_size + len(tokenizer.added_tokens_encoder), + # as handling of special and added tokens varies across tokenizers. + new_num_tokens=len(tokenizer.vocab), + pad_to_multiple_of=resize_to_multiple_of if resize_to_multiple_of is not None else None, + ) + # Update the model config to use the new eos & bos tokens + if getattr(model, "config", None) is not None: + model.config.pad_token_id = tokenizer.pad_token_id + model.config.bos_token_id = tokenizer.bos_token_id + model.config.eos_token_id = tokenizer.eos_token_id + # Update the generation config to use the new eos & bos token + if getattr(model, "generation_config", None) is not None: + model.generation_config.bos_token_id = tokenizer.bos_token_id + model.generation_config.eos_token_id = tokenizer.eos_token_id + model.generation_config.pad_token_id = tokenizer.pad_token_id + + return model, tokenizer + + +def clone_chat_template( + model: PreTrainedModel, + tokenizer: PreTrainedTokenizer, + source_tokenizer_path: str, + resize_to_multiple_of: Optional[int] = 64, +) -> tuple[PreTrainedModel, PreTrainedTokenizer]: + """ + Clones a chat template from a source tokenizer to the target tokenizer and updates the model accordingly. + + This function: + - Copies the chat template from a source tokenizer to the target tokenizer. + - Adds any new tokens from the source tokenizer to the target tokenizer. + - Sets and synchronizes the EOS token across the tokenizer and model. + - Resizes the model's token embeddings to match the new vocabulary size, optionally rounding it up to a multiple of + a specified value. + + Args: + model (`PreTrainedModel`): + Model to update. + tokenizer (`PreTrainedTokenizer`): + Tokenizer to update. + source_tokenizer_path (`str`): + Path or identifier of the pretrained tokenizer to clone from. + resize_to_multiple_of (`int` or `None`, *optional*, defaults to `64`): + The embedding layer will be resized to the new vocabulary size. If this is not `None`, it will round up the + new vocabulary size to the nearest multiple of this value. + + Returns: + model (`PreTrainedModel`): + Updated model with resized token embeddings and EOS token configured. + tokenizer (`~transformers.PreTrainedTokenizer`): + Updated tokenizer with the chat template and special tokens applied. + + Example: + ```python + from transformers import AutoModelForCausalLM, AutoTokenizer + from trl import clone_chat_template + + model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B") + tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B") + model, tokenizer = clone_chat_template(model, tokenizer, "Qwen/Qwen3-0.6B") + ``` + """ + # Load the source tokenizer containing the desired chat template + tokenizer_source = AutoTokenizer.from_pretrained(source_tokenizer_path) + + # Copy the chat template from the source tokenizer + tokenizer.chat_template = tokenizer_source.get_chat_template() + + # Ensure all added tokens from the source are available in the target tokenizer + tokenizer.add_tokens(list(tokenizer_source.added_tokens_decoder.values())) + + # Set the EOS token from the source tokenizer (important for generation) + tokenizer.eos_token = tokenizer_source.eos_token + model.config.eos_token_id = tokenizer.eos_token_id + model.generation_config.eos_token_id = tokenizer.eos_token_id + + # Resize model embeddings to include any new tokens, optionally rounding up to a multiple + model.resize_token_embeddings( + # After studying many tokenizers, we found that len(tokenizer.vocab) is the most reliable way to get the vocab + # size. Avoid using tokenizer.vocab_size or tokenizer.vocab_size + len(tokenizer.added_tokens_encoder), + # as handling of special and added tokens varies across tokenizers. + new_num_tokens=len(tokenizer.vocab), + pad_to_multiple_of=resize_to_multiple_of if resize_to_multiple_of is not None else None, + ) + + return model, tokenizer + + +def remove_hooks(model: "DeepSpeedEngine") -> None: + """Removes the optimizer hooks from a DeepSpeed ZeRO-3 model.""" + if not hasattr(model, "optimizer"): # before the first training step, the model has no optimizer + return + if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): + optimizer_offload = model.optimizer.parameter_offload + elif model.optimizer is not None: + optimizer_offload = model.optimizer + else: + raise RuntimeError("The model optimizer is None, which is not yet supported.") + + for param in iter_params(optimizer_offload.module, recurse=True): + param.ds_active_sub_modules.clear() + + for hook in optimizer_offload.forward_hooks: + hook.remove() + for hook in optimizer_offload.backward_hooks: + hook.remove() + + optimizer_offload.forward_hooks = [] + optimizer_offload.backward_hooks = [] + + +def get_all_parameters(sub_module, recurse=False): + return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters()) + + +def iter_params(module, recurse=False): + return [param for _, param in get_all_parameters(module, recurse)] + + +def add_hooks(model: "DeepSpeedEngine") -> None: + """Adds the optimizer hooks from a DeepSpeed ZeRO-3 model.""" + import deepspeed + + if not hasattr(model, "optimizer"): # before the first training step, the model has no optimizer + return + if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"): + optimizer_offload = model.optimizer.parameter_offload + elif model.optimizer is not None: + optimizer_offload = model.optimizer + else: + raise RuntimeError("The model optimizer is None, which is not yet supported.") + if version.parse(deepspeed.__version__) >= version.parse("0.16.4"): + # Account for renaming in https://github.com/deepspeedai/DeepSpeed/pull/6847 + optimizer_offload._register_deepspeed_module(optimizer_offload.module) + else: + optimizer_offload._register_hooks_recursively(optimizer_offload.module) + + +@contextmanager +def unwrap_model_for_generation( + model: Union["DistributedDataParallel", "DeepSpeedEngine"], + accelerator: "Accelerator", + gather_deepspeed3_params: bool = True, +): + """ + Context manager to unwrap distributed or accelerated models for generation tasks. + + Args: + model (`Union[DistributedDataParallel, DeepSpeedEngine]`): + Model to be unwrapped. + accelerator (`~accelerate.Accelerator`): + Accelerator instance managing the model. + gather_deepspeed3_params (`bool`, *optional*, defaults to `True`): + Whether to gather weights for DeepSpeed ZeRO Stage 3 models. If `False`, skips parameter gathering, which + can be more memory-efficient but may lead to slower generation times. + + Yields: + Unwrapped model. + + Example: + ```python + with unwrap_model_for_generation(model, accelerator) as unwrapped_model: + generated_outputs = unwrapped_model.generate(input_ids) + ``` + """ + unwrapped_model = accelerator.unwrap_model(model) + if accelerator.state.deepspeed_plugin is not None and accelerator.state.deepspeed_plugin.zero_stage == 3: + if not gather_deepspeed3_params: + yield accelerator.unwrap_model(model) + else: + import deepspeed + + with deepspeed.zero.GatheredParameters(model.parameters()): + remove_hooks(model) + yield accelerator.unwrap_model(model) + add_hooks(model) + else: + yield unwrapped_model + + +def prepare_deepspeed(model: "Module", accelerator: "Accelerator"): + """Prepares the model for DeepSpeed inference or evaluation by initializing it with the appropriate configuration. + + Adapted from accelerate: + https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1473 + """ + import deepspeed # local import (instead of top-level) to avoid DS init interfering with other backends (like vllm): https://github.com/deepspeedai/DeepSpeed/issues/7252 + + deepspeed_plugin = accelerator.state.deepspeed_plugin + config_kwargs = deepcopy(deepspeed_plugin.deepspeed_config) + stage = config_kwargs["zero_optimization"]["stage"] + + if model is not None: + hidden_size = ( + max(model.config.hidden_sizes) + if getattr(model.config, "hidden_sizes", None) + else getattr(model.config, "hidden_size", None) + ) + if hidden_size is not None and stage == 3: + # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache + # @ step 0: expected module 1, but got module 0` + # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081 + config_kwargs.update( + { + "zero_optimization.reduce_bucket_size": hidden_size * hidden_size, + "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size, + "zero_optimization.stage3_prefetch_bucket_size": 0.9 * hidden_size * hidden_size, + } + ) + + # If ZeRO-3 is used, we shard both the active and reference model. + # Otherwise, we assume the reference model fits in memory and is initialized on each device with ZeRO + # disabled (stage 0) + if stage != 3: + config_kwargs["zero_optimization"]["stage"] = 0 + model, *_ = deepspeed.initialize(model=model, config=config_kwargs) + model.eval() + return model + + +def prepare_fsdp(model, accelerator): + # Adapted from accelerate: https://github.com/huggingface/accelerate/blob/739b135f8367becb67ffaada12fe76e3aa60fefd/src/accelerate/accelerator.py#L1421 + from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP + + # Check if the model is already a FSDP model due to `Manual Wrapping` and if so, + # don't wrap it again + if not isinstance(model, FSDP): + accelerator.state.fsdp_plugin.set_auto_wrap_policy(model) + fsdp_plugin = accelerator.state.fsdp_plugin + kwargs = { + "sharding_strategy": fsdp_plugin.sharding_strategy or fsdp_plugin.reshard_after_forward, + "cpu_offload": fsdp_plugin.cpu_offload, + "auto_wrap_policy": fsdp_plugin.auto_wrap_policy, + "mixed_precision": fsdp_plugin.mixed_precision_policy, + "sync_module_states": fsdp_plugin.sync_module_states, + "backward_prefetch": fsdp_plugin.backward_prefetch, + "forward_prefetch": fsdp_plugin.forward_prefetch, + "use_orig_params": fsdp_plugin.use_orig_params, + "param_init_fn": fsdp_plugin.param_init_fn, + "ignored_modules": fsdp_plugin.ignored_modules, + "limit_all_gathers": fsdp_plugin.limit_all_gathers, + "device_id": accelerator.device, + } + model = FSDP(model, **kwargs) + model.eval() + return model + + +class _ForwardRedirection: + """Implements the `forward-redirection`. + + Taken from Pytorch-lightning: + https://github.com/Lightning-AI/pytorch-lightning/blob/02311d03fb982560246eead7c08104481fac9579/src/lightning/pytorch/strategies/strategy.py#L602 + + A method call to a wrapped module gets rerouted through the wrapper's `forward` method instead. + + """ + + def __call__( + self, wrapper_module: nn.Module, original_module: nn.Module, method: callable, *args: Any, **kwargs: Any + ): + """Reroutes a method call through the `wrapper_module`'s `forward` method. + + Args: + wrapper_module: The module that has `original_module` wrapped. + original_module: The module that was wrapped inside `wrapper_module`. + method_name: The name of the method that should be called on the `original_module` after inputs get + redirected through the `wrapper_module`'s `forward` method. + *args: The positional arguments to the method `method_name`. They will get passed to a patched + `forward` method instead. + **kwargs: The keyword arguments to the method `method_name`. They will get passed to a patched + `forward` method instead. + + """ + original_forward = original_module.forward + + def wrapped_forward(*_args: Any, **_kwargs: Any) -> Any: + # Unpatch ourselves immediately before calling the method `method_name` + # because itself may want to call the real `forward` + original_module.forward = original_forward # type: ignore[method-assign] + # Call the actual method e.g. `.training_step(...)` + out = method(*_args, **_kwargs) + self.on_after_inner_forward(wrapper_module, original_module) + return out + + # Patch the original_module's forward so we can redirect the arguments back to the real method + original_module.forward = wrapped_forward # type: ignore[method-assign] + + wrapper_output = wrapper_module(*args, **kwargs) + self.on_after_outer_forward(wrapper_module, original_module) + return wrapper_output + + def on_after_inner_forward(self, wrapper_module: nn.Module, original_module: nn.Module) -> None: + pass + + def on_after_outer_forward(self, wrapper_module: nn.Module, original_module: nn.Module) -> None: + pass diff --git a/lib/python3.12/site-packages/trl/rewards/__init__.py b/lib/python3.12/site-packages/trl/rewards/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..136d70af9cb5d95a855ce5bcedf5b5f9b526cb69 --- /dev/null +++ b/lib/python3.12/site-packages/trl/rewards/__init__.py @@ -0,0 +1,32 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import sys +from typing import TYPE_CHECKING + +from ..import_utils import _LazyModule + + +_import_structure = { + "format_rewards": ["think_format_reward"], +} + + +if TYPE_CHECKING: + from .format_rewards import think_format_reward + + +else: + sys.modules[__name__] = _LazyModule(__name__, __file__, _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/trl/rewards/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/trl/rewards/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4612023200ff49830542125800f97a6d2520bf91 Binary files /dev/null and b/lib/python3.12/site-packages/trl/rewards/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/rewards/__pycache__/format_rewards.cpython-312.pyc b/lib/python3.12/site-packages/trl/rewards/__pycache__/format_rewards.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..27548fb0ae6ef90c44319437e6dd1479aee23026 Binary files /dev/null and b/lib/python3.12/site-packages/trl/rewards/__pycache__/format_rewards.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/rewards/format_rewards.py b/lib/python3.12/site-packages/trl/rewards/format_rewards.py new file mode 100644 index 0000000000000000000000000000000000000000..4df6b0a3bcc539ab6cab270dfde2b1071e615260 --- /dev/null +++ b/lib/python3.12/site-packages/trl/rewards/format_rewards.py @@ -0,0 +1,50 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re + + +def think_format_reward(completions: list[list[dict[str, str]]], **kwargs) -> list[float]: + r""" + Reward function that checks if the reasoning process is enclosed within `""` and `""` tags. The + function returns a reward of 1.0 if the format is correct, otherwise 0.0. + + Args: + completions (`list[list[dict[str, str]]]`): + List of completions to be evaluated. Each completion must be a list of one message, i.e. a dictionary + containing the key `"content"` with the value being the text of the completion. + **kwargs: + Additional keyword arguments. This function does not use them, but they are required in the function + signature to ensure compatibility with trainers like [`GRPOTrainer`]. + + Returns: + `list[float]`: + A list of rewards, where each reward is 1.0 if the completion matches the expected format, otherwise 0.0. + + Example: + ```python + >>> from trl.rewards import think_format_reward + + >>> completions = [ + ... [{"content": "\nThis is my reasoning.\n\nThis is my answer."}], + ... [{"content": "\nThis is my reasoning.\nThis is my answer."}], + ... ] + >>> think_format_reward(completions) + [1.0, 0.0] + ``` + """ + pattern = r"^(?!.*)(.*?).*$" + completion_contents = [completion[0]["content"] for completion in completions] + matches = [re.match(pattern, content, re.DOTALL | re.MULTILINE) for content in completion_contents] + return [1.0 if match else 0.0 for match in matches] diff --git a/lib/python3.12/site-packages/trl/scripts/__init__.py b/lib/python3.12/site-packages/trl/scripts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..272a717db0e9a01d8c9ff23c21d44d6e804bc99c --- /dev/null +++ b/lib/python3.12/site-packages/trl/scripts/__init__.py @@ -0,0 +1,29 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ..import_utils import _LazyModule + + +_import_structure = { + "utils": ["init_zero_verbose", "ScriptArguments", "TrlParser"], +} + +if TYPE_CHECKING: + from .utils import ScriptArguments, TrlParser, init_zero_verbose +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/trl/scripts/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/trl/scripts/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77c82dd4f8d2818dd03ed250936ef485c606ad84 Binary files /dev/null and b/lib/python3.12/site-packages/trl/scripts/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/scripts/__pycache__/dpo.cpython-312.pyc b/lib/python3.12/site-packages/trl/scripts/__pycache__/dpo.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d33b7d1b6cbb48e5f2e4918e7c09704bea5c25d Binary files /dev/null and b/lib/python3.12/site-packages/trl/scripts/__pycache__/dpo.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/scripts/__pycache__/env.cpython-312.pyc b/lib/python3.12/site-packages/trl/scripts/__pycache__/env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d4c5c5324d63f746e96e104b2602152f482296f Binary files /dev/null and b/lib/python3.12/site-packages/trl/scripts/__pycache__/env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/scripts/__pycache__/grpo.cpython-312.pyc b/lib/python3.12/site-packages/trl/scripts/__pycache__/grpo.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79f54efc57dd83e12b7503f587fe9d1a8f25fe16 Binary files /dev/null and b/lib/python3.12/site-packages/trl/scripts/__pycache__/grpo.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/scripts/__pycache__/kto.cpython-312.pyc b/lib/python3.12/site-packages/trl/scripts/__pycache__/kto.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4033de6f3ee0435cfa353d0771145df050b6060 Binary files /dev/null and b/lib/python3.12/site-packages/trl/scripts/__pycache__/kto.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/scripts/__pycache__/sft.cpython-312.pyc b/lib/python3.12/site-packages/trl/scripts/__pycache__/sft.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..889e5a62aa81170de8476366f34a019c2acc72e2 Binary files /dev/null and b/lib/python3.12/site-packages/trl/scripts/__pycache__/sft.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/scripts/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/trl/scripts/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fc72b7da7c0e1135cea38bc8d8c42159f195479 Binary files /dev/null and b/lib/python3.12/site-packages/trl/scripts/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/scripts/__pycache__/vllm_serve.cpython-312.pyc b/lib/python3.12/site-packages/trl/scripts/__pycache__/vllm_serve.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f21bb134865e2cdb24c029581de8f9490eb2bf3 Binary files /dev/null and b/lib/python3.12/site-packages/trl/scripts/__pycache__/vllm_serve.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/scripts/dpo.py b/lib/python3.12/site-packages/trl/scripts/dpo.py new file mode 100644 index 0000000000000000000000000000000000000000..83fc9a43295eea26b493e39e8692adccdde6f588 --- /dev/null +++ b/lib/python3.12/site-packages/trl/scripts/dpo.py @@ -0,0 +1,161 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +# Full training +```bash +python trl/scripts/dpo.py \ + --dataset_name trl-lib/ultrafeedback_binarized \ + --dataset_streaming \ + --model_name_or_path Qwen/Qwen2-0.5B-Instruct \ + --learning_rate 5.0e-7 \ + --num_train_epochs 1 \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 8 \ + --gradient_checkpointing \ + --eval_strategy steps \ + --eval_steps 50 \ + --output_dir Qwen2-0.5B-DPO \ + --no_remove_unused_columns + --report_to wandb +``` + +# LoRA: +```bash +python trl/scripts/dpo.py \ + --dataset_name trl-lib/ultrafeedback_binarized \ + --dataset_streaming \ + --model_name_or_path Qwen/Qwen2-0.5B-Instruct \ + --learning_rate 5.0e-6 \ + --num_train_epochs 1 \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 8 \ + --gradient_checkpointing \ + --eval_strategy steps \ + --eval_steps 50 \ + --output_dir Qwen2-0.5B-DPO \ + --no_remove_unused_columns \ + --use_peft \ + --lora_r 32 \ + --lora_alpha 16 + --report_to wandb +``` +""" + +import argparse + +import torch +from datasets import load_dataset +from transformers import AutoModelForCausalLM, AutoTokenizer + +from trl import ( + DPOConfig, + DPOTrainer, + ModelConfig, + ScriptArguments, + TrlParser, + get_kbit_device_map, + get_peft_config, + get_quantization_config, +) +from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE + + +def main(script_args, training_args, model_args): + ################ + # Model & Tokenizer + ################### + torch_dtype = ( + model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype) + ) + quantization_config = get_quantization_config(model_args) + model_kwargs = dict( + revision=model_args.model_revision, + attn_implementation=model_args.attn_implementation, + torch_dtype=torch_dtype, + use_cache=False if training_args.gradient_checkpointing else True, + device_map=get_kbit_device_map() if quantization_config is not None else None, + quantization_config=quantization_config, + ) + model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code, **model_kwargs + ) + peft_config = get_peft_config(model_args) + if peft_config is None: + ref_model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code, **model_kwargs + ) + else: + ref_model = None + tokenizer = AutoTokenizer.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + if tokenizer.chat_template is None: + tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE + if script_args.ignore_bias_buffers: + # torch distributed hack + model._ddp_params_and_buffers_to_ignore = [ + name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool + ] + + ################ + # Dataset + ################ + dataset = load_dataset( + script_args.dataset_name, + name=script_args.dataset_config, + streaming=script_args.dataset_streaming, + ) + + ########## + # Training + ################ + trainer = DPOTrainer( + model, + ref_model, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + processing_class=tokenizer, + peft_config=peft_config, + ) + + trainer.train() + + if training_args.eval_strategy != "no": + metrics = trainer.evaluate() + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + + # Save and push to hub + trainer.save_model(training_args.output_dir) + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + + +def make_parser(subparsers: argparse._SubParsersAction = None): + dataclass_types = (ScriptArguments, DPOConfig, ModelConfig) + if subparsers is not None: + parser = subparsers.add_parser("dpo", help="Run the DPO training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + script_args, training_args, model_args = parser.parse_args_and_config() + main(script_args, training_args, model_args) diff --git a/lib/python3.12/site-packages/trl/scripts/env.py b/lib/python3.12/site-packages/trl/scripts/env.py new file mode 100644 index 0000000000000000000000000000000000000000..7eb81b4544e352065a7c45a1a41c5587c809a0a4 --- /dev/null +++ b/lib/python3.12/site-packages/trl/scripts/env.py @@ -0,0 +1,84 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import platform +from importlib.metadata import version + +import torch +from accelerate.commands.config import default_config_file, load_config_from_file +from transformers import is_bitsandbytes_available +from transformers.utils import is_openai_available, is_peft_available + +from .. import __version__ +from ..import_utils import ( + is_deepspeed_available, + is_diffusers_available, + is_liger_kernel_available, + is_llm_blender_available, + is_vllm_available, +) +from .utils import get_git_commit_hash + + +def print_env(): + devices = None + if torch.cuda.is_available(): + devices = [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())] + elif torch.backends.mps.is_available(): + devices = ["MPS"] + elif torch.xpu.is_available(): + devices = [torch.xpu.get_device_name(i) for i in range(torch.xpu.device_count())] + + accelerate_config = accelerate_config_str = "not found" + + # Get the default from the config file. + if os.path.isfile(default_config_file): + accelerate_config = load_config_from_file(default_config_file).to_dict() + + accelerate_config_str = ( + "\n" + "\n".join([f" - {prop}: {val}" for prop, val in accelerate_config.items()]) + if isinstance(accelerate_config, dict) + else accelerate_config + ) + + commit_hash = get_git_commit_hash("trl") + + info = { + "Platform": platform.platform(), + "Python version": platform.python_version(), + "TRL version": f"{__version__}+{commit_hash[:7]}" if commit_hash else __version__, + "PyTorch version": version("torch"), + "accelerator(s)": ", ".join(devices) if devices is not None else "cpu", + "Transformers version": version("transformers"), + "Accelerate version": version("accelerate"), + "Accelerate config": accelerate_config_str, + "Datasets version": version("datasets"), + "HF Hub version": version("huggingface_hub"), + "bitsandbytes version": version("bitsandbytes") if is_bitsandbytes_available() else "not installed", + "DeepSpeed version": version("deepspeed") if is_deepspeed_available() else "not installed", + "Diffusers version": version("diffusers") if is_diffusers_available() else "not installed", + "Liger-Kernel version": version("liger_kernel") if is_liger_kernel_available() else "not installed", + "LLM-Blender version": version("llm_blender") if is_llm_blender_available() else "not installed", + "OpenAI version": version("openai") if is_openai_available() else "not installed", + "PEFT version": version("peft") if is_peft_available() else "not installed", + "vLLM version": version("vllm") if is_vllm_available() else "not installed", + } + + info_str = "\n".join([f"- {prop}: {val}" for prop, val in info.items()]) + print(f"\nCopy-paste the following information when reporting an issue:\n\n{info_str}\n") # noqa + + +if __name__ == "__main__": + print_env() diff --git a/lib/python3.12/site-packages/trl/scripts/grpo.py b/lib/python3.12/site-packages/trl/scripts/grpo.py new file mode 100644 index 0000000000000000000000000000000000000000..a606d9e3f3ff40f15b57f39cea77b25cd9094281 --- /dev/null +++ b/lib/python3.12/site-packages/trl/scripts/grpo.py @@ -0,0 +1,132 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import importlib +import os +import sys +from dataclasses import dataclass, field +from typing import Optional + +from datasets import load_dataset +from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer + +from trl import GRPOConfig, GRPOTrainer, ModelConfig, ScriptArguments, TrlParser, get_peft_config +from trl.rewards import think_format_reward + + +reward_funcs_registry = { + "think_format_reward": think_format_reward, +} + + +@dataclass +class GRPOScriptArguments(ScriptArguments): + """ + Script arguments for the GRPO training script. + + Args: + reward_model_name_or_path (`str` or `None`, *optional*, defaults to `None`): + Reward model id of a pretrained model hosted inside a model repo on huggingface.co or local path to a + directory containing model weights saved using [`~transformers.PreTrainedModel.save_pretrained`]. + reward_funcs (`list[str]` or `None`, *optional*, defaults to `None`): + Reward functions to use. It can be either one of `"think_format_reward"`; or a dotted import path " (e.g., + `'my_lib.rewards.custom_reward'`). + """ + + reward_model_name_or_path: Optional[str] = field( + default=None, + metadata={ + "help": "Reward model id of a pretrained model hosted inside a model repo on huggingface.co or " + "local path to a directory containing model weights saved using `PreTrainedModel.save_pretrained`." + }, + ) + reward_funcs: Optional[list[str]] = field( + default=None, + metadata={ + "help": "Reward functions to use. It can be either one of 'think_format_reward'; or a dotted " + "import path. (e.g., 'my_lib.rewards.custom_reward')." + }, + ) + + +def main(script_args, training_args, model_args): + # Load a pretrained model + model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + tokenizer = AutoTokenizer.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + + # Get the reward models and functions + reward_funcs = [] + if script_args.reward_model_name_or_path: + reward_model = AutoModelForSequenceClassification.from_pretrained( + script_args.reward_model_name_or_path, trust_remote_code=model_args.trust_remote_code, num_labels=1 + ) + reward_funcs.append(reward_model) + + if script_args.reward_funcs: + for func_name in script_args.reward_funcs: + if func_name in reward_funcs_registry: + reward_funcs.append(reward_funcs_registry[func_name]) + elif "." in func_name: + module_path, func_name = func_name.rsplit(".", 1) + sys.path.insert(0, os.getcwd()) + module = importlib.import_module(module_path) + reward_func = getattr(module, func_name) + reward_funcs.append(reward_func) + else: + raise ValueError( + f"Could not load reward function '{func_name}'. Expected one of " + f"{list(reward_funcs_registry.keys())} or a valid import path." + ) + + # Load the dataset + dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config) + + # Initialize the GRPO trainer + trainer = GRPOTrainer( + model=model, + reward_funcs=reward_model, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + processing_class=tokenizer, + peft_config=get_peft_config(model_args), + ) + + # Train and push the model to the Hub + trainer.train() + + # Save and push to hub + trainer.save_model(training_args.output_dir) + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + + +def make_parser(subparsers: argparse._SubParsersAction = None): + dataclass_types = (GRPOScriptArguments, GRPOConfig, ModelConfig) + if subparsers is not None: + parser = subparsers.add_parser("grpo", help="Run the GRPO training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + script_args, training_args, model_args = parser.parse_args_and_config() + main(script_args, training_args, model_args) diff --git a/lib/python3.12/site-packages/trl/scripts/kto.py b/lib/python3.12/site-packages/trl/scripts/kto.py new file mode 100644 index 0000000000000000000000000000000000000000..0058621f048a717d71f5bf98b62244b5f55473ec --- /dev/null +++ b/lib/python3.12/site-packages/trl/scripts/kto.py @@ -0,0 +1,132 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Run the KTO training script with the commands below. In general, the optimal configuration for KTO will be similar to +that of DPO. + +# Full training: +```bash +python trl/scripts/kto.py \ + --dataset_name trl-lib/kto-mix-14k \ + --model_name_or_path=trl-lib/qwen1.5-1.8b-sft \ + --per_device_train_batch_size 16 \ + --num_train_epochs 1 \ + --learning_rate 5e-7 \ + --lr_scheduler_type=cosine \ + --gradient_accumulation_steps 1 \ + --eval_steps 500 \ + --output_dir=kto-aligned-model \ + --warmup_ratio 0.1 \ + --report_to wandb \ + --bf16 \ + --logging_first_step +``` + +# QLoRA: +```bash +# QLoRA: +python trl/scripts/kto.py \ + --dataset_name trl-lib/kto-mix-14k \ + --model_name_or_path=trl-lib/qwen1.5-1.8b-sft \ + --per_device_train_batch_size 8 \ + --num_train_epochs 1 \ + --learning_rate 5e-7 \ + --lr_scheduler_type=cosine \ + --gradient_accumulation_steps 1 \ + --eval_steps 500 \ + --output_dir=kto-aligned-model-lora \ + --warmup_ratio 0.1 \ + --report_to wandb \ + --bf16 \ + --logging_first_step \ + --use_peft \ + --load_in_4bit \ + --lora_target_modules=all-linear \ + --lora_r=16 \ + --lora_alpha=16 +``` +""" + +import argparse + +from datasets import load_dataset +from transformers import AutoModelForCausalLM, AutoTokenizer + +from trl import ( + KTOConfig, + KTOTrainer, + ModelConfig, + ScriptArguments, + TrlParser, + get_peft_config, + setup_chat_format, +) + + +def main(script_args, training_args, model_args): + # Load a pretrained model + model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + ref_model = AutoModelForCausalLM.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + + tokenizer = AutoTokenizer.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # If we are aligning a base model, we use ChatML as the default template + if tokenizer.chat_template is None: + model, tokenizer = setup_chat_format(model, tokenizer) + + # Load the dataset + dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config) + + # Initialize the KTO trainer + trainer = KTOTrainer( + model, + ref_model, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + processing_class=tokenizer, + peft_config=get_peft_config(model_args), + ) + + # Train and push the model to the Hub + trainer.train() + + # Save and push to hub + trainer.save_model(training_args.output_dir) + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + + +def make_parser(subparsers: argparse._SubParsersAction = None): + dataclass_types = (ScriptArguments, KTOConfig, ModelConfig) + if subparsers is not None: + parser = subparsers.add_parser("kto", help="Run the KTO training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + script_args, training_args, model_args = parser.parse_args_and_config() + main(script_args, training_args, model_args) diff --git a/lib/python3.12/site-packages/trl/scripts/sft.py b/lib/python3.12/site-packages/trl/scripts/sft.py new file mode 100644 index 0000000000000000000000000000000000000000..7ab6c5099c996e98d6483b3cc93cf34c8a665270 --- /dev/null +++ b/lib/python3.12/site-packages/trl/scripts/sft.py @@ -0,0 +1,152 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +# Full training +``` +python trl/scripts/sft.py \ + --model_name_or_path Qwen/Qwen2-0.5B \ + --dataset_name trl-lib/Capybara \ + --learning_rate 2.0e-5 \ + --num_train_epochs 1 \ + --packing \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 8 \ + --gradient_checkpointing \ + --eos_token '<|im_end|>' \ + --eval_strategy steps \ + --eval_steps 100 \ + --output_dir Qwen2-0.5B-SFT \ + --push_to_hub +``` + +# LoRA +``` +python trl/scripts/sft.py \ + --model_name_or_path Qwen/Qwen2-0.5B \ + --dataset_name trl-lib/Capybara \ + --learning_rate 2.0e-4 \ + --num_train_epochs 1 \ + --packing \ + --per_device_train_batch_size 2 \ + --gradient_accumulation_steps 8 \ + --gradient_checkpointing \ + --eos_token '<|im_end|>' \ + --eval_strategy steps \ + --eval_steps 100 \ + --use_peft \ + --lora_r 32 \ + --lora_alpha 16 \ + --output_dir Qwen2-0.5B-SFT \ + --push_to_hub +``` +""" + +import argparse + +from datasets import load_dataset +from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer +from transformers.models.auto.modeling_auto import MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES + +from trl import ( + ModelConfig, + ScriptArguments, + SFTConfig, + SFTTrainer, + TrlParser, + clone_chat_template, + get_kbit_device_map, + get_peft_config, + get_quantization_config, +) + + +def main(script_args, training_args, model_args): + ################ + # Model init kwargs & Tokenizer + ################ + quantization_config = get_quantization_config(model_args) + model_kwargs = dict( + revision=model_args.model_revision, + trust_remote_code=model_args.trust_remote_code, + attn_implementation=model_args.attn_implementation, + torch_dtype=model_args.torch_dtype, + use_cache=False if training_args.gradient_checkpointing else True, + device_map=get_kbit_device_map() if quantization_config is not None else None, + quantization_config=quantization_config, + ) + + # Create model + config = AutoConfig.from_pretrained(model_args.model_name_or_path) + valid_image_text_architectures = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.values() + + if config.architectures and any(arch in valid_image_text_architectures for arch in config.architectures): + from transformers import AutoModelForImageTextToText + + model_kwargs.pop("use_cache", None) # Image models do not support cache + model = AutoModelForImageTextToText.from_pretrained(model_args.model_name_or_path, **model_kwargs) + else: + model = AutoModelForCausalLM.from_pretrained(model_args.model_name_or_path, **model_kwargs) + + # Create tokenizer + tokenizer = AutoTokenizer.from_pretrained( + model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code, use_fast=True + ) + + # Set default chat template if needed + if tokenizer.chat_template is None: + # TODO: source should be passed as an argument + model, tokenizer = clone_chat_template(model, tokenizer, "Qwen/Qwen3-0.6B") + + ################ + # Dataset + ################ + dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config) + + ################ + # Training + ################ + trainer = SFTTrainer( + model=model, + args=training_args, + train_dataset=dataset[script_args.dataset_train_split], + eval_dataset=dataset[script_args.dataset_test_split] if training_args.eval_strategy != "no" else None, + processing_class=tokenizer, + peft_config=get_peft_config(model_args), + ) + + trainer.train() + + # Save and push to hub + trainer.save_model(training_args.output_dir) + if training_args.push_to_hub: + trainer.push_to_hub(dataset_name=script_args.dataset_name) + + +def make_parser(subparsers: argparse._SubParsersAction = None): + dataclass_types = (ScriptArguments, SFTConfig, ModelConfig) + if subparsers is not None: + parser = subparsers.add_parser("sft", help="Run the SFT training script", dataclass_types=dataclass_types) + else: + parser = TrlParser(dataclass_types) + return parser + + +if __name__ == "__main__": + parser = make_parser() + # When using the trl cli, this script may be run with additional arguments, corresponding accelerate arguments. + # To ensure that their parsing does not interfere with the script arguments, parse the arguments with + # `return_remaining_strings=True`, then ignore the remaining strings. + script_args, training_args, model_args, _ = parser.parse_args_and_config(return_remaining_strings=True) + main(script_args, training_args, model_args) diff --git a/lib/python3.12/site-packages/trl/scripts/utils.py b/lib/python3.12/site-packages/trl/scripts/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..030719ab1e4587f76232085df867597d0730215d --- /dev/null +++ b/lib/python3.12/site-packages/trl/scripts/utils.py @@ -0,0 +1,284 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import importlib +import inspect +import logging +import os +import subprocess +import sys +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Optional, Union + +import yaml +from transformers import HfArgumentParser +from transformers.hf_argparser import DataClass, DataClassType +from transformers.utils import is_rich_available + + +logger = logging.getLogger(__name__) + + +@dataclass +class ScriptArguments: + """ + Arguments common to all scripts. + + Args: + dataset_name (`str`): + Dataset name. + dataset_config (`str` or `None`, *optional*, defaults to `None`): + Dataset configuration name. Corresponds to the `name` argument of the [`~datasets.load_dataset`] function. + dataset_train_split (`str`, *optional*, defaults to `"train"`): + Dataset split to use for training. + dataset_test_split (`str`, *optional*, defaults to `"test"`): + Dataset split to use for evaluation. + dataset_streaming (`bool`, *optional*, defaults to `False`): + Whether to stream the dataset. If True, the dataset will be loaded in streaming mode. + gradient_checkpointing_use_reentrant (`bool`, *optional*, defaults to `False`): + Whether to apply `use_reentrant` for gradient checkpointing. + ignore_bias_buffers (`bool`, *optional*, defaults to `False`): + Debug argument for distributed training. Fix for DDP issues with LM bias/mask buffers - invalid scalar + type, inplace operation. See + https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992. + """ + + dataset_name: Optional[str] = field(default=None, metadata={"help": "Dataset name."}) + dataset_config: Optional[str] = field( + default=None, + metadata={ + "help": "Dataset configuration name. Corresponds to the `name` argument of the `datasets.load_dataset` " + "function." + }, + ) + dataset_train_split: str = field(default="train", metadata={"help": "Dataset split to use for training."}) + dataset_test_split: str = field(default="test", metadata={"help": "Dataset split to use for evaluation."}) + dataset_streaming: bool = field( + default=False, + metadata={"help": "Whether to stream the dataset. If True, the dataset will be loaded in streaming mode."}, + ) + gradient_checkpointing_use_reentrant: bool = field( + default=False, + metadata={"help": "Whether to apply `use_reentrant` for gradient checkpointing."}, + ) + ignore_bias_buffers: bool = field( + default=False, + metadata={ + "help": "Debug argument for distributed training. Fix for DDP issues with LM bias/mask buffers - invalid " + "scalar type, inplace operation. See " + "https://github.com/huggingface/transformers/issues/22482#issuecomment-1595790992." + }, + ) + + +def init_zero_verbose(): + """ + Perform zero verbose init - use this method on top of the CLI modules to make logging and warning output cleaner. + Uses Rich if available, falls back otherwise. + """ + import logging + import warnings + + FORMAT = "%(message)s" + + if is_rich_available(): + from rich.logging import RichHandler + + handler = RichHandler() + else: + handler = logging.StreamHandler() + + logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[handler], level=logging.ERROR) + + # Custom warning handler to redirect warnings to the logging system + def warning_handler(message, category, filename, lineno, file=None, line=None): + logging.warning(f"{filename}:{lineno}: {category.__name__}: {message}") + + # Add the custom warning handler - we need to do that before importing anything to make sure the loggers work well + warnings.showwarning = warning_handler + + +class TrlParser(HfArgumentParser): + """ + A subclass of [`transformers.HfArgumentParser`] designed for parsing command-line arguments with dataclass-backed + configurations, while also supporting configuration file loading and environment variable management. + + Args: + dataclass_types (`Union[DataClassType, Iterable[DataClassType]]` or `None`, *optional*, defaults to `None`): + Dataclass types to use for argument parsing. + **kwargs: + Additional keyword arguments passed to the [`transformers.HfArgumentParser`] constructor. + + Examples: + + ```yaml + # config.yaml + env: + VAR1: value1 + arg1: 23 + ``` + + ```python + # main.py + import os + from dataclasses import dataclass + from trl import TrlParser + + + @dataclass + class MyArguments: + arg1: int + arg2: str = "alpha" + + + parser = TrlParser(dataclass_types=[MyArguments]) + training_args = parser.parse_args_and_config() + + print(training_args, os.environ.get("VAR1")) + ``` + + ```bash + $ python main.py --config config.yaml + (MyArguments(arg1=23, arg2='alpha'),) value1 + + $ python main.py --arg1 5 --arg2 beta + (MyArguments(arg1=5, arg2='beta'),) None + ``` + """ + + def __init__( + self, + dataclass_types: Optional[Union[DataClassType, Iterable[DataClassType]]] = None, + **kwargs, + ): + # Make sure dataclass_types is an iterable + if dataclass_types is None: + dataclass_types = [] + elif not isinstance(dataclass_types, Iterable): + dataclass_types = [dataclass_types] + + # Check that none of the dataclasses have the "config" field + for dataclass_type in dataclass_types: + if "config" in dataclass_type.__dataclass_fields__: + raise ValueError( + f"Dataclass {dataclass_type.__name__} has a field named 'config'. This field is reserved for the " + f"config file path and should not be used in the dataclass." + ) + + super().__init__(dataclass_types=dataclass_types, **kwargs) + + def parse_args_and_config( + self, + args: Optional[Iterable[str]] = None, + return_remaining_strings: bool = False, + fail_with_unknown_args: bool = True, + ) -> tuple[DataClass, ...]: + """ + Parse command-line args and config file into instances of the specified dataclass types. + + This method wraps [`transformers.HfArgumentParser.parse_args_into_dataclasses`] and also parses the config file + specified with the `--config` flag. The config file (in YAML format) provides argument values that replace the + default values in the dataclasses. Command line arguments can override values set by the config file. The + method also sets any environment variables specified in the `env` field of the config file. + """ + args = list(args) if args is not None else sys.argv[1:] + if "--config" in args: + # Get the config file path from + config_index = args.index("--config") + args.pop(config_index) # remove the --config flag + config_path = args.pop(config_index) # get the path to the config file + with open(config_path) as yaml_file: + config = yaml.safe_load(yaml_file) + + # Set the environment variables specified in the config file + if "env" in config: + env_vars = config.pop("env", {}) + if not isinstance(env_vars, dict): + raise ValueError("`env` field should be a dict in the YAML file.") + for key, value in env_vars.items(): + os.environ[key] = str(value) + + # Set the defaults from the config values + config_remaining_strings = self.set_defaults_with_config(**config) + else: + config_remaining_strings = [] + + # Parse the arguments from the command line + output = self.parse_args_into_dataclasses(args=args, return_remaining_strings=return_remaining_strings) + + # Merge remaining strings from the config file with the remaining strings from the command line + if return_remaining_strings: + args_remaining_strings = output[-1] + return output[:-1] + (config_remaining_strings + args_remaining_strings,) + elif fail_with_unknown_args and config_remaining_strings: + raise ValueError( + f"Unknown arguments from config file: {config_remaining_strings}. Please remove them, add them to the " + "dataclass, or set `fail_with_unknown_args=False`." + ) + else: + return output + + def set_defaults_with_config(self, **kwargs) -> list[str]: + """ + Overrides the parser's default values with those provided via keyword arguments, including for subparsers. + + Any argument with an updated default will also be marked as not required if it was previously required. + + Returns a list of strings that were not consumed by the parser. + """ + + def apply_defaults(parser, kw): + used_keys = set() + for action in parser._actions: + # Handle subparsers recursively + if isinstance(action, argparse._SubParsersAction): + for subparser in action.choices.values(): + used_keys.update(apply_defaults(subparser, kw)) + elif action.dest in kw: + action.default = kw[action.dest] + action.required = False + used_keys.add(action.dest) + return used_keys + + used_keys = apply_defaults(self, kwargs) + # Remaining args not consumed by the parser + remaining = [ + item for key, value in kwargs.items() if key not in used_keys for item in (f"--{key}", str(value)) + ] + return remaining + + +def get_git_commit_hash(package_name): + try: + # Import the package to locate its path + package = importlib.import_module(package_name) + # Get the path to the package using inspect + package_path = os.path.dirname(inspect.getfile(package)) + + # Navigate up to the Git repository root if the package is inside a subdirectory + git_repo_path = os.path.abspath(os.path.join(package_path, "..")) + git_dir = os.path.join(git_repo_path, ".git") + + if os.path.isdir(git_dir): + # Run the git command to get the current commit hash + commit_hash = ( + subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=git_repo_path).strip().decode("utf-8") + ) + return commit_hash + else: + return None + except Exception as e: + return f"Error: {str(e)}" diff --git a/lib/python3.12/site-packages/trl/scripts/vllm_serve.py b/lib/python3.12/site-packages/trl/scripts/vllm_serve.py new file mode 100644 index 0000000000000000000000000000000000000000..eb4a6e269ca000e929f7bdc4a4d42859bf99160f --- /dev/null +++ b/lib/python3.12/site-packages/trl/scripts/vllm_serve.py @@ -0,0 +1,610 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import logging +import os +from collections.abc import Sequence +from contextlib import asynccontextmanager +from dataclasses import dataclass, field +from itertools import chain +from multiprocessing import Pipe, Process +from multiprocessing.connection import Connection +from typing import Optional + +import torch + +from trl import TrlParser +from trl.import_utils import ( + is_fastapi_available, + is_pydantic_available, + is_uvicorn_available, + is_vllm_ascend_available, + is_vllm_available, +) + + +if is_fastapi_available(): + from fastapi import FastAPI + + +if is_pydantic_available(): + from pydantic import BaseModel + + +if is_uvicorn_available(): + import uvicorn + + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.parallel_state import get_world_group + from vllm.distributed.utils import StatelessProcessGroup + from vllm.sampling_params import GuidedDecodingParams + from vllm.utils import get_open_port + + if is_vllm_ascend_available(): + from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator + + +logger = logging.getLogger(__name__) + +# We use CUDA with multiprocessing, so we must use the 'spawn' start method. Otherwise, we will get the following +# error: RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use +# the 'spawn' start method +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + + +class WeightSyncWorkerExtension: + """ + A vLLM worker extension that enables weight synchronization between a client and multiple server workers. + + This worker uses a `StatelessProcessGroup` to establish communication and a `PyNcclCommunicator` to handle + efficient GPU-based communication using NCCL. The primary purpose of this class is to receive updated model weights + from a client process and distribute them to all worker processes participating in model inference. + """ + + # The following attributes are initialized when `init_communicator` method is called. + pynccl_comm = None # Communicator for weight updates + client_rank = None # Source rank for broadcasting updated weights + + def init_communicator(self, host: str, port: int, world_size: int) -> None: + """ + Initializes the weight update communicator using a stateless process group. + + This method creates a `StatelessProcessGroup` that allows external training processes to communicate with vLLM + workers without interfering with the global torch distributed group. + + Args: + host (`str`): + Hostname or IP address of the master node. + port (`int`): + Port number to be used for communication. + world_size (`int`): + Total number of participating processes in the update group. + """ + if self.pynccl_comm is not None: + raise RuntimeError("Weight update group already initialized. Call close_communicator first.") + + # Get the rank of the current worker in the global world group. + rank = get_world_group().rank + + # Create a stateless process group to manage communication between training processes and vLLM workers. + pg = StatelessProcessGroup.create(host=host, port=port, rank=rank, world_size=world_size) + + # Initialize the NCCL-based communicator for weight synchronization. + self.pynccl_comm = PyNcclCommunicator(pg, device=self.device) + + # The client process that sends updated weights has the highest rank (world_size - 1). + self.client_rank = world_size - 1 + + def update_named_param(self, name: str, dtype: torch.dtype, shape: Sequence[int]) -> None: + """ + Receives updated weights from the client process and updates the named parameter in the model. + + Args: + name (`str`): + Name of the weight tensor being updated. + dtype (`torch.dtype`): + Data type of the weight tensor (e.g., `torch.float32`). + shape (`Sequence[int]`): + Shape of the weight tensor. + """ + if self.pynccl_comm is None: + raise RuntimeError("Communicator not initialized. Call `init_communicator` first.") + + # Allocate memory for the incoming weight tensor on the correct device. + weight = torch.empty(shape, dtype=dtype, device=self.device) + + # Use NCCL to broadcast the updated weights from the client (src) to all workers. + self.pynccl_comm.broadcast(weight, src=self.client_rank) + self.pynccl_comm.group.barrier() + + # Load the received weights into the model. + self.model_runner.model.load_weights(weights=[(name, weight)]) + + def close_communicator(self) -> None: + """ + Closes the communicator when weight synchronization is no longer needed. + + This method deletes the NCCL communicator to release associated resources. + """ + + if self.pynccl_comm is not None: + del self.pynccl_comm + self.pynccl_comm = None # Ensure attribute is reset to None + self.client_rank = None # Ensure attribute is reset to None + + +@dataclass +class ScriptArguments: + r""" + Arguments for the script. + + Args: + model (`str`): + Model name or path to load the model from. + revision (`str` or `None`, *optional*, defaults to `None`): + Revision to use for the model. If not specified, the default branch will be used. + tensor_parallel_size (`int`, *optional*, defaults to `1`): + Number of tensor parallel workers to use. + data_parallel_size (`int`, *optional*, defaults to `1`): + Number of data parallel workers to use. + host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host address to run the server on. + port (`int`, *optional*, defaults to `8000`): + Port to run the server on. + gpu_memory_utilization (`float`, *optional*, defaults to `0.9`): + Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV cache on the + device dedicated to generation powered by vLLM. Higher values will increase the KV cache size and thus + improve the model's throughput. However, if the value is too high, it may cause out-of-memory (OOM) errors + during initialization. + dtype (`str`, *optional*, defaults to `"auto"`): + Data type to use for vLLM generation. If set to `"auto"`, the data type will be automatically determined + based on the model configuration. Find the supported values in the vLLM documentation. + max_model_len (`int` or `None`, *optional*, defaults to `None`): + If set, the `max_model_len` to use for vLLM. This can be useful when running with reduced + `vllm_gpu_memory_utilization`, leading to a reduced KV cache size. If not set, vLLM will use the model + context size, which might be much larger than the KV cache, leading to inefficiencies. + enable_prefix_caching (`bool` or `None`, *optional*, defaults to `None`): + Whether to enable prefix caching in vLLM. If set to `True`, ensure that the model and the hardware support + this feature. + enforce_eager (`bool`, *optional*, defaults to `False`): + Whether to enforce eager execution. If set to `True`, we will disable CUDA graph and always execute the + model in eager mode. If `False` (default behavior), we will use CUDA graph and eager execution in hybrid. + kv_cache_dtype (`str`, *optional*, defaults to `"auto"`): + Data type to use for KV cache. If set to `"auto"`, the dtype will default to the model data type. + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to trust remote code when loading models. Set to `True` to allow executing code from model + repositories. This is required for some custom models but introduces security risks. + log_level (`str`, *optional*, defaults to `"info"`): + Log level for uvicorn. Possible choices: `"critical"`, `"error"`, `"warning"`, `"info"`, `"debug"`, + `"trace"`. + """ + + model: str = field( + metadata={"help": "Model name or path to load the model from."}, + ) + revision: Optional[str] = field( + default=None, + metadata={"help": "Revision to use for the model. If not specified, the default branch will be used."}, + ) + tensor_parallel_size: int = field( + default=1, + metadata={"help": "Number of tensor parallel workers to use."}, + ) + data_parallel_size: int = field( + default=1, + metadata={"help": "Number of data parallel workers to use."}, + ) + host: str = field( + default="0.0.0.0", + metadata={"help": "Host address to run the server on."}, + ) + port: int = field( + default=8000, + metadata={"help": "Port to run the server on."}, + ) + gpu_memory_utilization: float = field( + default=0.9, + metadata={ + "help": "Ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV " + "cache on the device dedicated to generation powered by vLLM. Higher values will increase the KV cache " + "size and thus improve the model's throughput. However, if the value is too high, it may cause " + "out-of-memory (OOM) errors during initialization." + }, + ) + dtype: str = field( + default="auto", + metadata={ + "help": "Data type to use for vLLM generation. If set to 'auto', the data type will be automatically " + "determined based on the model configuration. Find the supported values in the vLLM documentation." + }, + ) + max_model_len: Optional[int] = field( + default=None, + metadata={ + "help": "If set, the `max_model_len` to use for vLLM. This can be useful when running with reduced " + "`vllm_gpu_memory_utilization`, leading to a reduced KV cache size. If not set, vLLM will use the model " + "context size, which might be much larger than the KV cache, leading to inefficiencies." + }, + ) + enable_prefix_caching: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to enable prefix caching in vLLM. If set to `True`, ensure that the model and the " + "hardware support this feature." + }, + ) + enforce_eager: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether to enforce eager execution. If set to `True`, we will disable CUDA graph and always " + "execute the model in eager mode. If `False` (default behavior), we will use CUDA graph and eager " + "execution in hybrid." + }, + ) + kv_cache_dtype: str = field( + default="auto", + metadata={ + "help": "Data type to use for KV cache. If set to 'auto', the dtype will default to the model data type." + }, + ) + trust_remote_code: bool = field( + default=False, + metadata={ + "help": "Whether to trust remote code when loading models. Set to True to allow executing code from model " + "repositories. This is required for some custom models but introduces security risks." + }, + ) + log_level: str = field( + default="info", + metadata={ + "help": "Log level for uvicorn. Possible choices: 'critical', 'error', 'warning', 'info', 'debug', " + "'trace'." + }, + ) + + +def llm_worker( + script_args: ScriptArguments, data_parallel_rank: int, master_port: int, connection: Connection +) -> None: + # Set required environment variables for DP to work with vLLM + os.environ["VLLM_DP_RANK"] = str(data_parallel_rank) + os.environ["VLLM_DP_RANK_LOCAL"] = str(data_parallel_rank) + os.environ["VLLM_DP_SIZE"] = str(script_args.data_parallel_size) + os.environ["VLLM_DP_MASTER_PORT"] = str(master_port) + + llm = LLM( + model=script_args.model, + revision=script_args.revision, + tensor_parallel_size=script_args.tensor_parallel_size, + gpu_memory_utilization=script_args.gpu_memory_utilization, + enforce_eager=script_args.enforce_eager, + dtype=script_args.dtype, + # Automatic Prefix Caching caches the KV cache of existing queries, so that a new query can + # directly reuse the KV cache if it shares the same prefix with one of the existing queries. + # This is particularly useful here because we generate completions from the same prompts. + enable_prefix_caching=script_args.enable_prefix_caching, + kv_cache_dtype=script_args.kv_cache_dtype, + max_model_len=script_args.max_model_len, + worker_extension_cls="trl.scripts.vllm_serve.WeightSyncWorkerExtension", + trust_remote_code=script_args.trust_remote_code, + ) + + # Send ready signal to parent process + connection.send({"status": "ready"}) + + while True: + # Wait for commands from the parent process + try: + command = connection.recv() + except KeyboardInterrupt: + llm.collective_rpc(method="close_communicator") + break + + # Handle commands + if command["type"] in ["call", "fire_and_forget"]: + method_name = command["method"] + args, kwargs = command.get("args", ()), command.get("kwargs", {}) + method = getattr(llm, method_name) + result = method(*args, **kwargs) + if command["type"] == "call": + connection.send(result) + elif command["type"] == "shutdown": + break + + +def chunk_list(lst: list, n: int) -> list[list]: + """ + Split list `lst` into `n` evenly distributed sublists. + + Example: + ```python + >>> chunk_list([1, 2, 3, 4, 5, 6], 2) + [[1, 2, 3], [4, 5, 6]] + + >>> chunk_list([1, 2, 3, 4, 5, 6], 4) + [[1, 2], [3, 4], [5], [6]] + + >>> chunk_list([1, 2, 3, 4, 5, 6], 8) + [[1], [2], [3], [4], [5], [6], [], []] + ``` + """ + k, r = divmod(len(lst), n) + return [lst[i * k + min(i, r) : (i + 1) * k + min(i + 1, r)] for i in range(n)] + + +def main(script_args: ScriptArguments): + if not is_fastapi_available(): + raise ImportError( + "FastAPI is required to run the vLLM serve script. Please install it using `pip install fastapi`." + ) + + if not is_pydantic_available(): + raise ImportError( + "Pydantic is required to run the vLLM serve script. Please install it using `pip install pydantic`." + ) + + if not is_uvicorn_available(): + raise ImportError( + "Uvicorn is required to run the vLLM serve script. Please install it using `pip install uvicorn`." + ) + + if not is_vllm_available(): + raise ImportError("vLLM is required to run the vLLM serve script. Please install it using `pip install vllm`.") + + # Spawn dp workers, and setup pipes for communication + master_port = get_open_port() + connections = [] + processes = [] + for data_parallel_rank in range(script_args.data_parallel_size): + parent_connection, child_connection = Pipe() + process = Process(target=llm_worker, args=(script_args, data_parallel_rank, master_port, child_connection)) + process.start() + connections.append(parent_connection) + processes.append(process) + + @asynccontextmanager + async def lifespan(app: FastAPI): + # Wait for all workers to send "ready" + ready_connections = set() + while len(ready_connections) < script_args.data_parallel_size: + for connection in connections: + msg = connection.recv() + if isinstance(msg, dict) and msg.get("status") == "ready": + ready_connections.add(connection) + + yield + + # Wait for processes to terminate + for process in processes: + process.join(timeout=10) # Wait for 10 seconds for the process to terminate + if process.is_alive(): + logger.warning(f"Process {process} is still alive after 10 seconds, attempting to terminate...") + process.terminate() + process.join() # ensure process termination after calling terminate() + + app = FastAPI(lifespan=lifespan) + + # Define the endpoints for the model server + @app.get("/health/") + async def health(): + """ + Health check endpoint to verify that the server is running. + """ + return {"status": "ok"} + + @app.get("/get_world_size/") + async def get_world_size(): + """ + Retrieves the world size of the LLM engine, which is `tensor_parallel_size * data_parallel_size`. + + Returns: + `dict`: + A dictionary containing the world size. + + Example response: + ```json + {"world_size": 8} + ``` + """ + return {"world_size": script_args.tensor_parallel_size * script_args.data_parallel_size} + + class GenerateRequest(BaseModel): + prompts: list[str] + n: int = 1 + repetition_penalty: float = 1.0 + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 + min_p: float = 0.0 + max_tokens: int = 16 + guided_decoding_regex: Optional[str] = None + generation_kwargs: dict = field(default_factory=dict) + + class GenerateResponse(BaseModel): + completion_ids: list[list[int]] + + @app.post("/generate/", response_model=GenerateResponse) + async def generate(request: GenerateRequest): + """ + Generates completions for the provided prompts. + + Args: + request (`GenerateRequest`): + - `prompts` (list of `str`): A list of prompts (text strings) for the model to generate completions. + - `n` (`int`, *optional*, defaults to `1`): Number of completions to generate for each prompt. + - `repetition_penalty` (`float`, *optional*, defaults to `1.0`): Repetition penalty to apply during generation. + - `temperature` (`float`, *optional*, defaults to `1.0`): Temperature for sampling. Higher values lead to more random outputs. + - `top_p` (`float`, *optional*, defaults to `1.0`): Top-p (nucleus) sampling parameter. It controls the diversity of the generated text. + - `top_k` (`int`, *optional*, defaults to `-1`): Top-k sampling parameter. If set to `-1`, it disables top-k sampling. + - `min_p` (`float`, *optional*, defaults to `0.0`): Minimum probability threshold for sampling. + - `max_tokens` (`int`, *optional*, defaults to `16`): Maximum number of tokens to generate for each completion. + - `guided_decoding_regex` (`str`, *optional*): A regex pattern for guided decoding. If provided, the model will only generate tokens that match this regex pattern. + - `generation_kwargs` (`dict`, *optional*): Additional generation parameters to pass to the vLLM `SamplingParams`. This can include parameters like `seed`, `frequency_penalty`, etc. If it contains keys that conflict with the other parameters, they will override them. + + Returns: + `GenerateResponse`: + - `completion_ids` (list of list of `int`): A list of lists of token IDs for each generated completion. + + Example request: + ```json + {"prompts": ["Hello world", "What is AI?"]} + ``` + + Example response: + ```json + {"completion_ids": [[101, 102, 103], [201, 202, 203]]} + ``` + """ + + # Guided decoding, if enabled + if request.guided_decoding_regex is not None: + guided_decoding = GuidedDecodingParams(backend="outlines", regex=request.guided_decoding_regex) + else: + guided_decoding = None + + generation_kwargs = { + "n": request.n, + "repetition_penalty": request.repetition_penalty, + "temperature": request.temperature, + "top_p": request.top_p, + "top_k": request.top_k, + "min_p": request.min_p, + "max_tokens": request.max_tokens, + "guided_decoding": guided_decoding, + } + generation_kwargs.update(request.generation_kwargs) + sampling_params = SamplingParams(**generation_kwargs) + + # Evenly distribute prompts across DP ranks + chunked_prompts = chunk_list(request.prompts, script_args.data_parallel_size) + + # Send the prompts to each worker + for connection, prompts in zip(connections, chunked_prompts): + # When the number of prompts is less than data_parallel_size, some workers will receive empty prompts. + # However, vLLM requires that we always send at least one prompt. So we send a placeholder prompt to comply + # with vLLM's requirement, and we later ignore the result. + if not prompts: + prompts = [""] + kwargs = {"prompts": prompts, "sampling_params": sampling_params} + connection.send({"type": "call", "method": "generate", "kwargs": kwargs}) + + # Receive results + all_outputs = [connection.recv() for connection in connections] + + # Handle empty prompts (see above) + all_outputs = [output for output, prompts in zip(all_outputs, chunked_prompts) if prompts] + + # Flatten and combine all results + all_outputs = list(chain.from_iterable(all_outputs)) # from list of list to single list + completion_ids = [list(output.token_ids) for outputs in all_outputs for output in outputs.outputs] + return {"completion_ids": completion_ids} + + class InitCommunicatorRequest(BaseModel): + host: str + port: int + world_size: int + + @app.post("/init_communicator/") + async def init_communicator(request: InitCommunicatorRequest): + """ + Initializes the communicator for synchronizing model weights between a client and multiple server workers. + + Args: + request (`InitCommunicatorRequest`): + - `host` (`str`): Hostname or IP address of the master node. + - `port` (`int`): Port number to be used for communication. + - `world_size` (`int`): Total number of participating processes in the group. + """ + world_size = script_args.tensor_parallel_size * script_args.data_parallel_size + 1 + + # The function init_communicator is called this way: init_communicator(host, port, world_size) + # So with collective_rpc we need to call it this way: + # llm.collective_rpc(method="init_communicator", args=(host, port, world_size)) + kwargs = {"method": "init_communicator", "args": (request.host, request.port, world_size)} + for connection in connections: + connection.send({"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}) + + return {"message": "Request received, initializing communicator"} + + class UpdateWeightsRequest(BaseModel): + name: str + dtype: str + shape: list[int] + + @app.post("/update_named_param/") + async def update_named_param(request: UpdateWeightsRequest): + """ + Updates the model weights with the provided tensor. + + Once this endpoint is called, the client process should broadcast the updated weights to all server workers. + + Args: + request (`UpdateWeightsRequest`): + - `name` (`str`): Name of the weight tensor being updated. + - `dtype` (`str`): Data type of the weight tensor (e.g., `"torch.float32"`). + - `shape` (list of `int`): Shape of the weight + + """ + # The function update_named_param is called this way: update_named_param("name", torch.float32, (10, 10)) + # So with collective_rpc we need to call it this way: + # llm.collective_rpc("update_named_param", args=("name", torch.float32, (10, 10))) + dtype = torch.__getattribute__(request.dtype.split(".")[-1]) + kwargs = {"method": "update_named_param", "args": (request.name, dtype, tuple(request.shape))} + for connection in connections: + connection.send({"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}) + + return {"message": "Request received, updating named parameter"} + + @app.post("/reset_prefix_cache/") + async def reset_prefix_cache(): + """ + Resets the prefix cache for the model. + """ + for connection in connections: + connection.send({"type": "call", "method": "reset_prefix_cache"}) + # Wait for and collect all results + all_outputs = [connection.recv() for connection in connections] + success = all(output for output in all_outputs) + return {"message": "Request received, resetting prefix cache status: " + str(success)} + + @app.post("/close_communicator/") + async def close_communicator(): + """ + Closes the weight update group and cleans up associated resources. + """ + kwargs = {"method": "close_communicator"} + for connection in connections: + connection.send({"type": "fire_and_forget", "method": "collective_rpc", "kwargs": kwargs}) + return {"message": "Request received, closing communicator"} + + # Start the server + uvicorn.run(app, host=script_args.host, port=script_args.port, log_level=script_args.log_level) + + +def make_parser(subparsers: argparse._SubParsersAction = None): + if subparsers is not None: + parser = subparsers.add_parser("vllm-serve", help="Run the vLLM serve script", dataclass_types=ScriptArguments) + else: + parser = TrlParser(ScriptArguments) + return parser + + +if __name__ == "__main__": + parser = make_parser() + (script_args,) = parser.parse_args_and_config() + main(script_args) diff --git a/lib/python3.12/site-packages/trl/templates/lm_model_card.md b/lib/python3.12/site-packages/trl/templates/lm_model_card.md new file mode 100644 index 0000000000000000000000000000000000000000..1583c123f7569163cc276baaedeabc8c1ae69d1c --- /dev/null +++ b/lib/python3.12/site-packages/trl/templates/lm_model_card.md @@ -0,0 +1,55 @@ +--- +{{ card_data }} +--- + +# Model Card for {{ model_name }} + +This model is a fine-tuned version of [{{ base_model }}](https://huggingface.co/{{ base_model }}){% if dataset_name %} on the [{{ dataset_name }}](https://huggingface.co/datasets/{{ dataset_name }}) dataset{% endif %}. +It has been trained using [TRL](https://github.com/huggingface/trl). + +## Quick start + +```python +from transformers import pipeline + +question = "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?" +generator = pipeline("text-generation", model="{{ hub_model_id }}", device="cuda") +output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0] +print(output["generated_text"]) +``` + +## Training procedure + +{% if wandb_url %}[Visualize in Weights & Biases]({{ wandb_url }}){% endif %} +{% if comet_url %}[Visualize in Comet]({{ comet_url }}){% endif %} + +This model was trained with {{ trainer_name }}{% if paper_id %}, a method introduced in [{{ paper_title }}](https://huggingface.co/papers/{{ paper_id }}){% endif %}. + +### Framework versions + +- TRL: {{ trl_version }} +- Transformers: {{ transformers_version }} +- Pytorch: {{ pytorch_version }} +- Datasets: {{ datasets_version }} +- Tokenizers: {{ tokenizers_version }} + +## Citations + +{% if trainer_citation %}Cite {{ trainer_name }} as: + +```bibtex +{{ trainer_citation }} +```{% endif %} + +Cite TRL as: + +```bibtex +{% raw %}@misc{vonwerra2022trl, + title = {{TRL: Transformer Reinforcement Learning}}, + author = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec}, + year = 2020, + journal = {GitHub repository}, + publisher = {GitHub}, + howpublished = {\url{https://github.com/huggingface/trl}} +}{% endraw %} +``` diff --git a/lib/python3.12/site-packages/trl/trainer/__init__.py b/lib/python3.12/site-packages/trl/trainer/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a773f911ad54496a938e012bb3dffd2b4cbe51f5 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/__init__.py @@ -0,0 +1,165 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ..import_utils import OptionalDependencyNotAvailable, _LazyModule, is_diffusers_available + + +_import_structure = { + "alignprop_config": ["AlignPropConfig"], + "alignprop_trainer": ["AlignPropTrainer"], + "bco_config": ["BCOConfig"], + "bco_trainer": ["BCOTrainer"], + "callbacks": [ + "LogCompletionsCallback", + "MergeModelCallback", + "RichProgressCallback", + "SyncRefModelCallback", + "WinRateCallback", + ], + "cpo_config": ["CPOConfig"], + "cpo_trainer": ["CPOTrainer"], + "ddpo_config": ["DDPOConfig"], + "dpo_config": ["DPOConfig", "FDivergenceConstants", "FDivergenceType"], + "dpo_trainer": ["DPOTrainer"], + "gkd_config": ["GKDConfig"], + "gkd_trainer": ["GKDTrainer"], + "grpo_config": ["GRPOConfig"], + "grpo_trainer": ["GRPOTrainer"], + "iterative_sft_config": ["IterativeSFTConfig"], + "iterative_sft_trainer": ["IterativeSFTTrainer"], + "judges": [ + "AllTrueJudge", + "BaseBinaryJudge", + "BaseJudge", + "BasePairwiseJudge", + "BaseRankJudge", + "HfPairwiseJudge", + "OpenAIPairwiseJudge", + "PairRMJudge", + ], + "kto_config": ["KTOConfig"], + "kto_trainer": ["KTOTrainer"], + "model_config": ["ModelConfig"], + "nash_md_config": ["NashMDConfig"], + "nash_md_trainer": ["NashMDTrainer"], + "online_dpo_config": ["OnlineDPOConfig"], + "online_dpo_trainer": ["OnlineDPOTrainer"], + "orpo_config": ["ORPOConfig"], + "orpo_trainer": ["ORPOTrainer"], + "ppo_config": ["PPOConfig"], + "ppo_trainer": ["PPOTrainer"], + "prm_config": ["PRMConfig"], + "prm_trainer": ["PRMTrainer"], + "reward_config": ["RewardConfig"], + "reward_trainer": ["RewardTrainer"], + "rloo_config": ["RLOOConfig"], + "rloo_trainer": ["RLOOTrainer"], + "sft_config": ["SFTConfig"], + "sft_trainer": ["SFTTrainer"], + "utils": [ + "ConstantLengthDataset", + "DataCollatorForCompletionOnlyLM", + "RunningMoments", + "compute_accuracy", + "disable_dropout_in_model", + "empty_cache", + "peft_module_casting_to_bf16", + ], + "xpo_config": ["XPOConfig"], + "xpo_trainer": ["XPOTrainer"], +} +try: + if not is_diffusers_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + pass +else: + _import_structure["ddpo_trainer"] = ["DDPOTrainer"] + +if TYPE_CHECKING: + from .alignprop_config import AlignPropConfig + from .alignprop_trainer import AlignPropTrainer + from .bco_config import BCOConfig + from .bco_trainer import BCOTrainer + from .callbacks import ( + LogCompletionsCallback, + MergeModelCallback, + RichProgressCallback, + SyncRefModelCallback, + WinRateCallback, + ) + from .cpo_config import CPOConfig + from .cpo_trainer import CPOTrainer + from .ddpo_config import DDPOConfig + from .dpo_config import DPOConfig, FDivergenceConstants, FDivergenceType + from .dpo_trainer import DPOTrainer + from .gkd_config import GKDConfig + from .gkd_trainer import GKDTrainer + from .grpo_config import GRPOConfig + from .grpo_trainer import GRPOTrainer + from .iterative_sft_trainer import IterativeSFTConfig, IterativeSFTTrainer + from .judges import ( + AllTrueJudge, + BaseBinaryJudge, + BaseJudge, + BasePairwiseJudge, + BaseRankJudge, + HfPairwiseJudge, + OpenAIPairwiseJudge, + PairRMJudge, + ) + from .kto_config import KTOConfig + from .kto_trainer import KTOTrainer + from .model_config import ModelConfig + from .nash_md_config import NashMDConfig + from .nash_md_trainer import NashMDTrainer + from .online_dpo_config import OnlineDPOConfig + from .online_dpo_trainer import OnlineDPOTrainer + from .orpo_config import ORPOConfig + from .orpo_trainer import ORPOTrainer + from .ppo_config import PPOConfig + from .ppo_trainer import PPOTrainer + from .prm_config import PRMConfig + from .prm_trainer import PRMTrainer + from .reward_config import RewardConfig + from .reward_trainer import RewardTrainer + from .rloo_config import RLOOConfig + from .rloo_trainer import RLOOTrainer + from .sft_config import SFTConfig + from .sft_trainer import SFTTrainer + from .utils import ( + ConstantLengthDataset, + DataCollatorForCompletionOnlyLM, + RunningMoments, + compute_accuracy, + disable_dropout_in_model, + empty_cache, + peft_module_casting_to_bf16, + ) + from .xpo_config import XPOConfig + from .xpo_trainer import XPOTrainer + + try: + if not is_diffusers_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + pass + else: + from .ddpo_trainer import DDPOTrainer +else: + import sys + + sys.modules[__name__] = _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__) diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21e3ace9b80f0281d61435f4b7a6dbadabf96a37 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/alignprop_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/alignprop_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78dd76e788b0ac59e3b9e6bbbf67d4d3f3d94f49 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/alignprop_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/alignprop_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/alignprop_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5449be3ba2c80f7e9eea6a8f69d7cd55df64b29d Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/alignprop_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/bco_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/bco_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e887df544a226a0ed25080f17c6203287c68e803 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/bco_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/bco_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/bco_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57958f88f486840e4217a46b35668632e50f5f00 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/bco_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/callbacks.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/callbacks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..825d99e0791125acfa9a222bef24fbe40051428f Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/callbacks.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/cpo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/cpo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f615b5169fc9908f2154549553818f44b24ded6a Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/cpo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/cpo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/cpo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6be377ad534e1b92b337dc403ee08543222fa329 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/cpo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/ddpo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/ddpo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a1e5babf864fd6a979c992cd1d687ec39748a4cb Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/ddpo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/ddpo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/ddpo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..676e8f67a5ab614d28276426f6212ad0289b2cd0 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/ddpo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/dpo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/dpo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb4373aaec4a2953d6716e6adbbf7a029de1cf0e Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/dpo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/dpo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/dpo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3bb500ce34dc85281e8105ca7696d214a0da861e Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/dpo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/gkd_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/gkd_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a3630055f2cbe9ef89342236fa16a66074b2f32 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/gkd_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/gkd_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/gkd_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dad58f4d4d6d9254f4b3a364be00ac1478d000de Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/gkd_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/grpo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/grpo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8762baa4acda5a9a3e36aca9849ea7b3918145dc Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/grpo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/grpo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/grpo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..066ec6c3dbadeeb3c9dea0755836161f10e5d347 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/grpo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/iterative_sft_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/iterative_sft_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e00c4e145a3cc53cb6838d7608cb162e2237a64 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/iterative_sft_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/iterative_sft_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/iterative_sft_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..47ad0e574f0afe4671cc63edae442acf93a8cdfe Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/iterative_sft_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/judges.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/judges.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca8bd93c0fc64c72fb405ee2be2203bd67ba4865 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/judges.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/kto_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/kto_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4aa272351f1bbb364c787f8d342b3be46121f308 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/kto_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/kto_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/kto_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94b12add267491c66fe1a27dc002d8b0926147b5 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/kto_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/model_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/model_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3dbdd5d03d40225cf2329bc717ea332215da97ad Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/model_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/nash_md_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/nash_md_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54bf116f5504419c8984a7242e06a09bcef624c3 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/nash_md_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/nash_md_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/nash_md_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f6d85b5187f7d187a2446d2ee66f7afbd521933 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/nash_md_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/online_dpo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/online_dpo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41409d2a6d1845e6c4bd96f6772917a31e1881c7 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/online_dpo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/online_dpo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/online_dpo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b10ed5794ac347dd554f5f4eb27679a04097e140 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/online_dpo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/orpo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/orpo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36d8d1d0568b99d780574489f623897dd658fe64 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/orpo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/orpo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/orpo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a2e2248d0f8b1ec22ae905075d0a295aeebfad5 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/orpo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/ppo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/ppo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e76c2dd116ff91c54b067d71c4c8db4e39ff5ff7 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/ppo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/ppo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/ppo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..af4d42acb7165b6519b3d92ff93eded5195fa747 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/ppo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/prm_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/prm_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..588cc32ee2908345c823d7ad193ea65841f4cbe7 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/prm_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/prm_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/prm_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c31820163c4708a48857d0eafeb5253d8c3e55d3 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/prm_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/reward_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/reward_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..128d5940936427fd379b3d9769ba70c10f1e9e93 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/reward_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/reward_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/reward_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f99464cce25863ddfa4a9f85498c5679df9beed8 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/reward_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/rloo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/rloo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e81c1795fc0d28fc674710fe491224e01652146 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/rloo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/rloo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/rloo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aebd366c26fb00c59b6ba82a930fcebd7e312f28 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/rloo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/sft_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/sft_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0560ef7e737f79e8c1cdcb2b6c6d34802ad95a24 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/sft_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/sft_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/sft_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72363d5da0e1842f61253e7323165a0fd889c969 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/sft_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..28ef871f7880c5cfc871a7764c7033ac1bca6c4e Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/xpo_config.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/xpo_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..05cd4c33b85543307d35fcea5109f9005006226f Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/xpo_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/__pycache__/xpo_trainer.cpython-312.pyc b/lib/python3.12/site-packages/trl/trainer/__pycache__/xpo_trainer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e2dc88db7a0fade22bd43963bc89a9517922ebf0 Binary files /dev/null and b/lib/python3.12/site-packages/trl/trainer/__pycache__/xpo_trainer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/trl/trainer/alignprop_config.py b/lib/python3.12/site-packages/trl/trainer/alignprop_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6eef04df1d7d441a0819f00682db38c427d89f7c --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/alignprop_config.py @@ -0,0 +1,192 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +from dataclasses import dataclass, field +from typing import Any, Optional + +from transformers import is_bitsandbytes_available + +from ..core import flatten_dict + + +@dataclass +class AlignPropConfig: + r""" + Configuration class for the [`AlignPropTrainer`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + exp_name (`str`, *optional*, defaults to `os.path.basename(sys.argv[0])[: -len(".py")]`): + Name of this experiment (defaults to the file name without the extension). + run_name (`str`, *optional*, defaults to `""`): + Name of this run. + seed (`int`, *optional*, defaults to `0`): + Random seed for reproducibility. + log_with (`str` or `None`, *optional*, defaults to `None`): + Log with either `"wandb"` or `"tensorboard"`. Check + [tracking](https://huggingface.co/docs/accelerate/usage_guides/tracking) for more details. + log_image_freq (`int`, *optional*, defaults to `1`): + Frequency for logging images. + tracker_kwargs (`dict[str, Any]`, *optional*, defaults to `{}`): + Keyword arguments for the tracker (e.g., `wandb_project`). + accelerator_kwargs (`dict[str, Any]`, *optional*, defaults to `{}`): + Keyword arguments for the accelerator. + project_kwargs (`dict[str, Any]`, *optional*, defaults to `{}`): + Keyword arguments for the accelerator project config (e.g., `logging_dir`). + tracker_project_name (`str`, *optional*, defaults to `"trl"`): + Name of project to use for tracking. + logdir (`str`, *optional*, defaults to `"logs"`): + Top-level logging directory for checkpoint saving. + num_epochs (`int`, *optional*, defaults to `100`): + Number of epochs to train. + save_freq (`int`, *optional*, defaults to `1`): + Number of epochs between saving model checkpoints. + num_checkpoint_limit (`int`, *optional*, defaults to `5`): + Number of checkpoints to keep before overwriting old ones. + mixed_precision (`str`, *optional*, defaults to `"fp16"`): + Mixed precision training. + allow_tf32 (`bool`, *optional*, defaults to `True`): + Allow `tf32` on Ampere GPUs. + resume_from (`str`, *optional*, defaults to `""`): + Path to resume training from a checkpoint. + sample_num_steps (`int`, *optional*, defaults to `50`): + Number of sampler inference steps. + sample_eta (`float`, *optional*, defaults to `1.0`): + Eta parameter for the DDIM sampler. + sample_guidance_scale (`float`, *optional*, defaults to `5.0`): + Classifier-free guidance weight. + train_batch_size (`int`, *optional*, defaults to `1`): + Batch size for training. + train_use_8bit_adam (`bool`, *optional*, defaults to `False`): + Whether to use the 8bit Adam optimizer from `bitsandbytes`. + train_learning_rate (`float`, *optional*, defaults to `1e-3`): + Learning rate. + train_adam_beta1 (`float`, *optional*, defaults to `0.9`): + Beta1 for Adam optimizer. + train_adam_beta2 (`float`, *optional*, defaults to `0.999`): + Beta2 for Adam optimizer. + train_adam_weight_decay (`float`, *optional*, defaults to `1e-4`): + Weight decay for Adam optimizer. + train_adam_epsilon (`float`, *optional*, defaults to `1e-8`): + Epsilon value for Adam optimizer. + train_gradient_accumulation_steps (`int`, *optional*, defaults to `1`): + Number of gradient accumulation steps. + train_max_grad_norm (`float`, *optional*, defaults to `1.0`): + Maximum gradient norm for gradient clipping. + negative_prompts (`str` or `None`, *optional*, defaults to `None`): + Comma-separated list of prompts to use as negative examples. + truncated_backprop_rand (`bool`, *optional*, defaults to `True`): + If `True`, randomized truncation to different diffusion timesteps is used. + truncated_backprop_timestep (`int`, *optional*, defaults to `49`): + Absolute timestep to which the gradients are backpropagated. Used only if `truncated_backprop_rand=False`. + truncated_rand_backprop_minmax (`tuple[int, int]`, *optional*, defaults to `(0, 50)`): + Range of diffusion timesteps for randomized truncated backpropagation. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the final model to the Hub. + """ + + exp_name: str = field( + default=os.path.basename(sys.argv[0])[: -len(".py")], + metadata={"help": "Name of this experiment (defaults to the file name without the extension)."}, + ) + run_name: str = field(default="", metadata={"help": "Name of this run."}) + seed: int = field(default=0, metadata={"help": "Random seed for reproducibility."}) + log_with: Optional[str] = field( + default=None, + metadata={"help": "Log with either 'wandb' or 'tensorboard'.", "choices": ["wandb", "tensorboard"]}, + ) + log_image_freq: int = field(default=1, metadata={"help": "Frequency for logging images."}) + tracker_kwargs: dict[str, Any] = field( + default_factory=dict, + metadata={"help": "Keyword arguments for the tracker (e.g., `wandb_project`)."}, + ) + accelerator_kwargs: dict[str, Any] = field( + default_factory=dict, metadata={"help": "Keyword arguments for the accelerator."} + ) + project_kwargs: dict[str, Any] = field( + default_factory=dict, + metadata={"help": "Keyword arguments for the accelerator project config (e.g., `logging_dir`)."}, + ) + tracker_project_name: str = field(default="trl", metadata={"help": "Name of project to use for tracking."}) + logdir: str = field(default="logs", metadata={"help": "Top-level logging directory for checkpoint saving."}) + num_epochs: int = field(default=100, metadata={"help": "Number of epochs to train."}) + save_freq: int = field(default=1, metadata={"help": "Number of epochs between saving model checkpoints."}) + num_checkpoint_limit: int = field( + default=5, metadata={"help": "Number of checkpoints to keep before overwriting old ones."} + ) + mixed_precision: str = field( + default="fp16", + metadata={ + "help": "Mixed precision training. Possible values are 'fp16', 'bf16', 'none'.", + "choices": ["fp16", "bf16", "none"], + }, + ) + allow_tf32: bool = field(default=True, metadata={"help": "Allow `tf32` on Ampere GPUs."}) + resume_from: str = field(default="", metadata={"help": "Path to resume training from a checkpoint."}) + sample_num_steps: int = field(default=50, metadata={"help": "Number of sampler inference steps."}) + sample_eta: float = field(default=1.0, metadata={"help": "Eta parameter for the DDIM sampler."}) + sample_guidance_scale: float = field(default=5.0, metadata={"help": "Classifier-free guidance weight."}) + train_batch_size: int = field(default=1, metadata={"help": "Batch size for training."}) + train_use_8bit_adam: bool = field( + default=False, metadata={"help": "Whether to use the 8bit Adam optimizer from `bitsandbytes`."} + ) + train_learning_rate: float = field(default=1e-3, metadata={"help": "Learning rate."}) + train_adam_beta1: float = field(default=0.9, metadata={"help": "Beta1 for Adam optimizer."}) + train_adam_beta2: float = field(default=0.999, metadata={"help": "Beta2 for Adam optimizer."}) + train_adam_weight_decay: float = field(default=1e-4, metadata={"help": "Weight decay for Adam optimizer."}) + train_adam_epsilon: float = field(default=1e-8, metadata={"help": "Epsilon value for Adam optimizer."}) + train_gradient_accumulation_steps: int = field( + default=1, metadata={"help": "Number of gradient accumulation steps."} + ) + train_max_grad_norm: float = field(default=1.0, metadata={"help": "Maximum gradient norm for gradient clipping."}) + negative_prompts: Optional[str] = field( + default=None, + metadata={"help": "Comma-separated list of prompts to use as negative examples."}, + ) + truncated_backprop_rand: bool = field( + default=True, + metadata={"help": "If `True`, randomized truncation to different diffusion timesteps is used."}, + ) + truncated_backprop_timestep: int = field( + default=49, + metadata={ + "help": "Absolute timestep to which the gradients are backpropagated. Used only if " + "`truncated_backprop_rand=False`." + }, + ) + truncated_rand_backprop_minmax: tuple[int, int] = field( + default=(0, 50), + metadata={ + "help": "Range of diffusion timesteps for randomized truncated backpropagation.", + }, + ) + push_to_hub: bool = field(default=False, metadata={"help": "Whether to push the final model to the Hub."}) + + def to_dict(self): + output_dict = {} + for key, value in self.__dict__.items(): + output_dict[key] = value + return flatten_dict(output_dict) + + def __post_init__(self): + if self.train_use_8bit_adam and not is_bitsandbytes_available(): + raise ImportError( + "You need to install bitsandbytes to use 8bit Adam. " + "You can install it with `pip install bitsandbytes`." + ) diff --git a/lib/python3.12/site-packages/trl/trainer/alignprop_trainer.py b/lib/python3.12/site-packages/trl/trainer/alignprop_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..fcabb47d67228e16b757e32a564d0225531751dc --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/alignprop_trainer.py @@ -0,0 +1,465 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import textwrap +from collections import defaultdict +from pathlib import Path +from typing import Any, Callable, Optional, Union +from warnings import warn + +import torch +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import ProjectConfiguration, set_seed +from huggingface_hub import PyTorchModelHubMixin +from transformers import is_wandb_available + +from ..models import DDPOStableDiffusionPipeline +from .alignprop_config import AlignPropConfig +from .utils import generate_model_card, get_comet_experiment_url + + +if is_wandb_available(): + import wandb + +logger = get_logger(__name__) + + +class AlignPropTrainer(PyTorchModelHubMixin): + """ + The AlignPropTrainer uses Deep Diffusion Policy Optimization to optimise diffusion models. Note, this trainer is + heavily inspired by the work here: https://github.com/mihirp1998/AlignProp/ As of now only Stable Diffusion based + pipelines are supported + + Attributes: + config (`AlignPropConfig`): + Configuration object for AlignPropTrainer. Check the documentation of `PPOConfig` for more details. + reward_function (`Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor]`): + Reward function to be used + prompt_function (`Callable[[], tuple[str, Any]]`): + Function to generate prompts to guide model + sd_pipeline (`DDPOStableDiffusionPipeline`): + Stable Diffusion pipeline to be used for training. + image_samples_hook (`Optional[Callable[[Any, Any, Any], Any]]`): + Hook to be called to log images + """ + + _tag_names = ["trl", "alignprop"] + + def __init__( + self, + config: AlignPropConfig, + reward_function: Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor], + prompt_function: Callable[[], tuple[str, Any]], + sd_pipeline: DDPOStableDiffusionPipeline, + image_samples_hook: Optional[Callable[[Any, Any, Any], Any]] = None, + ): + if image_samples_hook is None: + warn("No image_samples_hook provided; no images will be logged") + + self.prompt_fn = prompt_function + self.reward_fn = reward_function + self.config = config + self.image_samples_callback = image_samples_hook + + accelerator_project_config = ProjectConfiguration(**self.config.project_kwargs) + + if self.config.resume_from: + self.config.resume_from = os.path.normpath(os.path.expanduser(self.config.resume_from)) + if "checkpoint_" not in os.path.basename(self.config.resume_from): + # get the most recent checkpoint in this directory + checkpoints = list( + filter( + lambda x: "checkpoint_" in x, + os.listdir(self.config.resume_from), + ) + ) + if len(checkpoints) == 0: + raise ValueError(f"No checkpoints found in {self.config.resume_from}") + checkpoint_numbers = sorted([int(x.split("_")[-1]) for x in checkpoints]) + self.config.resume_from = os.path.join( + self.config.resume_from, + f"checkpoint_{checkpoint_numbers[-1]}", + ) + + accelerator_project_config.iteration = checkpoint_numbers[-1] + 1 + + self.accelerator = Accelerator( + log_with=self.config.log_with, + mixed_precision=self.config.mixed_precision, + project_config=accelerator_project_config, + # we always accumulate gradients across timesteps; we want config.train.gradient_accumulation_steps to be the + # number of *samples* we accumulate across, so we need to multiply by the number of training timesteps to get + # the total number of optimizer steps to accumulate across. + gradient_accumulation_steps=self.config.train_gradient_accumulation_steps, + **self.config.accelerator_kwargs, + ) + + is_using_tensorboard = config.log_with is not None and config.log_with == "tensorboard" + + if self.accelerator.is_main_process: + self.accelerator.init_trackers( + self.config.tracker_project_name, + config=dict(alignprop_trainer_config=config.to_dict()) + if not is_using_tensorboard + else config.to_dict(), + init_kwargs=self.config.tracker_kwargs, + ) + + logger.info(f"\n{config}") + + set_seed(self.config.seed, device_specific=True) + + self.sd_pipeline = sd_pipeline + + self.sd_pipeline.set_progress_bar_config( + position=1, + disable=not self.accelerator.is_local_main_process, + leave=False, + desc="Timestep", + dynamic_ncols=True, + ) + + # For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + if self.accelerator.mixed_precision == "fp16": + inference_dtype = torch.float16 + elif self.accelerator.mixed_precision == "bf16": + inference_dtype = torch.bfloat16 + else: + inference_dtype = torch.float32 + + self.sd_pipeline.vae.to(self.accelerator.device, dtype=inference_dtype) + self.sd_pipeline.text_encoder.to(self.accelerator.device, dtype=inference_dtype) + self.sd_pipeline.unet.to(self.accelerator.device, dtype=inference_dtype) + + trainable_layers = self.sd_pipeline.get_trainable_layers() + + self.accelerator.register_save_state_pre_hook(self._save_model_hook) + self.accelerator.register_load_state_pre_hook(self._load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if self.config.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + self.optimizer = self._setup_optimizer( + trainable_layers.parameters() if not isinstance(trainable_layers, list) else trainable_layers + ) + + self.neg_prompt_embed = self.sd_pipeline.text_encoder( + self.sd_pipeline.tokenizer( + [""] if self.config.negative_prompts is None else self.config.negative_prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.sd_pipeline.tokenizer.model_max_length, + ).input_ids.to(self.accelerator.device) + )[0] + + # NOTE: for some reason, autocast is necessary for non-lora training but for lora training it isn't necessary and it uses + # more memory + self.autocast = self.sd_pipeline.autocast or self.accelerator.autocast + + if hasattr(self.sd_pipeline, "use_lora") and self.sd_pipeline.use_lora: + unet, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer) + self.trainable_layers = list(filter(lambda p: p.requires_grad, unet.parameters())) + else: + self.trainable_layers, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer) + + if config.resume_from: + logger.info(f"Resuming from {config.resume_from}") + self.accelerator.load_state(config.resume_from) + self.first_epoch = int(config.resume_from.split("_")[-1]) + 1 + else: + self.first_epoch = 0 + + def compute_rewards(self, prompt_image_pairs): + reward, reward_metadata = self.reward_fn( + prompt_image_pairs["images"], prompt_image_pairs["prompts"], prompt_image_pairs["prompt_metadata"] + ) + return reward + + def step(self, epoch: int, global_step: int): + """ + Perform a single step of training. + + Args: + epoch (int): The current epoch. + global_step (int): The current global step. + + Side Effects: + - Model weights are updated + - Logs the statistics to the accelerator trackers. + - If `self.image_samples_callback` is not None, it will be called with the prompt_image_pairs, global_step, + and the accelerator tracker. + + Returns: + global_step (int): The updated global step. + """ + info = defaultdict(list) + + self.sd_pipeline.unet.train() + + for _ in range(self.config.train_gradient_accumulation_steps): + with self.accelerator.accumulate(self.sd_pipeline.unet), self.autocast(), torch.enable_grad(): + prompt_image_pairs = self._generate_samples( + batch_size=self.config.train_batch_size, + ) + + rewards = self.compute_rewards(prompt_image_pairs) + + prompt_image_pairs["rewards"] = rewards + + rewards_vis = self.accelerator.gather(rewards).detach().cpu().numpy() + + loss = self.calculate_loss(rewards) + + self.accelerator.backward(loss) + + if self.accelerator.sync_gradients: + self.accelerator.clip_grad_norm_( + self.trainable_layers.parameters() + if not isinstance(self.trainable_layers, list) + else self.trainable_layers, + self.config.train_max_grad_norm, + ) + + self.optimizer.step() + self.optimizer.zero_grad() + + info["reward_mean"].append(rewards_vis.mean()) + info["reward_std"].append(rewards_vis.std()) + info["loss"].append(loss.item()) + + # Checks if the accelerator has performed an optimization step behind the scenes + if self.accelerator.sync_gradients: + # log training-related stuff + info = {k: torch.mean(torch.tensor(v)) for k, v in info.items()} + info = self.accelerator.reduce(info, reduction="mean") + info.update({"epoch": epoch}) + self.accelerator.log(info, step=global_step) + global_step += 1 + info = defaultdict(list) + else: + raise ValueError( + "Optimization step should have been performed by this point. Please check calculated gradient accumulation settings." + ) + # Logs generated images + if self.image_samples_callback is not None and global_step % self.config.log_image_freq == 0: + self.image_samples_callback(prompt_image_pairs, global_step, self.accelerator.trackers[0]) + + if epoch != 0 and epoch % self.config.save_freq == 0 and self.accelerator.is_main_process: + self.accelerator.save_state() + + return global_step + + def calculate_loss(self, rewards): + """ + Calculate the loss for a batch of an unpacked sample + + Args: + rewards (torch.Tensor): + Differentiable reward scalars for each generated image, shape: [batch_size] + + Returns: + loss (torch.Tensor) (all of these are of shape (1,)) + """ + # Loss is specific to Aesthetic Reward function used in AlignProp (https://huggingface.co/papers/2310.03739) + loss = 10.0 - (rewards).mean() + return loss + + def loss( + self, + advantages: torch.Tensor, + clip_range: float, + ratio: torch.Tensor, + ): + unclipped_loss = -advantages * ratio + clipped_loss = -advantages * torch.clamp( + ratio, + 1.0 - clip_range, + 1.0 + clip_range, + ) + return torch.mean(torch.maximum(unclipped_loss, clipped_loss)) + + def _setup_optimizer(self, trainable_layers_parameters): + if self.config.train_use_8bit_adam: + import bitsandbytes + + optimizer_cls = bitsandbytes.optim.AdamW8bit + else: + optimizer_cls = torch.optim.AdamW + + return optimizer_cls( + trainable_layers_parameters, + lr=self.config.train_learning_rate, + betas=(self.config.train_adam_beta1, self.config.train_adam_beta2), + weight_decay=self.config.train_adam_weight_decay, + eps=self.config.train_adam_epsilon, + ) + + def _save_model_hook(self, models, weights, output_dir): + self.sd_pipeline.save_checkpoint(models, weights, output_dir) + weights.pop() # ensures that accelerate doesn't try to handle saving of the model + + def _load_model_hook(self, models, input_dir): + self.sd_pipeline.load_checkpoint(models, input_dir) + models.pop() # ensures that accelerate doesn't try to handle loading of the model + + def _generate_samples(self, batch_size, with_grad=True, prompts=None): + """ + Generate samples from the model + + Args: + batch_size (int): Batch size to use for sampling + with_grad (bool): Whether the generated RGBs should have gradients attached to it. + + Returns: + prompt_image_pairs (dict[Any]) + """ + prompt_image_pairs = {} + + sample_neg_prompt_embeds = self.neg_prompt_embed.repeat(batch_size, 1, 1) + + if prompts is None: + prompts, prompt_metadata = zip(*[self.prompt_fn() for _ in range(batch_size)]) + else: + prompt_metadata = [{} for _ in range(batch_size)] + + prompt_ids = self.sd_pipeline.tokenizer( + prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.sd_pipeline.tokenizer.model_max_length, + ).input_ids.to(self.accelerator.device) + + prompt_embeds = self.sd_pipeline.text_encoder(prompt_ids)[0] + + if with_grad: + sd_output = self.sd_pipeline.rgb_with_grad( + prompt_embeds=prompt_embeds, + negative_prompt_embeds=sample_neg_prompt_embeds, + num_inference_steps=self.config.sample_num_steps, + guidance_scale=self.config.sample_guidance_scale, + eta=self.config.sample_eta, + truncated_backprop_rand=self.config.truncated_backprop_rand, + truncated_backprop_timestep=self.config.truncated_backprop_timestep, + truncated_rand_backprop_minmax=self.config.truncated_rand_backprop_minmax, + output_type="pt", + ) + else: + sd_output = self.sd_pipeline( + prompt_embeds=prompt_embeds, + negative_prompt_embeds=sample_neg_prompt_embeds, + num_inference_steps=self.config.sample_num_steps, + guidance_scale=self.config.sample_guidance_scale, + eta=self.config.sample_eta, + output_type="pt", + ) + + images = sd_output.images + + prompt_image_pairs["images"] = images + prompt_image_pairs["prompts"] = prompts + prompt_image_pairs["prompt_metadata"] = prompt_metadata + + return prompt_image_pairs + + def train(self, epochs: Optional[int] = None): + """ + Train the model for a given number of epochs + """ + global_step = 0 + if epochs is None: + epochs = self.config.num_epochs + for epoch in range(self.first_epoch, epochs): + global_step = self.step(epoch, global_step) + + def _save_pretrained(self, save_directory): + self.sd_pipeline.save_pretrained(save_directory) + self.create_model_card() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{prabhudesai2024aligning, + title = {{Aligning Text-to-Image Diffusion Models with Reward Backpropagation}}, + author = {Mihir Prabhudesai and Anirudh Goyal and Deepak Pathak and Katerina Fragkiadaki}, + year = 2024, + eprint = {arXiv:2310.03739} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="AlignProp", + trainer_citation=citation, + paper_title="Aligning Text-to-Image Diffusion Models with Reward Backpropagation", + paper_id="2310.03739", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/bco_config.py b/lib/python3.12/site-packages/trl/trainer/bco_config.py new file mode 100644 index 0000000000000000000000000000000000000000..10e6263a6c46664edd4d1eb9a4c3099205afd501 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/bco_config.py @@ -0,0 +1,206 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Any, Optional + +from transformers import TrainingArguments + + +@dataclass +class BCOConfig(TrainingArguments): + r""" + Configuration class for the [`BCOTrainer`]. + + This class includes only the parameters that are specific to BCO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from both the model and the reference model to W&B or Comet + during evaluation. + is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): + Whether to precompute reference model log probabilities for training and evaluation datasets. This is + useful when training without the reference model to reduce the total GPU memory needed. + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + ref_model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the reference model + from a string. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + prompt_sample_size (`int`, *optional*, defaults to `1024`): + Number of prompts that are fed to density ratio classifier. + min_density_ratio (`float`, *optional*, defaults to `0.5`): + Minimum value of the density ratio. The estimated density ratio is clamped to this value. + max_density_ratio (`float`, *optional*, defaults to `10.0`): + Maximum value of the density ratio. The estimated density ratio is clamped to this value. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs", "ref_model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + max_length: Optional[int] = field( + default=1024, + metadata={ + "help": "Maximum length of the sequences (prompt + completion) in the batch. " + "This argument is required if you want to use the default data collator." + }, + ) + max_prompt_length: Optional[int] = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. " + "This argument is required if you want to use the default data collator." + }, + ) + max_completion_length: Optional[int] = field( + default=None, + metadata={ + "help": "Maximum length of the completion. This argument is required if you want to use the " + "default data collator and your model is an encoder-decoder." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the deviation from the reference model. " + "Higher β means less deviation from the reference model." + }, + ) + label_pad_token_id: int = field( + default=-100, + metadata={ + "help": "Label pad token id. This argument is required if you want to use the default data collator." + }, + ) + padding_value: Optional[int] = field( + default=None, + metadata={"help": "Padding value to use. If `None`, the padding value of the tokenizer is used."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the prompt is too long. Possible values are " + "`keep_end` or `keep_start`. This argument is required if you want to use the " + "default data collator." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model and reference model."}, + ) + generate_during_eval: bool = field( + default=False, + metadata={ + "help": "If `True`, generates and logs completions from both the model and the reference model " + "to W&B during evaluation." + }, + ) + is_encoder_decoder: Optional[bool] = field( + default=None, + metadata={ + "help": "When using the `model_init` argument (callable) to instantiate the model instead of the " + "`model` argument, you need to specify if the model returned by the callable is an " + "encoder-decoder model." + }, + ) + precompute_ref_log_probs: bool = field( + default=False, + metadata={ + "help": "Whether to precompute reference model log probabilities for training and evaluation datasets. " + "This is useful when training without the reference model to reduce the total GPU memory " + "needed." + }, + ) + model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " + "model from a string." + }, + ) + ref_model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " + "reference model from a string." + }, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + prompt_sample_size: int = field( + default=1024, + metadata={"help": "Number of prompts that are fed to density ratio classifier."}, + ) + min_density_ratio: float = field( + default=0.5, + metadata={"help": "Minimum value of the density ratio. The estimated density ratio is clamped to this value."}, + ) + max_density_ratio: float = field( + default=10.0, + metadata={"help": "Maximum value of the density ratio. The estimated density ratio is clamped to this value."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/lib/python3.12/site-packages/trl/trainer/bco_trainer.py b/lib/python3.12/site-packages/trl/trainer/bco_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..3a8a051f00b8dc46e3b82784e6ccbd28a4aec99b --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/bco_trainer.py @@ -0,0 +1,1543 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import random +import textwrap +import warnings +from collections import defaultdict +from contextlib import contextmanager, nullcontext +from operator import itemgetter +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState +from accelerate.logging import get_logger +from accelerate.utils import tqdm +from datasets import Dataset +from torch import autocast +from torch.utils.data import DataLoader, SequentialSampler +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + TrainingArguments, + is_comet_available, + is_sklearn_available, + is_wandb_available, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalLoopOutput, has_length +from transformers.utils import is_peft_available + +from ..data_utils import maybe_apply_chat_template +from ..import_utils import is_joblib_available +from ..models import create_reference_model, prepare_deepspeed +from .bco_config import BCOConfig +from .utils import ( + DPODataCollatorWithPadding, + RunningMoments, + disable_dropout_in_model, + generate_model_card, + get_comet_experiment_url, + log_table_to_comet_experiment, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + +if is_wandb_available(): + import wandb + +if is_sklearn_available(): + from sklearn.linear_model import LogisticRegression + +if is_joblib_available(): + import joblib + +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizer + +logger = get_logger(__name__) + +RUNNING_NAME = "running.json" +CLF_NAME = "clf.pkl" + + +def _tokenize( + batch: dict[str, list[Any]], + tokenizer: "PreTrainedTokenizer", + embedding_tokenizer: Optional["PreTrainedTokenizer"] = None, +) -> dict[str, list[Any]]: + """Tokenize a batch from a BCO specific dataset.""" + prompt_tokenized = tokenizer(batch["prompt"], add_special_tokens=False) + prompt_input_ids = prompt_tokenized["input_ids"] + prompt_attention_mask = prompt_tokenized["attention_mask"] + prompt_and_completion = [prompt + completion for prompt, completion in zip(batch["prompt"], batch["completion"])] + full_tokenized = tokenizer(prompt_and_completion, add_special_tokens=False) + full_input_ids = full_tokenized["input_ids"] + full_attention_mask = full_tokenized["attention_mask"] + + answer_input_ids = [f[len(p) :] for f, p in zip(full_input_ids, prompt_input_ids)] + answer_attention_mask = [f[len(p) :] for f, p in zip(full_attention_mask, prompt_attention_mask)] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = [np.concatenate([p, a]) for p, a in zip(prompt_input_ids, answer_input_ids)] + # Prepare input tokens for token by token comparison + full_input_ids = [np.array(f) for f in full_input_ids] + for full, concat in zip(full_input_ids, full_concat_input_ids): + if len(full) != len(concat): + raise ValueError( + "The elements in 'full_input_ids' and 'full_concat_input_ids' must have the same pairwise length." + ) + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = [len(p) for p in prompt_input_ids] + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + for idx, (p, f, r) in enumerate(zip(prompt_input_ids, full_input_ids, response_token_ids_start_idx)): + if not np.array_equal(p, f[:r]): + response_token_ids_start_idx[idx] -= 1 + + prompt_input_ids = [f[:r] for f, r in zip(full_input_ids, response_token_ids_start_idx)] + prompt_attention_mask = [f[:r] for f, r in zip(full_attention_mask, response_token_ids_start_idx)] + + for p, m in zip(prompt_input_ids, prompt_attention_mask): + if len(p) != len(m): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = [f[r:] for f, r in zip(full_input_ids, response_token_ids_start_idx)] + answer_attention_mask = [f[r:] for f, r in zip(full_attention_mask, response_token_ids_start_idx)] + + output = dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + answer_input_ids=answer_input_ids, + answer_attention_mask=answer_attention_mask, + ) + + if embedding_tokenizer is not None: + embedding_tokenized = embedding_tokenizer(batch["prompt"], truncation=True, add_special_tokens=False) + + output.update( + { + "embedding_input_ids": embedding_tokenized["input_ids"], + "embedding_attention_mask": embedding_tokenized["attention_mask"], + } + ) + + return output + + +def _process_tokens(example: dict[str, Any], model: "PreTrainedModel" = None, **kwargs) -> dict: + """Process tokens of a BCO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + completion responses is/are too long. First we truncate the prompt; if we're still too long, we truncate the + completion. + + We also create the labels for the completion responses, which are of length equal to the sum of the length of the + prompt and the completion response, with label_pad_token_id for the prompt tokens. + """ + prompt = example["prompt"] + completion = example["completion"] + + batch = { + f"{kwargs['prefix']}prompt": prompt, + f"{kwargs['prefix']}completion": completion, + f"{kwargs['prefix']}label": example["label"], + } + + if not kwargs["is_encoder_decoder"]: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + + if not isinstance(completion, str): + raise ValueError(f"completion should be an str but got {type(completion)}") + + # keys of format prompt_* refers to just the prompt and answer_* refers to just the answer + all_tokens = { + "prompt_input_ids": example["prompt_input_ids"], + "prompt_attention_mask": example["prompt_attention_mask"], + "answer_input_ids": example["answer_input_ids"], + "answer_attention_mask": example["answer_attention_mask"], + } + + # calculate max length by checking if BOS/EOS is already there + max_length = kwargs["max_length"] + bos_token_id = kwargs["tokenizer"].bos_token_id + eos_token_id = kwargs["tokenizer"].eos_token_id + if bos_token_id != all_tokens["prompt_input_ids"][0]: + max_length -= 1 + if eos_token_id != all_tokens["answer_input_ids"][-1]: + max_length -= 1 + + # if combined sequence is too long (> max_length - 1 for BOS token - 1 for EOS), truncate the prompt + if len(all_tokens["prompt_input_ids"]) + len(all_tokens["answer_input_ids"]) > max_length: + for k in ["prompt_input_ids", "prompt_attention_mask"]: + if kwargs["truncation_mode"] == "keep_start": + all_tokens[k] = all_tokens[k][: kwargs["max_prompt_length"]] + elif kwargs["truncation_mode"] == "keep_end": + all_tokens[k] = all_tokens[k][-kwargs["max_prompt_length"] :] + else: + raise ValueError(f"Unknown truncation mode: {kwargs['truncation_mode']}") + + # if that's still too long, truncate the response + if len(all_tokens["prompt_input_ids"]) + len(all_tokens["answer_input_ids"]) > max_length: + for k in ["answer_input_ids", "answer_attention_mask"]: + all_tokens[k] = all_tokens[k][: max_length - kwargs["max_prompt_length"]] + + # all input_ids and attention mask as is. We then check if we need to add BOS/EOS tokens + batch[f"{kwargs['prefix']}prompt_input_ids"] = all_tokens["prompt_input_ids"] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = all_tokens["prompt_attention_mask"] + batch[f"{kwargs['prefix']}completion_input_ids"] = ( + all_tokens["prompt_input_ids"] + all_tokens["answer_input_ids"] + ) + batch[f"{kwargs['prefix']}completion_attention_mask"] = ( + all_tokens["prompt_attention_mask"] + all_tokens["answer_attention_mask"] + ) + + # add BOS, which affects both prompt and the full completion + if bos_token_id is not None: + if len(all_tokens["prompt_input_ids"]) == 0 or bos_token_id != all_tokens["prompt_input_ids"][0]: + batch[f"{kwargs['prefix']}prompt_input_ids"] = [bos_token_id] + batch[ + f"{kwargs['prefix']}prompt_input_ids" + ] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = [1] + batch[ + f"{kwargs['prefix']}prompt_attention_mask" + ] + batch[f"{kwargs['prefix']}completion_input_ids"] = [bos_token_id] + batch[ + f"{kwargs['prefix']}completion_input_ids" + ] + batch[f"{kwargs['prefix']}completion_attention_mask"] = [1] + batch[ + f"{kwargs['prefix']}completion_attention_mask" + ] + # add EOS, which affects only the full completion + if len(all_tokens["answer_input_ids"]) == 0 or eos_token_id != all_tokens["answer_input_ids"][-1]: + batch[f"{kwargs['prefix']}completion_input_ids"] = batch[f"{kwargs['prefix']}completion_input_ids"] + [ + eos_token_id + ] + batch[f"{kwargs['prefix']}completion_attention_mask"] = batch[ + f"{kwargs['prefix']}completion_attention_mask" + ] + [1] + + batch[f"{kwargs['prefix']}completion_labels"] = batch[f"{kwargs['prefix']}completion_input_ids"][:] + batch[f"{kwargs['prefix']}completion_labels"][: len(batch[f"{kwargs['prefix']}prompt_input_ids"])] = [ + kwargs["label_pad_token_id"] + ] * len(batch[f"{kwargs['prefix']}prompt_input_ids"]) + else: + completion_tokens = kwargs["tokenizer"]( + completion, truncation=True, max_length=kwargs["max_completion_length"], add_special_tokens=True + ) + prompt_tokens = kwargs["tokenizer"]( + prompt, truncation=True, max_length=kwargs["max_prompt_length"], add_special_tokens=True + ) + + batch[f"{kwargs['prefix']}prompt_input_ids"] = prompt_tokens["input_ids"] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = prompt_tokens["attention_mask"] + + batch[f"{kwargs['prefix']}completion_labels"] = completion_tokens["input_ids"] + batch[f"{kwargs['prefix']}completion_attention_mask"] = completion_tokens["attention_mask"] + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch[f"{kwargs['prefix']}completion_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["completion_labels"]) + ) + + return batch + + +class BCOTrainer(Trainer): + r""" + Initialize BCOTrainer from [BCO](https://huggingface.co/papers/2404.04656) paper. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForSequenceClassification`. + ref_model (`PreTrainedModelWrapper`): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + args (`BCOConfig`): + The arguments to use for training. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + data_collator (`transformers.DataCollator`, *optional*, defaults to `None`): + The data collator to use for training. If None is specified, the default data collator + (`DPODataCollatorWithPadding`) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + model_adapter_name (`str`, defaults to `None`): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str`, defaults to `None`): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + """ + + _tag_names = ["trl", "bco"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module, str] = None, + ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: BCOConfig = None, + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + data_collator: Optional[DataCollator] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional[dict] = None, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + model_adapter_name: Optional[str] = None, + ref_adapter_name: Optional[str] = None, + embedding_func: Optional[Callable] = None, + embedding_tokenizer: Optional[PreTrainedTokenizerBase] = None, + ): + if embedding_func is not None and not (is_sklearn_available() and is_joblib_available()): + raise ImportError( + "BCOTrainer with UDM requires the scikit-learn and joblib libraries. Please install it with `pip install scikit-learn joblib`." + ) + + if type(args) is TrainingArguments: + raise ValueError("Please use `BCOConfig` instead `TrainingArguments`.") + + if not isinstance(model, str) and model is not None and ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, you must mass a copy of it, or `None` if you use peft." + ) + + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the BCOTrainer. But your model is already instantiated.") + else: + model_init_kwargs = args.model_init_kwargs + torch_dtype = model_init_kwargs.get("torch_dtype") + if torch_dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(torch_dtype, str) and torch_dtype != "auto": + torch_dtype = getattr(torch, torch_dtype) + if torch_dtype != "auto" and not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"Invalid `torch_dtype` passed to the BCOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}." + ) + model_init_kwargs["torch_dtype"] = torch_dtype + + if args.ref_model_init_kwargs is None: + ref_model_init_kwargs = {} + elif not isinstance(ref_model, str): + raise ValueError( + "You passed ref_model_kwargs to the BCOTrainer. But your ref_model is already instantiated." + ) + else: + ref_model_init_kwargs = args.ref_model_init_kwargs + torch_dtype = ref_model_init_kwargs.get("torch_dtype") + if torch_dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(torch_dtype, str) and torch_dtype != "auto": + torch_dtype = getattr(torch, torch_dtype) + if torch_dtype != "auto" and not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"Invalid `torch_dtype` passed to the BCOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}." + ) + ref_model_init_kwargs["torch_dtype"] = torch_dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + + if isinstance(ref_model, str): + ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **ref_model_init_kwargs) + + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it with `pip install peft` to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." + " Please install `wandb` or `comet-ml` to resolve." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + self.is_peft_model = is_peft_available() and isinstance(model, PeftModel) + self.model_adapter_name = model_adapter_name + self.ref_adapter_name = ref_adapter_name + + if ref_model: + self.ref_model = ref_model + elif self.is_peft_model or args.precompute_ref_log_probs: + # The `model` with adapters turned off will be used as the reference model + self.ref_model = None + else: + self.ref_model = create_reference_model(model) + + if processing_class is None: + raise ValueError( + "max_length or a processing_class must be specified when using the default DPODataCollatorWithPadding" + ) + if args.max_length is None: + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `max_length` in the `BCOConfig`. " + "It will be set to `512` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_length = 512 + if args.max_length is not None: + max_length = args.max_length + + if args.max_prompt_length is None: + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `max_prompt_length` in the `BCOConfig`. " + "It will be set to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_prompt_length = 128 + if args.max_prompt_length is not None: + max_prompt_length = args.max_prompt_length + + max_completion_length = None + if args.max_completion_length is None and self.is_encoder_decoder: + warnings.warn( + "When using DPODataCollatorWithPadding with an encoder decoder architecture, you should set `max_completion_length` in the BCOTrainer's init" + " it will be set to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_completion_length = 128 + if args.max_completion_length is not None and self.is_encoder_decoder: + max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your BCOConfig" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.max_completion_length = max_completion_length + self.precompute_ref_log_probs = args.precompute_ref_log_probs + + # Since ref_logs are precomputed on the first call to get_train/eval_dataloader + # keep track of first called to avoid computation of future calls + self._precomputed_train_ref_log_probs = False + self._precomputed_eval_ref_log_probs = False + + # metric + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # BCO parameter + self.beta = args.beta + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + warnings.warn( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + UserWarning, + ) + + # Underlying Distribution Matching argument + self.embedding_func = embedding_func + self.embedding_tokenizer = embedding_tokenizer + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in BCO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids" and "completion_input_ids". As a result, + # the trainer issues the warning: "Could not estimate the number of tokens of the input, floating-point + # operations will not be computed." To suppress this warning, we set the "estimate_tokens" key in the model's + # "warnings_issued" dictionary to True. This acts as a flag to indicate that the warning has already been + # issued. + model.warnings_issued["estimate_tokens"] = True + + with PartialState().main_process_first(): + # Apply the chat template if needed + train_dataset = train_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + ) + + # Tokenize and prepare the training datasets + train_dataset = train_dataset.map( + _tokenize, + batched=True, + fn_kwargs={"tokenizer": processing_class, "embedding_tokenizer": self.embedding_tokenizer}, + num_proc=args.dataset_num_proc, + desc="Tokenizing train dataset", + ) + + # Prepare the datasets + fn_kwargs = { + "prefix": "", + "is_encoder_decoder": self.is_encoder_decoder, + "tokenizer": processing_class, + "max_length": self.max_length, + "truncation_mode": self.truncation_mode, + "label_pad_token_id": self.label_pad_token_id, + "max_prompt_length": self.max_prompt_length, + "max_completion_length": self.max_completion_length, + } + train_dataset = train_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized train dataset", + ) + + if eval_dataset is not None: + # Tokenize + eval_dataset = eval_dataset.map( + _tokenize, + fn_kwargs={"tokenizer": processing_class, "embedding_tokenizer": self.embedding_tokenizer}, + batched=True, + num_proc=args.dataset_num_proc, + desc="Tokenizing eval dataset", + ) + + # Process + fn_kwargs = { + "prefix": "", + "is_encoder_decoder": self.is_encoder_decoder, + "tokenizer": processing_class, + "max_length": self.max_length, + "truncation_mode": self.truncation_mode, + "label_pad_token_id": self.label_pad_token_id, + "max_prompt_length": self.max_prompt_length, + "max_completion_length": self.max_completion_length, + } + eval_dataset = eval_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized eval dataset", + ) + + desirable = train_dataset.filter( + lambda x: x["label"], num_proc=args.dataset_num_proc, desc="Filtering desirable examples" + ) + undesirable = train_dataset.filter( + lambda x: not x["label"], num_proc=args.dataset_num_proc, desc="Filtering undesirable examples" + ) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + # Deepspeed Zero-3 does not support precompute_ref_log_probs + if self.is_deepspeed_enabled: + if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`." + ) + + if self.ref_model is None: + if not (self.is_peft_model or self.precompute_ref_log_probs): + raise ValueError( + "No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`" + ) + else: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + self.running = RunningMoments(accelerator=self.accelerator) + + if self.embedding_func is None or args.resume_from_checkpoint: + return + + chosen_embeddings = self._get_sample_prompt_embeddings(desirable, sample_size=self.args.prompt_sample_size) + rejected_embeddings = self._get_sample_prompt_embeddings(undesirable, sample_size=self.args.prompt_sample_size) + + embeddings = torch.cat((chosen_embeddings, rejected_embeddings), dim=0) + labels = torch.cat( + (torch.ones_like(chosen_embeddings[:, 0]), torch.zeros_like(rejected_embeddings[:, 0])), dim=0 + ) + + self.clf = LogisticRegression(class_weight="balanced").fit( + embeddings.cpu().float().numpy(), labels.cpu().numpy() + ) + chosen_mean = self.clf.score( + chosen_embeddings.cpu().float().numpy(), torch.ones_like(chosen_embeddings[:, 0]).cpu().numpy() + ) + rejected_mean = self.clf.score( + rejected_embeddings.cpu().float().numpy(), torch.zeros_like(rejected_embeddings[:, 0]).cpu().numpy() + ) + logger.info(f"UDM classifier training scores: chosen: {chosen_mean}, rejected: {rejected_mean}") + + @property + def match_underlying_distribution(self): + return self.embedding_func is not None and self.embedding_tokenizer is not None + + def _get_chosen_prob(self, prompt_embeddings: torch.FloatTensor) -> torch.FloatTensor: + """ + Calculates the probability if the given prompt embedding is from desirable dataset. This function calculates + the probability in the process and ensemble across processes. + """ + dtype = prompt_embeddings.dtype + device = prompt_embeddings.device + rank = self.accelerator.process_index + + padded_prompt_embeddings = self.accelerator.pad_across_processes( + prompt_embeddings, pad_index=self.embedding_tokenizer.pad_token_id + ) + sample_size = padded_prompt_embeddings.shape[0] + nonzero = padded_prompt_embeddings.mean(dim=1) != self.embedding_tokenizer.pad_token_id + prompt_embeddings = self.accelerator.gather(padded_prompt_embeddings) + + # cannot predict for all empty values + if prompt_embeddings.shape[0] == 0: + return torch.tensor([], device=device, dtype=dtype) + + prob = self.clf.predict_proba(prompt_embeddings.cpu().float().numpy())[:, 1] + prob = torch.as_tensor(prob, dtype=dtype, device=device) + prob = self.accelerator.reduce(prob, reduction="mean") + + prob = prob[sample_size * rank : sample_size * (rank + 1)] + prob = prob[nonzero] + + return prob + + def _vectorize_prompt(self, input_ids: torch.LongTensor, attention_mask: torch.LongTensor) -> torch.FloatTensor: + """ + Replaces processing_class.pad_token_id to embedding_tokenizer.pad_token_id and applies self.embedding_func + """ + input_ids = torch.where( + input_ids == self.processing_class.pad_token_id, + self.embedding_tokenizer.pad_token_id, + input_ids, + ) + + with torch.no_grad(): + embeddings = self.embedding_func( + input_ids=input_ids, + attention_mask=attention_mask, + ) + + return embeddings + + def _get_prompt_embeddings( + self, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor]: + """Extract embeddings from frozen embedding model""" + + if not self.match_underlying_distribution: + return None, None + + embeddings = self._vectorize_prompt( + input_ids=batch["embedding_input_ids"], + attention_mask=batch["embedding_attention_mask"], + ) + + chosen_idx = [i for i in range(len(batch["label"])) if batch["label"][i] is True] + rejected_idx = [i for i in range(len(batch["label"])) if batch["label"][i] is False] + + chosen_embeddings = embeddings[chosen_idx, ...] + rejected_embeddings = embeddings[rejected_idx, ...] + + return (chosen_embeddings, rejected_embeddings) + + def _get_sample_prompt_embeddings(self, dataset: Dataset, sample_size: int = 512) -> torch.FloatTensor: + """ + Sample instances from dataset and get prompt embeddings. Used for density ratio classifier training. + """ + n_samples = min(len(dataset), sample_size) + rand_indices = np.random.choice(len(dataset), size=(n_samples,)) + + embedding_dataset = dataset.select(rand_indices) + + dataloader_params = { + "batch_size": self.args.per_device_train_batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(embedding_dataset, **dataloader_params)) + + with torch.no_grad(): + all_embeddings = torch.empty(0) + for padded_batch in tqdm(iterable=data_loader, desc="Building sample prompt embeddings"): + embeddings = self._vectorize_prompt( + input_ids=padded_batch["embedding_input_ids"], + attention_mask=padded_batch["embedding_attention_mask"], + ) + embeddings = self.accelerator.gather_for_metrics(embeddings) + all_embeddings = torch.cat((all_embeddings, embeddings.cpu())) + + return all_embeddings + + def _save_optimizer_and_scheduler(self, output_dir): + output_dir = output_dir if output_dir is not None else self.args.output_dir + super()._save_optimizer_and_scheduler(output_dir) + + if self.accelerator.is_main_process: + # When saving optimizer and scheduler to checkpoint, save also the running delta object. + self.running.save_to_json(os.path.join(output_dir, RUNNING_NAME)) + + if self.match_underlying_distribution: + joblib.dump(self.clf, os.path.join(output_dir, CLF_NAME), compress=True) + + def _load_optimizer_and_scheduler(self, checkpoint): + if checkpoint is None: + logger.warning_once(f"Missing Checkpoint {checkpoint}") + return + + super()._load_optimizer_and_scheduler(checkpoint) + + # when loading optimizer and scheduler from checkpoint, also load the running delta object. + running_file = os.path.join(checkpoint, RUNNING_NAME) + if os.path.isfile(running_file): + self.running = RunningMoments.load_from_json(self.accelerator, running_file) + + if self.match_underlying_distribution: + clf_file = os.path.join(checkpoint, CLF_NAME) + if os.path.isfile(clf_file): + self.clf = joblib.load(clf_file) + + @contextmanager + def null_ref_context(self): + """Context manager for handling null reference model (that is, peft adapter manipulation).""" + with ( + self.accelerator.unwrap_model(self.model).disable_adapter() + if self.is_peft_model and not self.ref_adapter_name + else nullcontext() + ): + if self.ref_adapter_name: + self.model.set_adapter(self.ref_adapter_name) + yield + if self.ref_adapter_name: + self.model.set_adapter(self.model_adapter_name or "default") + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`. + """ + + if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs: + dataloader_params = { + "batch_size": self.args.per_device_train_batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params)) + reference_completion_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Train dataset reference log probs"): + reference_completion_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + self.train_dataset = self.train_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + + self._precomputed_train_ref_log_probs = True + + return super().get_train_dataloader() + + def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader: + """ + Returns the evaluation [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`. + + Args: + eval_dataset (`torch.utils.data.Dataset`, *optional*): + If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted + by the `model.forward()` method are automatically removed. It must implement `__len__`. + """ + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset + + if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs: + dataloader_params = { + "batch_size": self.args.per_device_eval_batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params)) + + reference_completion_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Eval dataset reference log probs"): + reference_completion_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + eval_dataset = eval_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + + # Save calculated reference_chosen_logps and reference_rejected_logps to the eval_dataset for subsequent runs + if self.eval_dataset is not None: + self.eval_dataset = eval_dataset + self._precomputed_eval_ref_log_probs = True + + return super().get_eval_dataloader(eval_dataset=eval_dataset) + + def compute_reference_log_probs(self, padded_batch: dict) -> dict: + """Computes log probabilities of the reference model for a single padded batch of a BCO specific dataset.""" + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + if self.is_encoder_decoder: + completion_logits = self.model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + else: + completion_logits = self.model( + padded_batch["completion_input_ids"], + attention_mask=padded_batch["completion_attention_mask"], + ).logits + + else: + if self.is_encoder_decoder: + completion_logits = self.ref_model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + else: + completion_logits = self.ref_model( + padded_batch["completion_input_ids"], attention_mask=padded_batch["completion_attention_mask"] + ).logits + + completion_logps = self.get_batch_logps( + completion_logits, + padded_batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + return completion_logps + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + else: + # Fixes end-dec RuntimeError + labels = labels.clone() + + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels[labels == label_pad_token_id] = 0 + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + model_kwargs = ( + { + "labels": batch["completion_labels"], + "decoder_input_ids": batch.get("completion_decoder_input_ids"), + } + if self.is_encoder_decoder + else {} + ) + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + **model_kwargs, + ) + completion_logits = outputs.logits + + completion_logps = self.get_batch_logps( + completion_logits, + batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + if completion_logps.shape[0] != len(batch["label"]): + raise ValueError( + "There is a mismatch between the number of examples in this batch and the number of " + "examples for which an output sequence was predicted." + ) + + chosen_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is False] + + chosen_logps = completion_logps[chosen_idx, ...] + rejected_logps = completion_logps[rejected_idx, ...] + + chosen_logits = completion_logits[chosen_idx, ...] + rejected_logits = completion_logits[rejected_idx, ...] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, outputs.aux_loss) + else: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits) + + def _get_udm_weight(self, rejected_embeddings: torch.FloatTensor) -> torch.FloatTensor: + prob_desirable = self._get_chosen_prob(rejected_embeddings) + min_ratio = self.args.min_density_ratio + max_ratio = self.args.max_density_ratio + + weight = (prob_desirable / (1 - prob_desirable + 1e-8)).clamp(min=min_ratio, max=max_ratio) + + return weight + + def bco_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + reference_chosen_logps: torch.FloatTensor, + reference_rejected_logps: torch.FloatTensor, + chosen_embeddings: Optional[torch.FloatTensor], + rejected_embeddings: Optional[torch.FloatTensor], + do_train: bool = True, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute the BCO loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (num(chosen) in batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (num(rejected) in batch_size,) + reference_chosen_logps: + Log probabilities of the reference model for the chosen responses. Shape: (num(chosen) in batch_size,) + reference_rejected_logps: + Log probabilities of the reference model for the rejected responses. Shape: (num(rejected) in + batch_size,) + chosen_embeddings: embeddings of desirable prompts + rejected_embeddings: embeddings of undesirable prompts + + Returns: + A tuple of four tensors: (losses, chosen_rewards, rejected_rewards, delta). The losses tensor contains the + BCO loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards + for the chosen and rejected responses, respectively. The delta value contains the moving average of all + implicit rewards. + """ + + chosen_logratios = policy_chosen_logps - reference_chosen_logps + chosen_rewards = self.beta * chosen_logratios + + rejected_logratios = policy_rejected_logps - reference_rejected_logps + rejected_rewards = self.beta * rejected_logratios + + if do_train: + self.running.update(torch.cat((chosen_rewards, rejected_rewards), 0).detach()) + delta = torch.as_tensor(self.running.mean, device=chosen_rewards.device) + + chosen_losses = -F.logsigmoid(chosen_rewards - delta) + rejected_losses = -F.logsigmoid(-(rejected_rewards - delta)) + + if self.match_underlying_distribution: + chosen_weight = torch.ones_like(chosen_losses) + rejected_weight = self._get_udm_weight(rejected_embeddings) + + losses = torch.cat((chosen_weight * chosen_losses, rejected_weight * rejected_losses), dim=0) + else: + losses = torch.cat((chosen_losses, rejected_losses), dim=0) + + return losses, chosen_rewards, rejected_rewards, delta + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, Union[list, torch.LongTensor]], + do_train: bool = True, + ): + """Compute the BCO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} + + forward_output = self.forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + ) = forward_output[:4] + if self.aux_loss_enabled: + aux_loss = forward_output[4] + + # if reference_logps in batch use them, otherwise use the reference model + if "reference_logps" in batch: + chosen_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is False] + + reference_chosen_logps = batch["reference_logps"][chosen_idx, ...] + reference_rejected_logps = batch["reference_logps"][rejected_idx, ...] + else: + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + ) = self.forward(self.model, batch)[:4] + else: + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + ) = self.forward(self.ref_model, batch)[:4] + + chosen_embeddings, rejected_embeddings = self._get_prompt_embeddings(batch) + + losses, chosen_rewards, rejected_rewards, delta = self.bco_loss( + policy_chosen_logps, + policy_rejected_logps, + reference_chosen_logps, + reference_rejected_logps, + chosen_embeddings, + rejected_embeddings, + do_train=do_train, + ) + metrics["delta"] = self.accelerator.gather_for_metrics(delta).mean().item() + + num_chosen = torch.Tensor([len(chosen_rewards)]).to(self.accelerator.device) + num_rejected = torch.Tensor([len(rejected_rewards)]).to(self.accelerator.device) + + all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() + all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() + + if all_num_chosen > 0: + metrics["rewards/chosen_sum"] = ( + self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() + ) + metrics["logps/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logps.nansum()).nansum().item() + ) + metrics["logits/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logits.nansum()).nansum().item() + ) + metrics["count/chosen"] = all_num_chosen + + if all_num_rejected > 0: + metrics["rewards/rejected_sum"] = ( + self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() + ) + metrics["logps/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logps.nansum()).nansum().item() + ) + metrics["logits/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logits.nansum()).nansum().item() + ) + metrics["count/rejected"] = all_num_rejected + + loss = losses.nanmean() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, metrics + + def compute_loss( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + return_outputs=False, + num_items_in_batch=None, + ) -> Union[torch.Tensor, tuple[torch.Tensor, dict[str, torch.Tensor]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs) + + # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: + loss = loss.to(self.args.device) + # force log the metrics + if self.accelerator.is_main_process: + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def _get_train_sampler(self, dataset: Optional[Dataset] = None) -> Optional[torch.utils.data.Sampler]: + if dataset is None: + dataset = self.train_dataset + if dataset is None or not has_length(dataset): + return None + return SequentialSampler(dataset) + + def generate_from_model_and_ref(self, model, batch: dict[str, torch.LongTensor]) -> tuple[str, str]: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + # if reference_output in batch use that otherwise use the reference model + if "reference_output" in batch: + reference_output = batch["reference_output"] + else: + if self.ref_model is None: + with self.null_ref_context(): + reference_output = self.model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + else: + reference_output = self.ref_model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + reference_output = pad_to_length(reference_output, self.max_length, self.processing_class.pad_token_id) + reference_output_decoded = self.processing_class.batch_decode(reference_output, skip_special_tokens=True) + + return policy_output_decoded, reference_output_decoded + + def prediction_step( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[list[str]] = None, + ): + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, do_train=False) + + # force log the metrics + if self.accelerator.is_main_process: + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return (loss.detach(), None, None) + + # logits for the chosen and rejected samples from model + logits_dict = {} + if "logits/chosen_sum" in metrics: + logits_dict["eval_logits/chosen"] = metrics["logits/chosen_sum"] + if "logits/rejected_sum" in metrics: + logits_dict["eval_logits/rejected"] = metrics["logits/rejected_sum"] + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: Optional[bool] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + target_indicies = [i for i in range(len(random_batch["label"])) if random_batch["label"][i] is False] + target_batch = { + "prompt_input_ids": random_batch["prompt_input_ids"][target_indicies], + "prompt_attention_mask": random_batch["prompt_attention_mask"][target_indicies], + "prompt": itemgetter(*target_indicies)(random_batch["prompt"]), + } + policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, target_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy", "Ref Model"], + data=[ + [prompt, pol[len(prompt) :], ref[len(prompt) :]] + for prompt, pol, ref in zip(target_batch["prompt"], policy_output_decoded, ref_output_decoded) + ], + ) + if "wandb" in self.args.report_to: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float` or `None`, *optional*, defaults to `None`): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # train metrics should have no prefix, eval should have 'eval_' + prefix = "eval_" if train_eval == "eval" else "" + # accumulate average metrics from sums and lengths + for split in ["chosen", "rejected"]: + if f"count/{split}" in self._stored_metrics[train_eval]: + count_sum = torch.Tensor(self._stored_metrics[train_eval][f"count/{split}"]).sum().item() + for metric in ["rewards", "logps", "logits"]: + logs[f"{prefix}{metric}/{split}"] = ( + torch.Tensor(self._stored_metrics[train_eval][f"{metric}/{split}_sum"]).sum().item() + / count_sum + ) + # delete obsolete metric + del self._stored_metrics[train_eval][f"{metric}/{split}_sum"] + del self._stored_metrics[train_eval][f"count/{split}"] + # calculate reward margin + if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs: + logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"] + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[f"{prefix}{key}"] = torch.Tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{jung2024binary, + title = {{Binary Classifier Optimization for Large Language Model Alignment}}, + author = {Seungjae Jung and Gunsoo Han and Daniel Wontae Nam and Kyoung{-}Woon On}, + year = 2024, + eprint = {arXiv:2404.04656} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="BCO", + trainer_citation=citation, + paper_title="Binary Classifier Optimization for Large Language Model Alignment", + paper_id="2404.04656", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/callbacks.py b/lib/python3.12/site-packages/trl/trainer/callbacks.py new file mode 100644 index 0000000000000000000000000000000000000000..2a02d8bf04f7cc158c984d8021f53eb96e46a7f9 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/callbacks.py @@ -0,0 +1,572 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from typing import Optional, Union + +import pandas as pd +import torch +from accelerate import Accelerator +from accelerate.state import AcceleratorState +from accelerate.utils import gather_object, is_wandb_available +from transformers import ( + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + Trainer, + TrainerCallback, + TrainerControl, + TrainerState, + TrainingArguments, +) +from transformers.trainer_utils import has_length +from transformers.utils import is_rich_available + +from ..data_utils import maybe_apply_chat_template +from ..import_utils import is_mergekit_available +from ..mergekit_utils import MergeConfig, merge_models, upload_model_to_hf +from ..models.utils import unwrap_model_for_generation +from .judges import BasePairwiseJudge +from .utils import log_table_to_comet_experiment + + +if is_rich_available(): + from rich.console import Console, Group + from rich.live import Live + from rich.panel import Panel + from rich.progress import Progress + +if is_wandb_available(): + import wandb + + +def _generate_completions( + prompts: list[str], + model: PreTrainedModel, + tokenizer: PreTrainedTokenizerBase, + accelerator: Accelerator, + generation_config: Optional[GenerationConfig], + batch_size: int = 1, +) -> list[str]: + """ + Generates completions for a list of pre-formatted prompts from the given model. + + Args: + prompts (list[str]): A list of input prompts for which completions are to be generated. + model (PreTrainedModel): The pre-trained model to be used for generation. + tokenizer (PreTrainedTokenizerBase): The tokenizer to be used for encoding and decoding. + accelerator (Accelerator): The accelerator to be used for model execution. + generation_config (GenerationConfig): Configuration for text generation. + batch_size (int, optional): The number of prompts to process in each batch. Default is 1. + + Returns: + list[str]: A list of generated text completions corresponding to the input prompts. + """ + completions = [] + with unwrap_model_for_generation(model, accelerator) as unwrapped_model: + for idx in range(0, len(prompts), batch_size): + batch = prompts[idx : idx + batch_size] + tokenized_batch = tokenizer(batch, return_tensors="pt", padding=True, truncation=True).to(model.device) + generations = unwrapped_model.generate( + **tokenized_batch, + generation_config=generation_config, + ) + for prompt, generation in zip(tokenized_batch.input_ids, generations): + # Remove prompt from generation + generation = generation[len(prompt) :] + completion = tokenizer.decode(generation, skip_special_tokens=True) + completions.append(completion) + return completions + + +class SyncRefModelCallback(TrainerCallback): + """ + Callback to synchronize the model with a reference model. + """ + + def __init__( + self, + ref_model: Union[PreTrainedModel, torch.nn.Module], + accelerator: Optional[Accelerator], + ): + self.accelerator = accelerator + self.ref_model = ref_model + + @staticmethod + def _sync_target_model(model, target_model, alpha): + for target_param, copy_param in zip(target_model.parameters(), model.parameters()): + target_param.data.mul_(1.0 - alpha).add_(copy_param.data, alpha=alpha) + + @staticmethod + def sync_target_model(model, target_model, alpha): + deepspeed_plugin = AcceleratorState().deepspeed_plugin + if deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3: + import deepspeed + + with deepspeed.zero.GatheredParameters( + list(model.parameters()) + list(target_model.parameters()), modifier_rank=0 + ): + if deepspeed.comm.get_rank() == 0: + SyncRefModelCallback._sync_target_model(model, target_model, alpha) + else: + SyncRefModelCallback._sync_target_model(model, target_model, alpha) + + def on_step_end(self, args, state, control, **kwargs): + model: PreTrainedModel = kwargs["model"] + + if self.ref_model is not None and state.global_step % args.ref_model_sync_steps == 0: + if self.accelerator: + model = self.accelerator.unwrap_model(model) + self.sync_target_model(model, self.ref_model, args.ref_model_mixup_alpha) + + +class RichProgressCallback(TrainerCallback): + """ + A [`TrainerCallback`] that displays the progress of training or evaluation using Rich. + """ + + def __init__(self): + if not is_rich_available(): + raise ImportError("RichProgressCallback requires the `rich` extra. To install, run `pip install rich`.") + + self.training_bar = None + self.prediction_bar = None + + self.training_task_id = None + self.prediction_task_id = None + + self.rich_group = None + self.rich_console = None + + self.training_status = None + self.current_step = None + + def on_train_begin(self, args, state, control, **kwargs): + if state.is_world_process_zero: + self.training_bar = Progress() + self.prediction_bar = Progress() + + self.rich_console = Console() + + self.training_status = self.rich_console.status("Nothing to log yet ...") + + self.rich_group = Live(Panel(Group(self.training_bar, self.prediction_bar, self.training_status))) + self.rich_group.start() + + self.training_task_id = self.training_bar.add_task("[blue]Training the model", total=state.max_steps) + self.current_step = 0 + + def on_step_end(self, args, state, control, **kwargs): + if state.is_world_process_zero: + self.training_bar.update(self.training_task_id, advance=state.global_step - self.current_step, update=True) + self.current_step = state.global_step + + def on_prediction_step(self, args, state, control, eval_dataloader=None, **kwargs): + if state.is_world_process_zero and has_length(eval_dataloader): + if self.prediction_task_id is None: + self.prediction_task_id = self.prediction_bar.add_task( + "[blue]Predicting on the evaluation dataset", total=len(eval_dataloader) + ) + self.prediction_bar.update(self.prediction_task_id, advance=1, update=True) + + def on_evaluate(self, args, state, control, **kwargs): + if state.is_world_process_zero: + if self.prediction_task_id is not None: + self.prediction_bar.remove_task(self.prediction_task_id) + self.prediction_task_id = None + + def on_predict(self, args, state, control, **kwargs): + if state.is_world_process_zero: + if self.prediction_task_id is not None: + self.prediction_bar.remove_task(self.prediction_task_id) + self.prediction_task_id = None + + def on_log(self, args, state, control, logs=None, **kwargs): + if state.is_world_process_zero and self.training_bar is not None: + _ = logs.pop("total_flos", None) + self.training_status.update(f"[bold green]Status = {str(logs)}") + + def on_train_end(self, args, state, control, **kwargs): + if state.is_world_process_zero: + self.rich_group.stop() + + self.training_bar = None + self.prediction_bar = None + self.training_task_id = None + self.prediction_task_id = None + self.rich_group = None + self.rich_console = None + self.training_status = None + self.current_step = None + + +def _win_rate_completions_df( + state: TrainerState, prompts: list[str], completions: list[str], winner_indices: list[str] +) -> pd.DataFrame: + global_step = [str(state.global_step)] * len(prompts) + data = list(zip(global_step, prompts, completions, winner_indices)) + # Split completions from reference model and policy + split_data = [(item[0], item[1], item[2][0], item[2][1], item[3]) for item in data] + return pd.DataFrame(split_data, columns=["step", "prompt", "reference_model", "policy", "winner_index"]) + + +class WinRateCallback(TrainerCallback): + """ + A [`~transformers.TrainerCallback`] that computes the win rate of a model based on a reference. + + It generates completions using prompts from the evaluation dataset and compares the trained model's outputs against + a reference. The reference is either the initial version of the model (before training) or the reference model, if + available in the trainer. During each evaluation step, a judge determines how often the trained model's completions + win against the reference using a judge. The win rate is then logged in the trainer's logs under the key + `"eval_win_rate"`. + + Usage: + ```python + trainer = DPOTrainer(...) + judge = PairRMJudge() + win_rate_callback = WinRateCallback(judge=judge, trainer=trainer) + trainer.add_callback(win_rate_callback) + ``` + + Args: + judge (`BasePairwiseJudge`): + The judge to use for comparing completions. + trainer (`Trainer`): + Trainer to which the callback will be attached. The trainer's evaluation dataset must include a `"prompt"` + column containing the prompts for generating completions. If the `Trainer` has a reference model (via the + `ref_model` attribute), it will use this reference model for generating the reference completions; + otherwise, it defaults to using the initial model. + generation_config (`GenerationConfig`, *optional*): + The generation config to use for generating completions. + num_prompts (`int` or `None`, *optional*, defaults to `None`): + The number of prompts to generate completions for. If not provided, defaults to the number of examples in + the evaluation dataset. + shuffle_order (`bool`, *optional*, defaults to `True`): + Whether to shuffle the order of the completions before judging. + use_soft_judge (`bool`, *optional*, defaults to `False`): + Whether to use a soft judge that returns a win probability between 0 and 1 for the first completion vs the + second. + """ + + def __init__( + self, + judge: BasePairwiseJudge, + trainer: Trainer, + generation_config: Optional[GenerationConfig] = None, + num_prompts: Optional[int] = None, + shuffle_order: bool = True, + use_soft_judge: bool = False, + ): + self.judge = judge + self.trainer = trainer + self.shuffle_order = shuffle_order + self.generation_config = generation_config + self.ref_completions = [] + self.use_soft_judge = use_soft_judge + + if self.trainer.eval_dataset is None: + raise ValueError("Trainer must have an evaluation dataset to use the WinRateCallback.") + else: + self.eval_dataset = self.trainer.eval_dataset + + if num_prompts is not None: + self.eval_dataset = self.eval_dataset.select(range(num_prompts)) + + def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + # When the trainer is initialized, we generate completions for the reference model. + tokenizer = kwargs["processing_class"] + tokenizer.padding_side = "left" + accelerator = self.trainer.accelerator + # Use the reference model if available, otherwise use the initial model + model = getattr(self.trainer, "ref_model", None) + # At this point, there are two cases where `ref_model` is None: + # 1. The method doesn't require a reference model. + # 2. The method uses a reference model, but `ref_model` is set to None. + # This occurs when using PEFT, where the reference model can be obtained by simply disabling the model's adapter. + # In theory, we should disable the adapter here, but since it's zero-initialized at the start of training, + # the model behaves identically with or without the adapter. + # Therefore, there's no need to explicitly disable it at this point. + if model is None: + model = self.trainer.model_wrapped + with accelerator.split_between_processes(self.eval_dataset["prompt"]) as prompts: + self.ref_completions = _generate_completions( + prompts, + model=model, + tokenizer=tokenizer, + accelerator=accelerator, + generation_config=self.generation_config, + batch_size=args.per_device_eval_batch_size, + ) + # Compute initial win rate as a reference point + completions = list(zip(self.ref_completions, self.ref_completions)) + if self.use_soft_judge: + ref_win_probs = self.judge.judge(prompts, completions, self.shuffle_order, return_scores=True) + winner_indices = [0 if score > 0.5 else 1 for score in ref_win_probs] + ref_win_probs = gather_object(ref_win_probs) + else: + winner_indices = self.judge.judge(prompts, completions, self.shuffle_order) + prompts = gather_object(prompts) + completions = gather_object(completions) + winner_indices = gather_object(winner_indices) + + # Logging + if self.trainer.accelerator.is_main_process: + win_rate = sum(winner_idx == 1 for winner_idx in winner_indices) / len(winner_indices) + if self.use_soft_judge: + avg_win_prob = 1.0 - sum(ref_win_probs) / len(ref_win_probs) + self.trainer.log({"eval_avg_win_prob": avg_win_prob, "eval_win_rate": win_rate}) + else: + self.trainer.log({"eval_win_rate": win_rate}) + + if "wandb" in args.report_to: + import wandb + + if wandb.run is not None: + df = _win_rate_completions_df( + state=state, + prompts=prompts, + completions=completions, + winner_indices=winner_indices, + ) + wandb.log({"win_rate_completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + df = _win_rate_completions_df( + state=state, + prompts=prompts, + completions=completions, + winner_indices=winner_indices, + ) + log_table_to_comet_experiment( + name="win_rate_completions.csv", + table=df, + ) + + def on_evaluate(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + # At every evaluation step, we generate completions for the model and compare them with the reference + # completions that have been generated at the beginning of training. We then compute the win rate and log it to + # the trainer. + tokenizer = kwargs["processing_class"] + tokenizer.padding_side = "left" + accelerator = self.trainer.accelerator + model = self.trainer.model_wrapped + with accelerator.split_between_processes(self.eval_dataset["prompt"]) as prompts: + completions = _generate_completions( + prompts, + model=model, + tokenizer=tokenizer, + accelerator=accelerator, + generation_config=self.generation_config, + batch_size=args.per_device_eval_batch_size, + ) + + completions = list(zip(self.ref_completions, completions)) + + if self.use_soft_judge: + ref_win_probs = self.judge.judge(prompts, completions, self.shuffle_order, return_scores=True) + winner_indices = [0 if score > 0.5 else 1 for score in ref_win_probs] + ref_win_probs = gather_object(ref_win_probs) + else: + winner_indices = self.judge.judge(prompts, completions, self.shuffle_order) + prompts = gather_object(prompts) + completions = gather_object(completions) + winner_indices = gather_object(winner_indices) + + # Logging + if self.trainer.accelerator.is_main_process: + win_rate = sum(winner_idx == 1 for winner_idx in winner_indices) / len(winner_indices) + if self.use_soft_judge: + avg_win_prob = 1.0 - sum(ref_win_probs) / len(ref_win_probs) + self.trainer.log({"eval_avg_win_prob": avg_win_prob, "eval_win_rate": win_rate}) + else: + self.trainer.log({"eval_win_rate": win_rate}) + + if "wandb" in args.report_to: + import wandb + + if wandb.run is not None: + df = _win_rate_completions_df( + state=state, + prompts=prompts, + completions=completions, + winner_indices=winner_indices, + ) + wandb.log({"win_rate_completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + df = _win_rate_completions_df( + state=state, + prompts=prompts, + completions=completions, + winner_indices=winner_indices, + ) + log_table_to_comet_experiment( + name="win_rate_completions.csv", + table=df, + ) + + +class LogCompletionsCallback(TrainerCallback): + r""" + A [`~transformers.TrainerCallback`] that logs completions to Weights & Biases and/or Comet. + + Usage: + ```python + trainer = DPOTrainer(...) + completions_callback = LogCompletionsCallback(trainer=trainer) + trainer.add_callback(completions_callback) + ``` + + Args: + trainer (`Trainer`): + Trainer to which the callback will be attached. The trainer's evaluation dataset must include a `"prompt"` + column containing the prompts for generating completions. + generation_config (`GenerationConfig`, *optional*): + The generation config to use for generating completions. + num_prompts (`int` or `None`, *optional*): + The number of prompts to generate completions for. If not provided, defaults to the number of examples in + the evaluation dataset. + freq (`int` or `None`, *optional*): + The frequency at which to log completions. If not provided, defaults to the trainer's `eval_steps`. + """ + + def __init__( + self, + trainer: Trainer, + generation_config: Optional[GenerationConfig] = None, + num_prompts: Optional[int] = None, + freq: Optional[int] = None, + ): + self.trainer = trainer + self.generation_config = generation_config + self.freq = freq + self.table = [] + self._last_logged_step = -1 + + if self.trainer.eval_dataset is None: + raise ValueError("Trainer must have an evaluation dataset to use the LogCompletionsCallback.") + else: + self.eval_dataset = self.trainer.eval_dataset + + if num_prompts is not None: + self.eval_dataset = self.eval_dataset.select(range(num_prompts)) + + def on_step_end(self, args, state, control, **kwargs): + # Only log once per step (this method may be called multiple times) + if state.global_step == self._last_logged_step: + return + + # Only log every `freq` steps (if no `freq` is provided, log every `eval_steps` steps) + freq = self.freq or state.eval_steps + if state.global_step % freq != 0: + return + + tokenizer = kwargs["processing_class"] + tokenizer.padding_side = "left" + accelerator = self.trainer.accelerator + model = self.trainer.model_wrapped + with accelerator.split_between_processes(self.eval_dataset["prompt"]) as prompts: + prompts = [maybe_apply_chat_template({"prompt": prompt}, tokenizer)["prompt"] for prompt in prompts] + completions = _generate_completions( + prompts, + model=model, + tokenizer=tokenizer, + accelerator=accelerator, + generation_config=self.generation_config, + batch_size=args.per_device_eval_batch_size, + ) + completions = gather_object(completions) + prompts = gather_object(prompts) + + # Build the data to log + if self.trainer.accelerator.is_main_process: + global_step = [str(state.global_step)] * len(prompts) + data = list(zip(global_step, prompts, completions)) + self.table.extend(data) + table = pd.DataFrame(columns=["step", "prompt", "completion"], data=self.table) + + if "wandb" in args.report_to: + wandb.log({"completions": table}) + + if "comet_ml" in args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=table, + ) + + # Save the last logged step, so we don't log the same completions multiple times + self._last_logged_step = state.global_step + + +class MergeModelCallback(TrainerCallback): + r""" + A [`~transformers.TrainerCallback`] that merges the policy model (the model being trained) with another model based + on a merge configuration. + + Args: + merge_config ([`MergeConfig`], *optional*, defaults to `None`): + Configuration used for the merging process. If not provided, the default [`MergeConfig`] is used. + merge_at_every_checkpoint (`bool`, *optional*, defaults to `False`): + Whether to merge the model at every checkpoint. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the merged model to the Hub after merging. + + Example: + + ```python + # pip install mergekit + + from trl.mergekit_utils import MergeConfig + from trl import MergeModelCallback + + config = MergeConfig() + merge_callback = MergeModelCallback(config) + trainer = DPOTrainer(..., callbacks=[merge_callback]) + ``` + """ + + def __init__( + self, + merge_config: Optional["MergeConfig"] = None, + merge_at_every_checkpoint: bool = False, + push_to_hub: bool = False, + ): + if not is_mergekit_available(): + raise ImportError( + "MergeModelCallback requires the `mergekit` extra. To install, run `pip install mergekit`." + ) + self.merge_config = merge_config or MergeConfig() + self.merge_at_every_checkpoint = merge_at_every_checkpoint + self.push_to_hub = push_to_hub + + def _merge_and_maybe_push(self, output_dir, global_step, model): + checkpoint_path = os.path.join(output_dir, f"checkpoint-{global_step}") + self.merge_config.policy_model_path = checkpoint_path + if self.merge_config.target_model_path is None: + self.merge_config.target_model_path = model.config._name_or_path + merge_path = os.path.join(checkpoint_path, "merged") + + merge_models(self.merge_config.create(), merge_path) + + if self.push_to_hub: + repo_name = f"{output_dir}_checkpoint-{global_step}_merged" + upload_model_to_hf(merge_path, repo_name) + + def on_save(self, args, state, control, model=None, **kwargs): + if self.merge_at_every_checkpoint: + self._merge_and_maybe_push(args.output_dir, state.global_step, model) + + def on_train_end(self, args, state, control, model=None, **kwargs): + if not self.merge_at_every_checkpoint: + self._merge_and_maybe_push(args.output_dir, state.global_step, model) diff --git a/lib/python3.12/site-packages/trl/trainer/cpo_config.py b/lib/python3.12/site-packages/trl/trainer/cpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..6493ef45c402038471f3e69149dbcd4f2dd366db --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/cpo_config.py @@ -0,0 +1,192 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Any, Optional + +from transformers import TrainingArguments + + +@dataclass +class CPOConfig(TrainingArguments): + r""" + Configuration class for the [`CPOTrainer`]. + + This class includes only the parameters that are specific to CPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in + the [paper](https://huggingface.co/papers/2310.12036). + label_smoothing (`float`, *optional*, defaults to `0.0`): + Label smoothing factor. This argument is required if you want to use the default data collator. + loss_type (`str`, *optional*, defaults to `"sigmoid"`): + Type of loss to use. Possible values are: + + - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. + - `"hinge"`: hinge loss on the normalized likelihood from the + [SLiC](https://huggingface.co/papers/2305.10425) paper. + - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. + - `"simpo"`: SimPO loss from the [SimPO](https://huggingface.co/papers/2405.14734) paper. + + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + cpo_alpha (`float`, *optional*, defaults to `1.0`): + Weight of the BC regularizer in CPO training. + simpo_gamma (`float`, *optional*, defaults to `0.5`): + Target reward margin for the SimPO loss, used only when the `loss_type="simpo"`. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`,*optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from the model to W&B or Comet during evaluation. + is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + max_length: Optional[int] = field( + default=1024, + metadata={"help": "Maximum length of the sequences (prompt + completion) in the batch."}, + ) + max_prompt_length: Optional[int] = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + max_completion_length: Optional[int] = field( + default=None, + metadata={ + "help": "Maximum length of the completion. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the deviation from the reference model. Higher β means less deviation from " + "the reference model." + }, + ) + label_smoothing: float = field( + default=0.0, + metadata={"help": "Label smoothing factor."}, + ) + loss_type: str = field( + default="sigmoid", + metadata={ + "help": "Type of loss to use.", + "choices": ["sigmoid", "hinge", "ipo", "simpo"], + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + cpo_alpha: float = field( + default=1.0, + metadata={"help": "Weight of the BC regularizer in CPO training."}, + ) + simpo_gamma: float = field( + default=0.5, + metadata={"help": "Target reward margin for the SimPO loss, used only when the `loss_type='simpo'`."}, + ) + label_pad_token_id: int = field( + default=-100, + metadata={"help": "Label pad token id."}, + ) + padding_value: Optional[int] = field( + default=None, + metadata={"help": "Padding value to use. If `None`, the padding value of the tokenizer is used."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the prompt is too long.", + "choices": ["keep_end", "keep_start"], + }, + ) + generate_during_eval: bool = field( + default=False, + metadata={"help": "If `True`, generates and logs completions from the model to W&B during evaluation."}, + ) + is_encoder_decoder: Optional[bool] = field( + default=None, + metadata={"help": "Whether the model is an encoder-decoder model."}, + ) + model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model " + "from a string." + }, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/lib/python3.12/site-packages/trl/trainer/cpo_trainer.py b/lib/python3.12/site-packages/trl/trainer/cpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..215a7cf92e70eeb3b9bf2aaab53e3bcada4e3b42 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/cpo_trainer.py @@ -0,0 +1,1115 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import random +import textwrap +import warnings +from collections import defaultdict +from contextlib import nullcontext +from pathlib import Path +from typing import Any, Callable, Literal, Optional, Union + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState +from datasets import Dataset +from torch import autocast +from torch.utils.data import DataLoader +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + is_comet_available, + is_wandb_available, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalLoopOutput +from transformers.utils import is_peft_available, is_torch_fx_proxy + +from ..data_utils import maybe_apply_chat_template, maybe_extract_prompt +from .cpo_config import CPOConfig +from .utils import ( + DPODataCollatorWithPadding, + add_bos_token_if_needed, + add_eos_token_if_needed, + disable_dropout_in_model, + generate_model_card, + get_comet_experiment_url, + log_table_to_comet_experiment, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + + +if is_wandb_available(): + import wandb + + +class CPOTrainer(Trainer): + r""" + Initialize CPOTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForSequenceClassification`. + args (`CPOConfig`): + The CPO config arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`DPODataCollatorWithPadding`) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + """ + + _tag_names = ["trl", "cpo"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: Optional[CPOConfig] = None, + data_collator: Optional[DataCollator] = None, + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional[dict] = None, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + ): + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the CPOTrainer. But your model is already instantiated.") + else: + model_init_kwargs = args.model_init_kwargs + torch_dtype = model_init_kwargs.get("torch_dtype") + if torch_dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(torch_dtype, str) and torch_dtype != "auto": + torch_dtype = getattr(torch, torch_dtype) + if torch_dtype != "auto" and not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"Invalid `torch_dtype` passed to the CPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}." + ) + model_init_kwargs["torch_dtype"] = torch_dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." + " Please install `wandb` or `comet-ml` to resolve." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + if self.is_encoder_decoder: + self.decoder_start_token_id = model.config.decoder_start_token_id + self.pad_token_id = model.config.pad_token_id + + if processing_class is None: + raise ValueError("processing_class must be specified to tokenize a CPO dataset.") + if args.max_length is None: + warnings.warn( + "`max_length` is not set in the CPOConfig's init" + " it will default to `512` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_length = 512 + else: + max_length = args.max_length + if args.max_prompt_length is None: + warnings.warn( + "`max_prompt_length` is not set in the CPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_prompt_length = 128 + else: + max_prompt_length = args.max_prompt_length + + if not max_prompt_length < max_length: + raise ValueError( + f"max_prompt_length ({max_prompt_length}) should be strictly less than max_length ({max_length})." + ) + + if args.max_completion_length is None and self.is_encoder_decoder: + warnings.warn( + "When using an encoder decoder architecture, you should set `max_completion_length` in the CPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_completion_length = 128 + else: + max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.max_completion_length = max_completion_length + self.processing_class = processing_class + + if args.loss_type in ["hinge", "ipo"] and args.label_smoothing > 0: + warnings.warn( + f"You are using the {args.loss_type} loss type that does not support label smoothing. The " + "`label_smoothing` parameter will be ignored. Set `label_smoothing` to `0.0` to remove this warning.", + UserWarning, + ) + if args.loss_type == "kto_pair": + raise ValueError("Support for kto_pair has been removed in CPOTrainer. Please use KTOTrainer.") + + self.beta = args.beta + self.label_smoothing = args.label_smoothing + self.loss_type = args.loss_type + self.cpo_alpha = args.cpo_alpha + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + warnings.warn( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + UserWarning, + ) + + if args.loss_type == "simpo": + self.simpo_gamma = args.simpo_gamma + + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in CPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and + # "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens + # of the input, floating-point operations will not be computed." To suppress this warning, we set the + # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate + # that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed, and apply the chat template if needed + train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + train_dataset = train_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + ) + + # tokenize the dataset + train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + def build_tokenized_answer(self, prompt, answer): + """ + Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a + + b)[len(enc(a)):]`. Reference: + https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + """ + + full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False) + prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"] + + answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :] + answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids]) + + # Prepare input tokens for token by token comparison + full_input_ids = np.array(full_tokenized["input_ids"]) + + if len(full_input_ids) != len(full_concat_input_ids): + raise ValueError("Prompt input ids and answer input ids should have the same length.") + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = len(prompt_input_ids) + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]: + response_token_ids_start_idx -= 1 + + prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx] + prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx] + + if len(prompt_input_ids) != len(prompt_attention_mask): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:] + answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:] + + return dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + input_ids=answer_input_ids, + attention_mask=answer_attention_mask, + ) + + def tokenize_row(self, feature, model: Optional[Union[PreTrainedModel, nn.Module]] = None) -> dict: + """Tokenize a single row from a CPO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long, + we truncate the chosen/rejected. + + We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length + of the prompt and the chosen/rejected response, with label_pad_token_id for the prompt tokens. + """ + batch = {} + prompt = feature["prompt"] + chosen = feature["chosen"] + rejected = feature["rejected"] + + if not self.is_encoder_decoder: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + prompt_tokens = self.processing_class(prompt, add_special_tokens=False) + prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} + + if not isinstance(chosen, str): + raise ValueError(f"chosen should be an str but got {type(chosen)}") + chosen_tokens = self.build_tokenized_answer(prompt, chosen) + + if not isinstance(rejected, str): + raise ValueError(f"rejected should be an str but got {type(rejected)}") + rejected_tokens = self.build_tokenized_answer(prompt, rejected) + + # Last prompt token might get merged by tokenizer and + # it should not be included for generation if that happens + prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) + + chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) + rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) + prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids) + + for k, v in prompt_tokens.items(): + prompt_tokens[k] = v[:prompt_len_input_ids] + + # Make sure prompts only have one different token at most an + # and length only differs by 1 at most + num_diff_tokens = sum( + [a != b for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"])] + ) + num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids) + if num_diff_tokens > 1 or num_diff_len > 1: + raise ValueError( + "Chosen and rejected prompt_input_ids might only differ on the " + "last token due to tokenizer merge ops." + ) + + # add BOS token to head of prompt. Avoid adding if it's already there + prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( + self.processing_class.bos_token_id, + prompt_len_input_ids, + prompt_tokens, + chosen_prompt_len_input_ids, + chosen_tokens, + rejected_prompt_len_input_ids, + rejected_tokens, + ) + + # add EOS token to end of answer. Avoid adding if it's already there + chosen_tokens, rejected_tokens = add_eos_token_if_needed( + self.processing_class.eos_token_id, chosen_tokens, rejected_tokens + ) + + longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"])) + + # if combined sequence is too long, truncate the prompt + for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + if self.truncation_mode == "keep_start": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_prompt_length] + elif self.truncation_mode == "keep_end": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :] + else: + raise ValueError(f"Unknown truncation mode: {self.truncation_mode}") + + # if that's still too long, truncate the response + for answer_tokens in [chosen_tokens, rejected_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + for k in ["input_ids", "attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length] + + # Create labels + chosen_sequence_tokens = { + k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"] + } + rejected_sequence_tokens = { + k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"] + } + chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:] + chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(chosen_tokens["prompt_input_ids"]) + rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:] + rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(rejected_tokens["prompt_input_ids"]) + + for k, toks in { + "chosen_": chosen_sequence_tokens, + "rejected_": rejected_sequence_tokens, + "": prompt_tokens, + }.items(): + for type_key, tokens in toks.items(): + if type_key == "token_type_ids": + continue + batch[f"{k}{type_key}"] = tokens + + else: + chosen_tokens = self.processing_class( + chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + rejected_tokens = self.processing_class( + rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + prompt_tokens = self.processing_class( + prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True + ) + + batch["chosen_labels"] = chosen_tokens["input_ids"] + batch["rejected_labels"] = rejected_tokens["input_ids"] + batch["prompt_input_ids"] = prompt_tokens["input_ids"] + batch["prompt_attention_mask"] = prompt_tokens["attention_mask"] + + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["rejected_labels"]) + ) + batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["chosen_labels"]) + ) + + return batch + + @staticmethod + def concatenated_inputs( + batch: dict[str, Union[list, torch.LongTensor]], + is_encoder_decoder: bool = False, + label_pad_token_id: int = -100, + padding_value: int = 0, + device: Optional[torch.device] = None, + ) -> dict[str, torch.LongTensor]: + """Concatenate the chosen and rejected inputs into a single tensor. + + Args: + batch: + A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors + of shape (batch_size, sequence_length). + is_encoder_decoder: + Whether the model is an encoder-decoder model. + label_pad_token_id: + The label pad token id. + padding_value: + The padding value to use for the concatenated inputs_ids. + device: + The device for the concatenated inputs. + + Returns: + A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'. + """ + concatenated_batch = {} + + if is_encoder_decoder: + max_length = max(batch["chosen_labels"].shape[1], batch["rejected_labels"].shape[1]) + else: + max_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) + + for k in batch: + if k.startswith("chosen") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("chosen", "concatenated") + concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value) + for k in batch: + if k.startswith("rejected") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("rejected", "concatenated") + concatenated_batch[concatenated_key] = torch.cat( + ( + concatenated_batch[concatenated_key], + pad_to_length(batch[k], max_length, pad_value=pad_value), + ), + dim=0, + ).to(device=device) + + if is_encoder_decoder: + concatenated_batch["concatenated_input_ids"] = batch["prompt_input_ids"].repeat(2, 1).to(device=device) + concatenated_batch["concatenated_attention_mask"] = ( + batch["prompt_attention_mask"].repeat(2, 1).to(device=device) + ) + + return concatenated_batch + + def cpo_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute the CPO loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) + + Returns: + A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the CPO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. + """ + logits = (policy_chosen_logps - policy_rejected_logps).to(self.accelerator.device) + + # The beta is a temperature parameter for the CPO loss, typically something in the range of 0.1 to 0.5. + # We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the labels and + # calculates a conservative CPO loss. + + if self.loss_type == "simpo": + gamma_logratios = self.simpo_gamma / self.beta + logits = logits - gamma_logratios + # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0. + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) + elif self.loss_type == "sigmoid": + # This reduces to Equation 3 from the CPO paper when label_smoothing -> 0. + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) + elif self.loss_type == "hinge": + losses = torch.relu(1 - self.beta * logits) + elif self.loss_type == "ipo": + # eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper. + losses = (logits - 1 / (2 * self.beta)) ** 2 + else: + raise ValueError( + f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'simpo']" + ) + + chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach() + rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach() + + return losses, chosen_rewards, rejected_rewards + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + label_pad_token_id: The label pad token id. + is_encoder_decoder: Whether the model is an encoder-decoder model. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels[labels == label_pad_token_id] = 0 + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def concatenated_forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. + + We do this to avoid doing two forward passes, because it's faster for FSDP. + """ + concatenated_batch = self.concatenated_inputs( + batch, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + padding_value=self.padding_value, + device=self.accelerator.device, + ) + len_chosen = batch["chosen_labels"].shape[0] + + model_kwargs = ( + { + "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]), + } + if self.is_encoder_decoder + else {} + ) + + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + concatenated_batch["concatenated_input_ids"], + attention_mask=concatenated_batch["concatenated_attention_mask"], + use_cache=False, + **model_kwargs, + ) + all_logits = outputs.logits + + def cross_entropy_loss(logits, labels): + if not self.is_encoder_decoder: + # Shift so that tokens < n predict n + logits = logits[..., :-1, :].contiguous() + labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + logits = logits.view(-1, logits.shape[-1]) + labels = labels.view(-1) + # Enable model parallelism + labels = labels.to(logits.device) + loss = loss_fct(logits, labels) + return loss + + labels = concatenated_batch["concatenated_labels"].clone() + + if self.cpo_alpha == 0: + nll_loss = torch.tensor(0.0).to(self.accelerator.device) + else: + nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen]) + + all_logps = self.get_batch_logps( + all_logits, + concatenated_batch["concatenated_labels"], + average_log_prob=self.loss_type in ["ipo", "simpo"], + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + chosen_logps = all_logps[:len_chosen] + rejected_logps = all_logps[len_chosen:] + + chosen_logits = all_logits[:len_chosen] + rejected_logits = all_logits[len_chosen:] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss, outputs.aux_loss) + + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, nll_loss) + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, Union[list, torch.LongTensor]], + train_eval: Literal["train", "eval"] = "train", + ): + """Compute the CPO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + + forward_output = self.concatenated_forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_nll_loss, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + losses, chosen_rewards, rejected_rewards = self.cpo_loss( + policy_chosen_logps, + policy_rejected_logps, + ) + + loss = losses.mean() + self.cpo_alpha * policy_nll_loss + reward_accuracies = (chosen_rewards > rejected_rewards).float() + + prefix = "eval_" if train_eval == "eval" else "" + metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean().item() + metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean().item() + metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean().item() + metrics[f"{prefix}rewards/margins"] = ( + self.accelerator.gather_for_metrics(chosen_rewards - rejected_rewards).mean().item() + ) + metrics[f"{prefix}logps/rejected"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean().item() + ) + metrics[f"{prefix}logps/chosen"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean().item() + ) + metrics[f"{prefix}logits/rejected"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logits.detach().mean()).mean().item() + ) + metrics[f"{prefix}logits/chosen"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logits.detach().mean()).mean().item() + ) + metrics[f"{prefix}nll_loss"] = self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean().item() + + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, metrics + + def compute_loss( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + return_outputs=False, + num_items_in_batch=None, + ) -> Union[torch.Tensor, tuple[torch.Tensor, dict[str, torch.Tensor]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") + + # force log the metrics + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + def generate_from_model(self, model, batch: dict[str, torch.LongTensor]) -> str: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + return policy_output_decoded + + def prediction_step( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[list[str]] = None, + ): + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval") + + # force log the metrics + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return (loss.detach(), None, None) + + # logits for the chosen and rejected samples from model + logits_dict = { + "eval_logits/chosen": metrics["eval_logits/chosen"], + "eval_logits/rejected": metrics["eval_logits/rejected"], + } + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: Optional[bool] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + policy_output_decoded = self.generate_from_model(self.model, random_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy"], + data=[ + [prompt, pol[len(prompt) :]] for prompt, pol in zip(random_batch["prompt"], policy_output_decoded) + ], + ) + if "wandb" in self.args.report_to: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float` or `None`, *optional*, defaults to `None`): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[key] = torch.tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + def _shift_right(self, input_ids): + if self.decoder_start_token_id is None: + raise ValueError( + "model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id." + ) + + # shift inputs to the right + if is_torch_fx_proxy(input_ids): + # Item assignment is not supported natively for proxies. + shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id) + shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) + else: + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = self.decoder_start_token_id + + if self.pad_token_id is None: + raise ValueError("model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id) + + return shifted_input_ids + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{xu2024contrastive, + title = {{Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation}}, + author = {Haoran Xu and Amr Sharaf and Yunmo Chen and Weiting Tan and Lingfeng Shen and Benjamin Van Durme and Kenton Murray and Young Jin Kim}, + year = 2024, + booktitle = {Forty-first International Conference on Machine Learning, {ICML} 2024, Vienna, Austria, July 21-27, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=51iwkioZpn} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="CPO", + trainer_citation=citation, + paper_title="Contrastive Preference Optimization: Pushing the Boundaries of LLM Performance in Machine Translation", + paper_id="2401.08417", + ) + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/ddpo_config.py b/lib/python3.12/site-packages/trl/trainer/ddpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..5852e8582c0eab0438d45b7c20afeab2f1bfbc5e --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/ddpo_config.py @@ -0,0 +1,299 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import sys +from dataclasses import dataclass, field +from typing import Optional + +from transformers import is_bitsandbytes_available + +from ..core import flatten_dict + + +@dataclass +class DDPOConfig: + r""" + Configuration class for the [`DDPOTrainer`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + exp_name (`str`, *optional*, defaults to `os.path.basename(sys.argv[0])[: -len(".py")]`): + Name of this experiment (by default is the file name without the extension name). + run_name (`str`, *optional*, defaults to `""`): + Name of this run. + seed (`int`, *optional*, defaults to `0`): + Random seed. + log_with (`Literal["wandb", "tensorboard"]]` or `None`, *optional*, defaults to `None`): + Log with either 'wandb' or 'tensorboard', check + https://huggingface.co/docs/accelerate/usage_guides/tracking for more details. + tracker_kwargs (`Dict`, *optional*, defaults to `{}`): + Keyword arguments for the tracker (e.g. wandb_project). + accelerator_kwargs (`Dict`, *optional*, defaults to `{}`): + Keyword arguments for the accelerator. + project_kwargs (`Dict`, *optional*, defaults to `{}`): + Keyword arguments for the accelerator project config (e.g. `logging_dir`). + tracker_project_name (`str`, *optional*, defaults to `"trl"`): + Name of project to use for tracking. + logdir (`str`, *optional*, defaults to `"logs"`): + Top-level logging directory for checkpoint saving. + num_epochs (`int`, *optional*, defaults to `100`): + Number of epochs to train. + save_freq (`int`, *optional*, defaults to `1`): + Number of epochs between saving model checkpoints. + num_checkpoint_limit (`int`, *optional*, defaults to `5`): + Number of checkpoints to keep before overwriting old ones. + mixed_precision (`str`, *optional*, defaults to `"fp16"`): + Mixed precision training. + allow_tf32 (`bool`, *optional*, defaults to `True`): + Allow `tf32` on Ampere GPUs. + resume_from (`str`, *optional*, defaults to `""`): + Resume training from a checkpoint. + sample_num_steps (`int`, *optional*, defaults to `50`): + Number of sampler inference steps. + sample_eta (`float`, *optional*, defaults to `1.0`): + Eta parameter for the DDIM sampler. + sample_guidance_scale (`float`, *optional*, defaults to `5.0`): + Classifier-free guidance weight. + sample_batch_size (`int`, *optional*, defaults to `1`): + Batch size (per GPU) to use for sampling. + sample_num_batches_per_epoch (`int`, *optional*, defaults to `2`): + Number of batches to sample per epoch. + train_batch_size (`int`, *optional*, defaults to `1`): + Batch size (per GPU) to use for training. + train_use_8bit_adam (`bool`, *optional*, defaults to `False`): + Use 8bit Adam optimizer from bitsandbytes. + train_learning_rate (`float`, *optional*, defaults to `3e-4`): + Learning rate. + train_adam_beta1 (`float`, *optional*, defaults to `0.9`): + Adam beta1. + train_adam_beta2 (`float`, *optional*, defaults to `0.999`): + Adam beta2. + train_adam_weight_decay (`float`, *optional*, defaults to `1e-4`): + Adam weight decay. + train_adam_epsilon (`float`, *optional*, defaults to `1e-8`): + Adam epsilon. + train_gradient_accumulation_steps (`int`, *optional*, defaults to `1`): + Number of gradient accumulation steps. + train_max_grad_norm (`float`, *optional*, defaults to `1.0`): + Maximum gradient norm for gradient clipping. + train_num_inner_epochs (`int`, *optional*, defaults to `1`): + Number of inner epochs per outer epoch. + train_cfg (`bool`, *optional*, defaults to `True`): + Whether to use classifier-free guidance during training. + train_adv_clip_max (`float`, *optional*, defaults to `5.0`): + Clip advantages to the range. + train_clip_range (`float`, *optional*, defaults to `1e-4`): + PPO clip range. + train_timestep_fraction (`float`, *optional*, defaults to `1.0`): + Fraction of timesteps to train on. + per_prompt_stat_tracking (`bool`, *optional*, defaults to `False`): + Whether to track statistics for each prompt separately. + per_prompt_stat_tracking_buffer_size (`int`, *optional*, defaults to `16`): + Number of reward values to store in the buffer for each prompt. + per_prompt_stat_tracking_min_count (`int`, *optional*, defaults to `16`): + Minimum number of reward values to store in the buffer. + async_reward_computation (`bool`, *optional*, defaults to `False`): + Whether to compute rewards asynchronously. + max_workers (`int`, *optional*, defaults to `2`): + Maximum number of workers to use for async reward computation. + negative_prompts (`str`, *optional*, defaults to `""`): + Comma-separated list of prompts to use as negative examples. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the final model checkpoint to the Hub. + """ + + exp_name: str = field( + default=os.path.basename(sys.argv[0])[: -len(".py")], + metadata={"help": "Name of this experiment (by default is the file name without the extension name)."}, + ) + run_name: str = field( + default="", + metadata={"help": "Name of this run."}, + ) + seed: int = field( + default=0, + metadata={"help": "Random seed."}, + ) + log_with: Optional[str] = field( + default=None, + metadata={ + "help": "Log with either 'wandb' or 'tensorboard'.", + "choices": ["wandb", "tensorboard"], + }, + ) + tracker_kwargs: dict = field( + default_factory=dict, + metadata={"help": "Keyword arguments for the tracker (e.g. wandb_project)."}, + ) + accelerator_kwargs: dict = field( + default_factory=dict, + metadata={"help": "Keyword arguments for the accelerator."}, + ) + project_kwargs: dict = field( + default_factory=dict, + metadata={"help": "Keyword arguments for the accelerator project config (e.g. `logging_dir`)."}, + ) + tracker_project_name: str = field( + default="trl", + metadata={"help": "Name of project to use for tracking."}, + ) + logdir: str = field( + default="logs", + metadata={"help": "Top-level logging directory for checkpoint saving."}, + ) + num_epochs: int = field( + default=100, + metadata={"help": "Number of epochs to train."}, + ) + save_freq: int = field( + default=1, + metadata={"help": "Number of epochs between saving model checkpoints."}, + ) + num_checkpoint_limit: int = field( + default=5, + metadata={"help": "Number of checkpoints to keep before overwriting old ones."}, + ) + mixed_precision: str = field( + default="fp16", + metadata={"help": "Mixed precision training."}, + ) + allow_tf32: bool = field( + default=True, + metadata={"help": "Allow `tf32` on Ampere GPUs."}, + ) + resume_from: str = field( + default="", + metadata={"help": "Resume training from a checkpoint."}, + ) + sample_num_steps: int = field( + default=50, + metadata={"help": "Number of sampler inference steps."}, + ) + sample_eta: float = field( + default=1.0, + metadata={"help": "Eta parameter for the DDIM sampler."}, + ) + sample_guidance_scale: float = field( + default=5.0, + metadata={"help": "Classifier-free guidance weight."}, + ) + sample_batch_size: int = field( + default=1, + metadata={"help": "Batch size (per GPU) to use for sampling."}, + ) + sample_num_batches_per_epoch: int = field( + default=2, + metadata={"help": "Number of batches to sample per epoch."}, + ) + train_batch_size: int = field( + default=1, + metadata={"help": "Batch size (per GPU) to use for training."}, + ) + train_use_8bit_adam: bool = field( + default=False, + metadata={"help": "Use 8bit Adam optimizer from bitsandbytes."}, + ) + train_learning_rate: float = field( + default=3e-4, + metadata={"help": "Learning rate."}, + ) + train_adam_beta1: float = field( + default=0.9, + metadata={"help": "Adam beta1."}, + ) + train_adam_beta2: float = field( + default=0.999, + metadata={"help": "Adam beta2."}, + ) + train_adam_weight_decay: float = field( + default=1e-4, + metadata={"help": "Adam weight decay."}, + ) + train_adam_epsilon: float = field( + default=1e-8, + metadata={"help": "Adam epsilon."}, + ) + train_gradient_accumulation_steps: int = field( + default=1, + metadata={"help": "Number of gradient accumulation steps."}, + ) + train_max_grad_norm: float = field( + default=1.0, + metadata={"help": "Maximum gradient norm for gradient clipping."}, + ) + train_num_inner_epochs: int = field( + default=1, + metadata={"help": "Number of inner epochs per outer epoch."}, + ) + train_cfg: bool = field( + default=True, + metadata={"help": "Whether to use classifier-free guidance during training."}, + ) + train_adv_clip_max: float = field( + default=5.0, + metadata={"help": "Clip advantages to the range."}, + ) + train_clip_range: float = field( + default=1e-4, + metadata={"help": "PPO clip range."}, + ) + train_timestep_fraction: float = field( + default=1.0, + metadata={"help": "Fraction of timesteps to train on."}, + ) + per_prompt_stat_tracking: bool = field( + default=False, + metadata={"help": "Whether to track statistics for each prompt separately."}, + ) + per_prompt_stat_tracking_buffer_size: int = field( + default=16, + metadata={"help": "Number of reward values to store in the buffer for each prompt."}, + ) + per_prompt_stat_tracking_min_count: int = field( + default=16, + metadata={"help": "Minimum number of reward values to store in the buffer."}, + ) + async_reward_computation: bool = field( + default=False, + metadata={"help": "Whether to compute rewards asynchronously."}, + ) + max_workers: int = field( + default=2, + metadata={"help": "Maximum number of workers to use for async reward computation."}, + ) + negative_prompts: str = field( + default="", + metadata={"help": "Comma-separated list of prompts to use as negative examples."}, + ) + push_to_hub: bool = field( + default=False, + metadata={"help": "Whether to push the final model checkpoint to the Hub."}, + ) + + def to_dict(self): + output_dict = {} + for key, value in self.__dict__.items(): + output_dict[key] = value + return flatten_dict(output_dict) + + def __post_init__(self): + if self.train_use_8bit_adam and not is_bitsandbytes_available(): + raise ImportError( + "You need to install bitsandbytes to use 8bit Adam. " + "You can install it with `pip install bitsandbytes`." + ) diff --git a/lib/python3.12/site-packages/trl/trainer/ddpo_trainer.py b/lib/python3.12/site-packages/trl/trainer/ddpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..f4d1d97c160b38868828d469fc9d4799c7c09f26 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/ddpo_trainer.py @@ -0,0 +1,669 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import textwrap +from collections import defaultdict +from concurrent import futures +from pathlib import Path +from typing import Any, Callable, Optional, Union +from warnings import warn + +import torch +from accelerate import Accelerator +from accelerate.logging import get_logger +from accelerate.utils import ProjectConfiguration, set_seed +from huggingface_hub import PyTorchModelHubMixin +from transformers import is_wandb_available + +from ..models import DDPOStableDiffusionPipeline +from .ddpo_config import DDPOConfig +from .utils import PerPromptStatTracker, generate_model_card, get_comet_experiment_url + + +if is_wandb_available(): + import wandb + + +logger = get_logger(__name__) + + +class DDPOTrainer(PyTorchModelHubMixin): + """ + The DDPOTrainer uses Deep Diffusion Policy Optimization to optimise diffusion models. Note, this trainer is heavily + inspired by the work here: https://github.com/kvablack/ddpo-pytorch As of now only Stable Diffusion based pipelines + are supported + + Attributes: + **config** (`DDPOConfig`) -- Configuration object for DDPOTrainer. Check the documentation of `PPOConfig` for more: + details. + **reward_function** (Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor]) -- Reward function to be used: + **prompt_function** (Callable[[], tuple[str, Any]]) -- Function to generate prompts to guide model + **sd_pipeline** (`DDPOStableDiffusionPipeline`) -- Stable Diffusion pipeline to be used for training. + **image_samples_hook** (Optional[Callable[[Any, Any, Any], Any]]) -- Hook to be called to log images + """ + + _tag_names = ["trl", "ddpo"] + + def __init__( + self, + config: DDPOConfig, + reward_function: Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor], + prompt_function: Callable[[], tuple[str, Any]], + sd_pipeline: DDPOStableDiffusionPipeline, + image_samples_hook: Optional[Callable[[Any, Any, Any], Any]] = None, + ): + if image_samples_hook is None: + warn("No image_samples_hook provided; no images will be logged") + + self.prompt_fn = prompt_function + self.reward_fn = reward_function + self.config = config + self.image_samples_callback = image_samples_hook + + accelerator_project_config = ProjectConfiguration(**self.config.project_kwargs) + + if self.config.resume_from: + self.config.resume_from = os.path.normpath(os.path.expanduser(self.config.resume_from)) + if "checkpoint_" not in os.path.basename(self.config.resume_from): + # get the most recent checkpoint in this directory + checkpoints = list( + filter( + lambda x: "checkpoint_" in x, + os.listdir(self.config.resume_from), + ) + ) + if len(checkpoints) == 0: + raise ValueError(f"No checkpoints found in {self.config.resume_from}") + checkpoint_numbers = sorted([int(x.split("_")[-1]) for x in checkpoints]) + self.config.resume_from = os.path.join( + self.config.resume_from, + f"checkpoint_{checkpoint_numbers[-1]}", + ) + + accelerator_project_config.iteration = checkpoint_numbers[-1] + 1 + + # number of timesteps within each trajectory to train on + self.num_train_timesteps = int(self.config.sample_num_steps * self.config.train_timestep_fraction) + + self.accelerator = Accelerator( + log_with=self.config.log_with, + mixed_precision=self.config.mixed_precision, + project_config=accelerator_project_config, + # we always accumulate gradients across timesteps; we want config.train.gradient_accumulation_steps to be the + # number of *samples* we accumulate across, so we need to multiply by the number of training timesteps to get + # the total number of optimizer steps to accumulate across. + gradient_accumulation_steps=self.config.train_gradient_accumulation_steps * self.num_train_timesteps, + **self.config.accelerator_kwargs, + ) + + is_okay, message = self._config_check() + if not is_okay: + raise ValueError(message) + + is_using_tensorboard = config.log_with is not None and config.log_with == "tensorboard" + + if self.accelerator.is_main_process: + self.accelerator.init_trackers( + self.config.tracker_project_name, + config=dict(ddpo_trainer_config=config.to_dict()) if not is_using_tensorboard else config.to_dict(), + init_kwargs=self.config.tracker_kwargs, + ) + + logger.info(f"\n{config}") + + set_seed(self.config.seed, device_specific=True) + + self.sd_pipeline = sd_pipeline + + self.sd_pipeline.set_progress_bar_config( + position=1, + disable=not self.accelerator.is_local_main_process, + leave=False, + desc="Timestep", + dynamic_ncols=True, + ) + + # For mixed precision training we cast all non-trainable weights (vae, non-lora text_encoder and non-lora unet) to half-precision + # as these weights are only used for inference, keeping weights in full precision is not required. + if self.accelerator.mixed_precision == "fp16": + inference_dtype = torch.float16 + elif self.accelerator.mixed_precision == "bf16": + inference_dtype = torch.bfloat16 + else: + inference_dtype = torch.float32 + + self.sd_pipeline.vae.to(self.accelerator.device, dtype=inference_dtype) + self.sd_pipeline.text_encoder.to(self.accelerator.device, dtype=inference_dtype) + self.sd_pipeline.unet.to(self.accelerator.device, dtype=inference_dtype) + + trainable_layers = self.sd_pipeline.get_trainable_layers() + + self.accelerator.register_save_state_pre_hook(self._save_model_hook) + self.accelerator.register_load_state_pre_hook(self._load_model_hook) + + # Enable TF32 for faster training on Ampere GPUs, + # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices + if self.config.allow_tf32: + torch.backends.cuda.matmul.allow_tf32 = True + + self.optimizer = self._setup_optimizer( + trainable_layers.parameters() if not isinstance(trainable_layers, list) else trainable_layers + ) + + self.neg_prompt_embed = self.sd_pipeline.text_encoder( + self.sd_pipeline.tokenizer( + [""] if self.config.negative_prompts is None else self.config.negative_prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.sd_pipeline.tokenizer.model_max_length, + ).input_ids.to(self.accelerator.device) + )[0] + + if config.per_prompt_stat_tracking: + self.stat_tracker = PerPromptStatTracker( + config.per_prompt_stat_tracking_buffer_size, + config.per_prompt_stat_tracking_min_count, + ) + + # NOTE: for some reason, autocast is necessary for non-lora training but for lora training it isn't necessary and it uses + # more memory + self.autocast = self.sd_pipeline.autocast or self.accelerator.autocast + + if hasattr(self.sd_pipeline, "use_lora") and self.sd_pipeline.use_lora: + unet, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer) + self.trainable_layers = list(filter(lambda p: p.requires_grad, unet.parameters())) + else: + self.trainable_layers, self.optimizer = self.accelerator.prepare(trainable_layers, self.optimizer) + + if self.config.async_reward_computation: + self.executor = futures.ThreadPoolExecutor(max_workers=config.max_workers) + + if config.resume_from: + logger.info(f"Resuming from {config.resume_from}") + self.accelerator.load_state(config.resume_from) + self.first_epoch = int(config.resume_from.split("_")[-1]) + 1 + else: + self.first_epoch = 0 + + def compute_rewards(self, prompt_image_pairs, is_async=False): + if not is_async: + rewards = [] + for images, prompts, prompt_metadata in prompt_image_pairs: + reward, reward_metadata = self.reward_fn(images, prompts, prompt_metadata) + rewards.append( + ( + torch.as_tensor(reward, device=self.accelerator.device), + reward_metadata, + ) + ) + else: + rewards = self.executor.map(lambda x: self.reward_fn(*x), prompt_image_pairs) + rewards = [ + (torch.as_tensor(reward.result(), device=self.accelerator.device), reward_metadata.result()) + for reward, reward_metadata in rewards + ] + + return zip(*rewards) + + def step(self, epoch: int, global_step: int): + """ + Perform a single step of training. + + Args: + epoch (int): The current epoch. + global_step (int): The current global step. + + Side Effects: + - Model weights are updated + - Logs the statistics to the accelerator trackers. + - If `self.image_samples_callback` is not None, it will be called with the prompt_image_pairs, global_step, + and the accelerator tracker. + + Returns: + global_step (int): The updated global step. + + """ + samples, prompt_image_data = self._generate_samples( + iterations=self.config.sample_num_batches_per_epoch, + batch_size=self.config.sample_batch_size, + ) + + # collate samples into dict where each entry has shape (num_batches_per_epoch * sample.batch_size, ...) + samples = {k: torch.cat([s[k] for s in samples]) for k in samples[0].keys()} + rewards, rewards_metadata = self.compute_rewards( + prompt_image_data, is_async=self.config.async_reward_computation + ) + + for i, image_data in enumerate(prompt_image_data): + image_data.extend([rewards[i], rewards_metadata[i]]) + + if self.image_samples_callback is not None: + self.image_samples_callback(prompt_image_data, global_step, self.accelerator.trackers[0]) + + rewards = torch.cat(rewards) + rewards = self.accelerator.gather(rewards).cpu().numpy() + + self.accelerator.log( + { + "reward": rewards, + "epoch": epoch, + "reward_mean": rewards.mean(), + "reward_std": rewards.std(), + }, + step=global_step, + ) + + if self.config.per_prompt_stat_tracking: + # gather the prompts across processes + prompt_ids = self.accelerator.gather(samples["prompt_ids"]).cpu().numpy() + prompts = self.sd_pipeline.tokenizer.batch_decode(prompt_ids, skip_special_tokens=True) + advantages = self.stat_tracker.update(prompts, rewards) + else: + advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-8) + + # ungather advantages; keep the entries corresponding to the samples on this process + samples["advantages"] = ( + torch.as_tensor(advantages) + .reshape(self.accelerator.num_processes, -1)[self.accelerator.process_index] + .to(self.accelerator.device) + ) + + del samples["prompt_ids"] + + total_batch_size, num_timesteps = samples["timesteps"].shape + + for inner_epoch in range(self.config.train_num_inner_epochs): + # shuffle samples along batch dimension + perm = torch.randperm(total_batch_size, device=self.accelerator.device) + samples = {k: v[perm] for k, v in samples.items()} + + # shuffle along time dimension independently for each sample + # still trying to understand the code below + perms = torch.stack( + [torch.randperm(num_timesteps, device=self.accelerator.device) for _ in range(total_batch_size)] + ) + + for key in ["timesteps", "latents", "next_latents", "log_probs"]: + samples[key] = samples[key][ + torch.arange(total_batch_size, device=self.accelerator.device)[:, None], + perms, + ] + + original_keys = samples.keys() + original_values = samples.values() + # rebatch them as user defined train_batch_size is different from sample_batch_size + reshaped_values = [v.reshape(-1, self.config.train_batch_size, *v.shape[1:]) for v in original_values] + + # Transpose the list of original values + transposed_values = zip(*reshaped_values) + # Create new dictionaries for each row of transposed values + samples_batched = [dict(zip(original_keys, row_values)) for row_values in transposed_values] + + self.sd_pipeline.unet.train() + global_step = self._train_batched_samples(inner_epoch, epoch, global_step, samples_batched) + # ensure optimization step at the end of the inner epoch + if not self.accelerator.sync_gradients: + raise ValueError( + "Optimization step should have been performed by this point. Please check calculated gradient accumulation settings." + ) + + if epoch != 0 and epoch % self.config.save_freq == 0 and self.accelerator.is_main_process: + self.accelerator.save_state() + + return global_step + + def calculate_loss(self, latents, timesteps, next_latents, log_probs, advantages, embeds): + """ + Calculate the loss for a batch of an unpacked sample + + Args: + latents (torch.Tensor): + The latents sampled from the diffusion model, shape: [batch_size, num_channels_latents, height, width] + timesteps (torch.Tensor): + The timesteps sampled from the diffusion model, shape: [batch_size] + next_latents (torch.Tensor): + The next latents sampled from the diffusion model, shape: [batch_size, num_channels_latents, height, + width] + log_probs (torch.Tensor): + The log probabilities of the latents, shape: [batch_size] + advantages (torch.Tensor): + The advantages of the latents, shape: [batch_size] + embeds (torch.Tensor): + The embeddings of the prompts, shape: [2*batch_size or batch_size, ...] Note: the "or" is because if + train_cfg is True, the expectation is that negative prompts are concatenated to the embeds + + Returns: + loss (torch.Tensor), approx_kl (torch.Tensor), clipfrac (torch.Tensor) (all of these are of shape (1,)) + """ + with self.autocast(): + if self.config.train_cfg: + noise_pred = self.sd_pipeline.unet( + torch.cat([latents] * 2), + torch.cat([timesteps] * 2), + embeds, + ).sample + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.config.sample_guidance_scale * ( + noise_pred_text - noise_pred_uncond + ) + else: + noise_pred = self.sd_pipeline.unet( + latents, + timesteps, + embeds, + ).sample + # compute the log prob of next_latents given latents under the current model + + scheduler_step_output = self.sd_pipeline.scheduler_step( + noise_pred, + timesteps, + latents, + eta=self.config.sample_eta, + prev_sample=next_latents, + ) + + log_prob = scheduler_step_output.log_probs + + advantages = torch.clamp( + advantages, + -self.config.train_adv_clip_max, + self.config.train_adv_clip_max, + ) + + ratio = torch.exp(log_prob - log_probs) + + loss = self.loss(advantages, self.config.train_clip_range, ratio) + + approx_kl = 0.5 * torch.mean((log_prob - log_probs) ** 2) + + clipfrac = torch.mean((torch.abs(ratio - 1.0) > self.config.train_clip_range).float()) + + return loss, approx_kl, clipfrac + + def loss( + self, + advantages: torch.Tensor, + clip_range: float, + ratio: torch.Tensor, + ): + unclipped_loss = -advantages * ratio + clipped_loss = -advantages * torch.clamp( + ratio, + 1.0 - clip_range, + 1.0 + clip_range, + ) + return torch.mean(torch.maximum(unclipped_loss, clipped_loss)) + + def _setup_optimizer(self, trainable_layers_parameters): + if self.config.train_use_8bit_adam: + import bitsandbytes + + optimizer_cls = bitsandbytes.optim.AdamW8bit + else: + optimizer_cls = torch.optim.AdamW + + return optimizer_cls( + trainable_layers_parameters, + lr=self.config.train_learning_rate, + betas=(self.config.train_adam_beta1, self.config.train_adam_beta2), + weight_decay=self.config.train_adam_weight_decay, + eps=self.config.train_adam_epsilon, + ) + + def _save_model_hook(self, models, weights, output_dir): + self.sd_pipeline.save_checkpoint(models, weights, output_dir) + weights.pop() # ensures that accelerate doesn't try to handle saving of the model + + def _load_model_hook(self, models, input_dir): + self.sd_pipeline.load_checkpoint(models, input_dir) + models.pop() # ensures that accelerate doesn't try to handle loading of the model + + def _generate_samples(self, iterations, batch_size): + """ + Generate samples from the model + + Args: + iterations (int): Number of iterations to generate samples for + batch_size (int): Batch size to use for sampling + + Returns: + samples (list[dict[str, torch.Tensor]]), prompt_image_pairs (list[list[Any]]) + """ + samples = [] + prompt_image_pairs = [] + self.sd_pipeline.unet.eval() + + sample_neg_prompt_embeds = self.neg_prompt_embed.repeat(batch_size, 1, 1) + + for _ in range(iterations): + prompts, prompt_metadata = zip(*[self.prompt_fn() for _ in range(batch_size)]) + + prompt_ids = self.sd_pipeline.tokenizer( + prompts, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=self.sd_pipeline.tokenizer.model_max_length, + ).input_ids.to(self.accelerator.device) + prompt_embeds = self.sd_pipeline.text_encoder(prompt_ids)[0] + + with self.autocast(): + sd_output = self.sd_pipeline( + prompt_embeds=prompt_embeds, + negative_prompt_embeds=sample_neg_prompt_embeds, + num_inference_steps=self.config.sample_num_steps, + guidance_scale=self.config.sample_guidance_scale, + eta=self.config.sample_eta, + output_type="pt", + ) + + images = sd_output.images + latents = sd_output.latents + log_probs = sd_output.log_probs + + latents = torch.stack(latents, dim=1) # (batch_size, num_steps + 1, ...) + log_probs = torch.stack(log_probs, dim=1) # (batch_size, num_steps, 1) + timesteps = self.sd_pipeline.scheduler.timesteps.repeat(batch_size, 1) # (batch_size, num_steps) + + samples.append( + { + "prompt_ids": prompt_ids, + "prompt_embeds": prompt_embeds, + "timesteps": timesteps, + "latents": latents[:, :-1], # each entry is the latent before timestep t + "next_latents": latents[:, 1:], # each entry is the latent after timestep t + "log_probs": log_probs, + "negative_prompt_embeds": sample_neg_prompt_embeds, + } + ) + prompt_image_pairs.append([images, prompts, prompt_metadata]) + + return samples, prompt_image_pairs + + def _train_batched_samples(self, inner_epoch, epoch, global_step, batched_samples): + """ + Train on a batch of samples. Main training segment + + Args: + inner_epoch (int): The current inner epoch + epoch (int): The current epoch + global_step (int): The current global step + batched_samples (list[dict[str, torch.Tensor]]): The batched samples to train on + + Side Effects: + - Model weights are updated + - Logs the statistics to the accelerator trackers. + + Returns: + global_step (int): The updated global step + """ + info = defaultdict(list) + for _i, sample in enumerate(batched_samples): + if self.config.train_cfg: + # concat negative prompts to sample prompts to avoid two forward passes + embeds = torch.cat([sample["negative_prompt_embeds"], sample["prompt_embeds"]]) + else: + embeds = sample["prompt_embeds"] + + for j in range(self.num_train_timesteps): + with self.accelerator.accumulate(self.sd_pipeline.unet): + loss, approx_kl, clipfrac = self.calculate_loss( + sample["latents"][:, j], + sample["timesteps"][:, j], + sample["next_latents"][:, j], + sample["log_probs"][:, j], + sample["advantages"], + embeds, + ) + info["approx_kl"].append(approx_kl) + info["clipfrac"].append(clipfrac) + info["loss"].append(loss) + + self.accelerator.backward(loss) + if self.accelerator.sync_gradients: + self.accelerator.clip_grad_norm_( + self.trainable_layers.parameters() + if not isinstance(self.trainable_layers, list) + else self.trainable_layers, + self.config.train_max_grad_norm, + ) + self.optimizer.step() + self.optimizer.zero_grad() + + # Checks if the accelerator has performed an optimization step behind the scenes + if self.accelerator.sync_gradients: + # log training-related stuff + info = {k: torch.mean(torch.stack(v)) for k, v in info.items()} + info = self.accelerator.reduce(info, reduction="mean") + info.update({"epoch": epoch, "inner_epoch": inner_epoch}) + self.accelerator.log(info, step=global_step) + global_step += 1 + info = defaultdict(list) + return global_step + + def _config_check(self) -> tuple[bool, str]: + samples_per_epoch = ( + self.config.sample_batch_size * self.accelerator.num_processes * self.config.sample_num_batches_per_epoch + ) + total_train_batch_size = ( + self.config.train_batch_size + * self.accelerator.num_processes + * self.config.train_gradient_accumulation_steps + ) + + if not self.config.sample_batch_size >= self.config.train_batch_size: + return ( + False, + f"Sample batch size ({self.config.sample_batch_size}) must be greater than or equal to the train batch size ({self.config.train_batch_size})", + ) + if not self.config.sample_batch_size % self.config.train_batch_size == 0: + return ( + False, + f"Sample batch size ({self.config.sample_batch_size}) must be divisible by the train batch size ({self.config.train_batch_size})", + ) + if not samples_per_epoch % total_train_batch_size == 0: + return ( + False, + f"Number of samples per epoch ({samples_per_epoch}) must be divisible by the total train batch size ({total_train_batch_size})", + ) + return True, "" + + def train(self, epochs: Optional[int] = None): + """ + Train the model for a given number of epochs + """ + global_step = 0 + if epochs is None: + epochs = self.config.num_epochs + for epoch in range(self.first_epoch, epochs): + global_step = self.step(epoch, global_step) + + def _save_pretrained(self, save_directory): + self.sd_pipeline.save_pretrained(save_directory) + self.create_model_card() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{black2024training, + title = {{Training Diffusion Models with Reinforcement Learning}}, + author = {Kevin Black and Michael Janner and Yilun Du and Ilya Kostrikov and Sergey Levine}, + year = 2024, + booktitle = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=YCWjhGrJFD}, + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="DDPO", + trainer_citation=citation, + paper_title="Training Diffusion Models with Reinforcement Learning", + paper_id="2305.13301", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/dpo_config.py b/lib/python3.12/site-packages/trl/trainer/dpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..835909698564dd50ce5e646744e55e70dcc2341a --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/dpo_config.py @@ -0,0 +1,443 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Callable, Optional, Union + +from transformers import TrainingArguments + + +class FDivergenceType(Enum): + REVERSE_KL = "reverse_kl" + JS_DIVERGENCE = "js_divergence" + ALPHA_DIVERGENCE = "alpha_divergence" + + +class FDivergenceConstants: + ALPHA_DIVERGENCE_COEF_KEY = "alpha_divergence_coef" + ALPHA_DIVERGENCE_COEF_DEFAULT = 1.0 + + +@dataclass +class DPOConfig(TrainingArguments): + r""" + Configuration class for the [`DPOTrainer`]. + + This class includes only the parameters that are specific to DPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model and reference model + + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of the + [`DPOTrainer`] is provided as a string. + ref_model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `ref_model` argument of the + [`DPOTrainer`] is provided as a string. + model_adapter_name (`str` or `None`, *optional*, defaults to `None`): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str` or `None`, *optional*, defaults to `None`): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + force_use_ref_model (`bool`, *optional*, defaults to `False`): + If you provide a PEFT model as the active model and wish to use a different model for the `ref_model`, set + this flag to `True`. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + use_logits_to_keep (`bool`, *optional*, defaults to `False`): + If `True`, only a specified number of logits are computed in the forward pass. This can be useful for + saving memory and speeding up training by not computing the logits for all tokens, especially in scenarios + when working with very long prompts where labels are ignored (-100). + + > Parameters that control the data preprocessing + + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Padding value to use for labels. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the full sequence (prompt + completion). + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the sequence exceeds `max_length`. Possible values are `"keep_end"` and + `"keep_start"`. + padding_free (`bool`, *optional*, defaults to `False`): + Whether to perform forward passes without padding by flattening all sequences in the batch into a single + continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only + supported with the `flash_attention_2` attention implementation, which can efficiently handle the flattened + batch structure. + precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): + Whether to precompute the log probabilities from the reference model. Setting this to `True` allows + training without needing the reference model during training, which can help reduce GPU memory usage. If + set to `False` (default), the reference model will be used during training to compute log probabilities + on-the-fly. + precompute_ref_batch_size (`int` or `None`, *optional*, defaults to `None`): + Batch size to use when precomputing reference model log probabilities. This can be set higher than the + training batch size to speed up preprocessing. If `None`, defaults to `per_device_train_batch_size` for + training and `per_device_eval_batch_size` for evaluation. + tools (`Optional[list[Union[dict, Callable]]]`, *optional*, defaults to `None`): + List of tools (callable functions) that will be accessible to the model. If the template does not support + function calling, this argument will have no effect. + + > Parameters that control the training + + loss_type (`str`, *optional*, defaults to `"sigmoid"`): + Type of loss to use. Possible values are: + + - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. + - `"hinge"`: hinge loss on the normalized likelihood from the + [SLiC](https://huggingface.co/papers/2305.10425) paper. + - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. + - `"exo_pair"`: pairwise EXO loss from the [EXO](https://huggingface.co/papers/2402.00856) paper. + - `"nca_pair"`: pairwise NCA loss from the [NCA](https://huggingface.co/papers/2402.05369) paper. + - `"robust"`: unbiased estimate of the DPO loss that is robust to preference noise from the [Robust + DPO](https://huggingface.co/papers/2403.00409) paper. + - `"bco_pair"`: pairwise BCO loss from the [BCO](https://huggingface.co/papers/2404.04656) paper. + - `"sppo_hard"`: SPPO loss with hard label from the [SPPO](https://huggingface.co/papers/2405.00675) + paper. + - `"aot"`: AOT loss for paired datasets from the [AOT](https://huggingface.co/papers/2406.05882) paper. + - `"aot_pair"`: AOT loss for unpaired datasets from the [AOT](https://huggingface.co/papers/2406.05882) + paper. + - `"discopop"`: DiscoPOP (a.k.a Log-Ratio Modulated Loss, LRML) loss from the + [DiscoPOP](https://huggingface.co/papers/2406.08414) paper. + - `"apo_zero"`: APO-zero loss from the [APO](https://huggingface.co/papers/2408.06266) paper. + - `"apo_down"`: APO-down loss from the [APO](https://huggingface.co/papers/2408.06266) paper. + + use_liger_loss (`bool`, *optional*, defaults to `False`): + Whether to use Liger loss. + base_model_attribute_name (`str`, *optional*, defaults to `"model"`): + Name of the attribute in the model that contains the base model. This is used to get the base model from + the model when the model does not have a `get_decoder` method in the case when `use_liger_loss` is `True`. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in + the [paper](https://huggingface.co/papers/2310.12036). + f_divergence_type (`str`, *optional*, defaults to `FDivergenceType.REVERSE_KL`): + Type of f-divergence regularization function to compute divergence between policy and reference model. + f_alpha_divergence_coef (`float`, *optional*, defaults to `1.0`): + α coefficient in the α-divergence u^-α regularization function for DPO loss. + reference_free (`bool`, *optional*, defaults to `False`): + Whether to ignore the provided reference model and implicitly use a reference model that assigns equal + probability to all responses. + label_smoothing (`float`, *optional*, defaults to `0.0`): + Robust DPO label smoothing parameter from the [cDPO report](https://ericmitchell.ai/cdpo.pdf) and [Robust + DPO](https://huggingface.co/papers/2403.00409) paper that should be between `0.0` and `0.5`. + use_weighting (`bool`, *optional*, defaults to `False`): + Whether to weight the loss as done in the [WPO paper](https://huggingface.co/papers/2406.11827). + rpo_alpha (`float`, *optional*, defaults to `None`): + α parameter from the [RPO paper](https://huggingface.co/papers/2404.19733) (v3), which controls the + weighting of the NLL term in the loss. If `None`, no weighting is applied and the loss is the same as the + DPO loss. The paper recommends `rpo_alpha=1.0`. + ld_alpha (`float` or `None`, *optional*, defaults to `None`): + α parameter from the [LD-DPO paper](https://huggingface.co/papers/2409.06411), which controls the weighting + of the verbose token log-probabilities in responses. If `None`, no weighting is applied to the verbose + part, and the loss is equivalent to the standard DPO loss. The paper recommends setting `ld_alpha` between + `0.0` and `1.0`. + discopop_tau (`float`, *optional*, defaults to `0.05`): + τ/temperature parameter from the [DiscoPOP](https://huggingface.co/papers/2406.08414) paper, which controls + the shape of log ratio modulated loss. The paper recommends the default value `discopop_tau=0.05`. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originites from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix + between the current policy and the previous reference policy during updates. The reference policy is + updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you + must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how + frequently the current policy is synchronized with the reference policy. To use this parameter, you must + set `sync_ref_model=True`. + + > Parameters that control the logging + + generate_during_eval (`bool`, *optional*, defaults to `False`): + Whether to generate and log completions from both the model and the reference model to W&B or Comet during + evaluation. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs", "ref_model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model and reference model + model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of " + "the `DPOTrainer` is provided as a string." + }, + ) + ref_model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `ref_model` argument " + "of the `DPOTrainer` is provided as a string." + }, + ) + model_adapter_name: Optional[str] = field( + default=None, + metadata={"help": "Name of the train target PEFT adapter, when using LoRA with multiple adapters."}, + ) + ref_adapter_name: Optional[str] = field( + default=None, + metadata={"help": "Name of the reference PEFT adapter, when using LoRA with multiple adapters."}, + ) + force_use_ref_model: bool = field( + default=False, + metadata={ + "help": "If you provide a PEFT model as the active model and wish to use a different model for the " + "`ref_model`, set this flag to `True`." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model and reference model."}, + ) + use_logits_to_keep: bool = field( + default=False, + metadata={ + "help": "If `True`, only a specified number of logits are computed in the forward pass. This can be " + "useful for saving memory and speeding up training by not computing the logits for all tokens, especially " + "in scenarios when working with very long prompts where labels are ignored (-100)." + }, + ) + + # Parameters that control the data preprocessing + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + padding_value: Optional[int] = field( + default=None, + metadata={"help": "Padding value to use. If `None`, the padding value of the tokenizer is used."}, + ) + label_pad_token_id: int = field( + default=-100, + metadata={"help": "Padding value to use for labels."}, + ) + max_prompt_length: Optional[int] = field( + default=512, + metadata={"help": "Maximum length of the prompt."}, + ) + max_completion_length: Optional[int] = field( + default=None, + metadata={"help": "Maximum length of the completion."}, + ) + max_length: Optional[int] = field( + default=1024, + metadata={"help": "Maximum length of the full sequence (prompt + completion)."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the sequence exceeds `max_length`. Possible values are `'keep_end'` " + "and `'keep_start'`.", + "choices": ["keep_end", "keep_start"], + }, + ) + padding_free: bool = field( + default=False, + metadata={ + "help": "Whether to perform forward passes without padding by flattening all sequences in the batch into " + "a single continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, " + "this is only supported with the `flash_attention_2` attention implementation, which can efficiently " + "handle the flattened batch structure." + }, + ) + precompute_ref_log_probs: bool = field( + default=False, + metadata={ + "help": "Whether to precompute the log probabilities from the reference model. Setting this to `True` " + "allows training without needing the reference model during training, which can help reduce GPU memory " + "usage. If set to `False` (default), the reference model will be used during training to compute log " + "probabilities on-the-fly." + }, + ) + precompute_ref_batch_size: Optional[int] = field( + default=None, + metadata={ + "help": "Batch size to use when precomputing reference model log probabilities. This can be set higher " + "than the training batch size to speed up preprocessing. If `None`, defaults to " + "`per_device_train_batch_size` for training and `per_device_eval_batch_size` for evaluation." + }, + ) + tools: Optional[list[Union[dict, Callable]]] = field( + default=None, + metadata={ + "help": "List of tools (callable functions) that will be accessible to the model. If the template does " + "not support function calling, this argument will have no effect." + }, + ) + + # Parameters that control the training + loss_type: str = field( + default="sigmoid", + metadata={ + "help": "Type of loss to use.", + "choices": [ + "sigmoid", + "hinge", + "ipo", + "exo_pair", + "nca_pair", + "robust", + "bco_pair", + "sppo_hard", + "aot", + "aot_pair", + "discopop", + "apo_zero", + "apo_down", + ], + }, + ) + use_liger_loss: bool = field( + default=False, + metadata={"help": "Whether to use Liger loss."}, + ) + base_model_attribute_name: str = field( + default="model", + metadata={ + "help": "Name of the attribute in the model that contains the base model. This is used to get the base " + "model from the model when the model does not have a `get_decoder` method in the case when " + "`use_liger_loss` is `True`." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the deviation from the reference model. " + "Higher β means less deviation from the reference model." + }, + ) + f_divergence_type: FDivergenceType = field( + default=FDivergenceType.REVERSE_KL, + metadata={ + "help": "Type of f-divergence regularization function to compute divergence between policy and reference " + "model." + }, + ) + f_alpha_divergence_coef: float = field( + default=1.0, + metadata={"help": "α coefficient in the α-divergence u^-α regularization function for DPO loss."}, + ) + reference_free: bool = field( + default=False, + metadata={ + "help": "Whether to ignore the provided reference model and implicitly use a reference model that assigns " + "equal probability to all responses." + }, + ) + label_smoothing: float = field( + default=0.0, + metadata={ + "help": "Robust DPO label smoothing parameter from the cDPO report and Robust DPO paper that should " + "be between `0.0` and `0.5`." + }, + ) + use_weighting: bool = field( + default=False, + metadata={"help": "Whether to weight the loss as done in the WPO paper."}, + ) + rpo_alpha: Optional[float] = field( + default=None, + metadata={ + "help": "α parameter from the RPO paper (v3), which controls the weighting of the NLL term in the loss. " + "If `None`, no weighting is applied and the loss is the same as the DPO loss. The paper recommends " + "`rpo_alpha=1.0`." + }, + ) + ld_alpha: Optional[float] = field( + default=None, + metadata={ + "help": "α parameter from the LD-DPO paper, which controls the weighting of the verbose token " + "log-probabilities in responses. If `None`, no weighting is applied to the verbose part, and the loss is " + "equivalent to the standard DPO loss. The paper recommends setting `ld_alpha` between `0.0` and `1.0`.", + }, + ) + discopop_tau: float = field( + default=0.05, + metadata={ + "help": "τ/temperature parameter from the DiscoPOP paper, which controls the shape of log ratio modulated " + "loss. The paper recommends the default value `discopop_tau=0.05`." + }, + ) + sync_ref_model: bool = field( + default=False, + metadata={ + "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` " + "steps, using the `ref_model_mixup_alpha` parameter." + }, + ) + ref_model_mixup_alpha: float = field( + default=0.6, + metadata={ + "help": "α parameter from the TR-DPO paper, which controls the mix between the current policy and the " + "previous reference policy during updates. The reference policy is updated according to the equation: " + "`π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + ref_model_sync_steps: int = field( + default=512, + metadata={ + "help": "τ parameter from the TR-DPO paper, which determines how frequently the current policy is " + "synchronized with the reference policy. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + + # Parameters that control the logging + generate_during_eval: bool = field( + default=False, + metadata={ + "help": "Whether to generate and log completions from both the model and the reference model to W&B or " + "Comet during evaluation." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/lib/python3.12/site-packages/trl/trainer/dpo_trainer.py b/lib/python3.12/site-packages/trl/trainer/dpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..f0ac07acaea1698d324bba46e13c2c7c106cad4a --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/dpo_trainer.py @@ -0,0 +1,1930 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import random +import textwrap +import warnings +from collections import defaultdict +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Literal, Optional, Union + +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState +from accelerate.utils import tqdm +from datasets import Dataset, IterableDataset +from torch import autocast +from torch.utils.data import DataLoader +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + is_comet_available, + is_wandb_available, +) +from transformers.data.data_collator import DataCollatorMixin +from transformers.models.auto.modeling_auto import MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalLoopOutput +from transformers.utils import is_liger_kernel_available, is_peft_available + +from ..data_utils import maybe_apply_chat_template, maybe_extract_prompt +from ..models import create_reference_model, prepare_deepspeed +from ..models.utils import prepare_fsdp +from .callbacks import SyncRefModelCallback +from .dpo_config import DPOConfig, FDivergenceConstants, FDivergenceType +from .utils import ( + RunningMoments, + cap_exp, + disable_dropout_in_model, + empty_cache, + flush_left, + flush_right, + generate_model_card, + get_comet_experiment_url, + log_table_to_comet_experiment, + pad, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import PeftConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearDPOLoss + + +if is_wandb_available(): + import wandb + + +def shift_tokens_right(input_ids: torch.Tensor, decoder_start_token_id: int) -> torch.Tensor: + """Shift input ids one token to the right, and pad with pad_token_id""" + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() + shifted_input_ids[:, 0] = decoder_start_token_id + + +@dataclass +class DataCollatorForPreference(DataCollatorMixin): + """ + Data collator used for preference data. Inputs are dynamically padded to the maximum length of a batch if they are + not all of the same length. + + Args: + pad_token_id (`int`): + Token ID to use for padding. + return_tensors (`str`, *optional*, defaults to `"pt"`): + Type of Tensor to return. Only `"pt"` is currently supported. + + Examples: + ```python + >>> from trl import DataCollatorForPreference + + >>> collator = DataCollatorForPreference(pad_token_id=0) + >>> examples = [ + ... {"prompt_input_ids": [1, 2, 3], "chosen_input_ids": [4, 5], "rejected_input_ids": [6]}, + ... {"prompt_input_ids": [7, 8], "chosen_input_ids": [9, 10], "rejected_input_ids": [11, 12, 13]}, + ... ] + >>> collator(examples) + {'prompt_input_ids': tensor([[1, 2, 3], + [0, 7, 8]]), + 'prompt_attention_mask': tensor([[1, 1, 1], + [0, 1, 1]]), + 'chosen_input_ids': tensor([[ 4, 5], + [ 9, 10]]), + 'chosen_attention_mask': tensor([[1, 1], + [1, 1]]), + 'rejected_input_ids': tensor([[ 6, 0, 0], + [11, 12, 13]]), + 'rejected_attention_mask': tensor([[1, 0, 0], + [1, 1, 1]]) + } + ``` + """ + + pad_token_id: int + return_tensors: str = "pt" + + def torch_call(self, examples: list[Union[list[int], Any, dict[str, Any]]]) -> dict[str, Any]: + # Convert to tensor + prompt_input_ids = [torch.tensor(example["prompt_input_ids"]) for example in examples] + prompt_attention_mask = [torch.ones_like(input_ids) for input_ids in prompt_input_ids] + chosen_input_ids = [torch.tensor(example["chosen_input_ids"]) for example in examples] + chosen_attention_mask = [torch.ones_like(input_ids) for input_ids in chosen_input_ids] + rejected_input_ids = [torch.tensor(example["rejected_input_ids"]) for example in examples] + rejected_attention_mask = [torch.ones_like(input_ids) for input_ids in rejected_input_ids] + if "pixel_values" in examples[0]: + pixel_values = [torch.tensor(example["pixel_values"]) for example in examples] + if "pixel_attention_mask" in examples[0]: + pixel_attention_mask = [torch.tensor(example["pixel_attention_mask"]) for example in examples] + if "ref_chosen_logps" in examples[0] and "ref_rejected_logps" in examples[0]: + ref_chosen_logps = torch.tensor([example["ref_chosen_logps"] for example in examples]) + ref_rejected_logps = torch.tensor([example["ref_rejected_logps"] for example in examples]) + + # Pad + output = {} + output["prompt_input_ids"] = pad(prompt_input_ids, padding_value=self.pad_token_id, padding_side="left") + output["prompt_attention_mask"] = pad(prompt_attention_mask, padding_value=0, padding_side="left") + output["chosen_input_ids"] = pad(chosen_input_ids, padding_value=self.pad_token_id) + output["chosen_attention_mask"] = pad(chosen_attention_mask, padding_value=0) + output["rejected_input_ids"] = pad(rejected_input_ids, padding_value=self.pad_token_id) + output["rejected_attention_mask"] = pad(rejected_attention_mask, padding_value=0) + if "pixel_values" in examples[0]: + output["pixel_values"] = pad(pixel_values, padding_value=0.0) + if "pixel_attention_mask" in examples[0]: + output["pixel_attention_mask"] = pad(pixel_attention_mask, padding_value=0) + if "image_sizes" in examples[0]: + output["image_sizes"] = torch.tensor([example["image_sizes"] for example in examples]) + if "ref_chosen_logps" in examples[0] and "ref_rejected_logps" in examples[0]: + output["ref_chosen_logps"] = ref_chosen_logps + output["ref_rejected_logps"] = ref_rejected_logps + + return output + + +class DPOTrainer(Trainer): + """ + Trainer for Direct Preference Optimization (DPO) method. + + This class is a wrapper around the [`transformers.Trainer`] class and inherits all of its attributes and methods. + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + ref_model (`PreTrainedModelWrapper`): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + args ([`DPOConfig`], *optional*, defaults to `None`): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator (`DataCollator`, *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`DataCollatorForPreference`]. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. DPO supports [preference](#preference) type and. The format of the samples can + be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoTokenizer.from_pretrained`]. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function that will be used to compute metrics at evaluation. Must take a [`EvalPrediction`] and return + a dictionary string to metric values. *Note* When passing TrainingArgs with `batch_eval_metrics` set to + `True`, your compute_metrics function must take a boolean `compute_result` argument. This will be triggered + after the last eval batch to signal that the function needs to calculate and return the global summary + statistics rather than accumulating the batch-level statistics. + callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*, defaults to `None`): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*, defaults to `None`): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + """ + + _tag_names = ["trl", "dpo"] + + def __init__( + self, + model: Union[str, nn.Module, PreTrainedModel], + ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: Optional[DPOConfig] = None, + data_collator: Optional[DataCollator] = None, # type: ignore + train_dataset: Optional[Union[Dataset, IterableDataset]] = None, + eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + optimizer_cls_and_kwargs: Optional[tuple[type[torch.optim.Optimizer], dict[str, Any]]] = None, + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional["PeftConfig"] = None, + ): + # Args + model_id = model if isinstance(model, str) else model.config._name_or_path + if args is None: + model_name = model_id.split("/")[-1] + args = DPOConfig(f"{model_name}-DPO") + + # Handle the tokenizer + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model_id) + + if args.padding_value is not None: + self.padding_value = args.padding_value + else: + if hasattr(processing_class, "pad_token_id") and processing_class.pad_token_id is not None: + self.padding_value = processing_class.pad_token_id + elif hasattr(processing_class, "tokenizer") and processing_class.tokenizer.pad_token_id is not None: + self.padding_value = processing_class.tokenizer.pad_token_id + else: + raise ValueError( + "`padding_value` is not specified in `DPOConfig`, and `pad_token_id` is missing in the " + "`processing_class`. Please either set the `padding_value` argument in `DPOConfig`, or set " + "`tokenizer.pad_token` (e.g., `tokenizer.pad_token = tokenizer.eos_token`) before instantiating " + "the trainer." + ) + + # Model + if not isinstance(model, str) and ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, you must mass a copy of it, or `None` if you use peft." + ) + + if args.model_init_kwargs is not None and not isinstance(model, str): + warnings.warn( + "You passed model_init_kwargs to the `DPOConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + if isinstance(model, str): + model = self._create_model_from_path(model, args) + + if args.ref_model_init_kwargs is not None and not isinstance(ref_model, str): + warnings.warn( + "You passed ref_model_init_kwargs to the `DPOConfig`, but your ref_model is already instantiated. " + "The `ref_model_init_kwargs` will be ignored." + ) + if isinstance(ref_model, str): + ref_model = self._create_model_from_path(ref_model, args, is_ref=True) + + # PEFT configuration and model wrapping + model = self._prepare_peft_model(model, ref_model, peft_config, args) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." + " Please install `wandb` or `comet-ml` to resolve." + ) + + self.is_encoder_decoder = model.config.is_encoder_decoder + self.is_vision_model = model.config.model_type in MODEL_FOR_VISION_2_SEQ_MAPPING_NAMES.keys() + self.is_peft_model = is_peft_available() and isinstance(model, PeftModel) + self.model_adapter_name = args.model_adapter_name + self.ref_adapter_name = args.ref_adapter_name + self.reference_free = args.reference_free + + if ref_model: + self.ref_model = ref_model + elif self.is_peft_model or args.precompute_ref_log_probs: + # The `model` with adapters turned off will be used as the reference model + self.ref_model = None + else: + self.ref_model = create_reference_model(model) + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Liger kernel + if args.use_liger_loss: + if not is_liger_kernel_available(): + raise ImportError( + "You set `use_liger_loss=True` but the liger kernel is not available. " + "Please install liger-kernel first: `pip install liger-kernel`" + ) + if args.loss_type != "sigmoid": + raise ValueError( + "You set `use_liger_loss=True` but the loss type is not `sigmoid`. " + "Please set `loss_type='sigmoid'` to use the liger kernel." + ) + self.dpo_loss_fn = LigerFusedLinearDPOLoss( + ignore_index=args.label_pad_token_id, + beta=args.beta, + use_ref_model=not args.reference_free, + average_log_prob=False, + ) + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in DPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and + # "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens + # of the input, floating-point operations will not be computed." To suppress this warning, we set the + # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate + # that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + # Data collator + if data_collator is None: + data_collator = DataCollatorForPreference(pad_token_id=self.padding_value) + + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.max_prompt_length = args.max_prompt_length + self.max_completion_length = args.max_completion_length + self.max_length = args.max_length + self.truncation_mode = args.truncation_mode + self.precompute_ref_log_probs = args.precompute_ref_log_probs + self.use_logits_to_keep = args.use_logits_to_keep + + if args.padding_free: + if model.config._attn_implementation != "flash_attention_2": + warnings.warn( + "Padding-free training is enabled, but the attention implementation is not set to " + "'flash_attention_2'. Padding-free training flattens batches into a single sequence, and " + "'flash_attention_2' is the only known attention mechanism that reliably supports this. Using " + "other implementations may lead to unexpected behavior. To ensure compatibility, set " + "`attn_implementation='flash_attention_2'` in the model configuration, or verify that your " + "attention mechanism can handle flattened sequences." + ) + if args.per_device_train_batch_size == 1: + warnings.warn( + "You are using a per_device_train_batch_size of 1 with padding-free training. Using a batch size " + "of 1 anihilate the benefits of padding-free training. Please consider increasing the batch size " + "to at least 2." + ) + self.padding_free = args.padding_free + + # Since ref_logs are precomputed on the first call to get_train/eval_dataloader + # keep track of first called to avoid computation of future calls + self._precomputed_train_ref_log_probs = False + self._precomputed_eval_ref_log_probs = False + + if ( + args.loss_type in ["hinge", "ipo", "bco_pair", "sppo_hard", "nca_pair", "apo_zero", "apo_down"] + and args.label_smoothing > 0 + ): + warnings.warn( + f"You are using the {args.loss_type} loss type that does not support label smoothing. The " + "`label_smoothing` parameter will be ignored. Set `label_smoothing` to `0.0` to remove this warning.", + UserWarning, + ) + if args.loss_type == "kto_pair": + raise ValueError("Support for kto_pair has been removed in DPOTrainer. Please use KTOTrainer.") + + self.beta = args.beta + self.label_smoothing = args.label_smoothing + self.loss_type = args.loss_type + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.use_weighting = args.use_weighting + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + warnings.warn( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + UserWarning, + ) + + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + self.f_divergence_type = args.f_divergence_type + self.f_divergence_params = {FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY: args.f_alpha_divergence_coef} + self.dataset_num_proc = args.dataset_num_proc + + # Dataset preparation + train_dataset = self._prepare_dataset(train_dataset, processing_class, args, "train") + if eval_dataset is not None: + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset(eval_dataset, processing_class, args, "eval") + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + # Deepspeed Zero-3 does not support precompute_ref_log_probs + if self.is_deepspeed_enabled: + if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`." + ) + + if self.ref_model is None: + if not (self.is_peft_model or self.precompute_ref_log_probs): + raise ValueError( + "No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`" + ) + if args.sync_ref_model: + raise ValueError( + "You currently cannot use `ref_model=None` with TR-DPO method. Please provide `ref_model`." + ) + else: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + if self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with TR-DPO method. Please set `precompute_ref_log_probs=False`." + ) + + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + if self.loss_type == "bco_pair": + self.running = RunningMoments(self.accelerator) + + def _create_model_from_path(self, model_path: str, args: DPOConfig, is_ref: bool = False) -> PreTrainedModel: + """Creates a model from a path or model identifier.""" + if not is_ref: + model_init_kwargs = args.model_init_kwargs or {} + else: + model_init_kwargs = args.ref_model_init_kwargs or {} + + # Handle torch dtype + torch_dtype = model_init_kwargs.get("torch_dtype") + if isinstance(torch_dtype, torch.dtype) or torch_dtype == "auto" or torch_dtype is None: + pass # torch_dtype is already a torch.dtype or "auto" or None + elif isinstance(torch_dtype, str): # it's a str, but not "auto" + torch_dtype = getattr(torch, torch_dtype) + model_init_kwargs["torch_dtype"] = torch_dtype + else: + raise ValueError( + "Invalid `torch_dtype` passed to `DPOConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}." + ) + # Disable caching if gradient checkpointing is enabled (not supported) + # if args.gradient_checkpointing: + # model_init_kwargs["use_cache"] = False + + # Create model + model = AutoModelForCausalLM.from_pretrained(model_path, **model_init_kwargs) + return model + + def _prepare_peft_model( + self, model: PreTrainedModel, ref_model: PreTrainedModel, peft_config: Any, args: DPOConfig + ) -> PreTrainedModel: + """Prepares a model for PEFT training.""" + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if ref_model is not None and not args.force_use_ref_model: + raise ValueError( + "You passed both a ref_model and a peft_config. For training PEFT adapters with DPO there is no need to pass a reference" + " model. Please pass `ref_model=None` in case you want to train PEFT adapters, or pass a ref_model with `force_use_ref_model=True` in DPOTrainer's init." + " if you want to use a different ref_model." + ) + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + + else: + model = self._prepare_gradient_checkpointing(model, args) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + else: + model = self._prepare_gradient_checkpointing(model, args) + + return model + + def _prepare_gradient_checkpointing(self, model: PreTrainedModel, args: DPOConfig): + """Prepare the gradienting checkpointing for the model.""" + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + if args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + return model + + def _prepare_dataset( + self, + dataset: Union[Dataset, IterableDataset], + processing_class: Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin], + args: DPOConfig, + dataset_name: str, + ) -> Union[Dataset, IterableDataset]: + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc nor writer_batch_size + map_kwargs["num_proc"] = args.dataset_num_proc + map_kwargs["writer_batch_size"] = 10 + + with PartialState().main_process_first(): + # Extract prompt if needed + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Extracting prompt in {dataset_name} dataset" + dataset = dataset.map(maybe_extract_prompt, **map_kwargs) + + # Apply the chat template if needed + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Applying chat template to {dataset_name} dataset" + dataset = dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class, "tools": args.tools}, **map_kwargs + ) + + # Tokenize the dataset + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" + + dataset = dataset.map( + self.tokenize_row if not self.is_vision_model else self.process_row, + remove_columns=["chosen", "rejected"], + fn_kwargs={ + "processing_class": processing_class, + "max_prompt_length": args.max_prompt_length, + "max_completion_length": args.max_completion_length, + # for enc-dec, we add the special tokens ([bos_token] + prompt + [eos_token]; completion + [eos_token]) + "add_special_tokens": False, + }, + **map_kwargs, + ) + + return dataset + + @staticmethod + def tokenize_row(features, processing_class, max_prompt_length, max_completion_length, add_special_tokens): + """ + Tokenize a row of the dataset. + + Args: + features (`dict[str, str]`): + Row of the dataset, should contain the keys `"prompt"`, `"chosen"`, and `"rejected"`. + processing_class (`PreTrainedTokenizerBase`): + Processing class used to process the data. + max_prompt_length (`int` or `None`): + Maximum length of the prompt sequence. If `None`, the prompt sequence is not truncated. + max_completion_length (`int` or `None`): + Maximum length of the completion sequences. If `None`, the completion sequences are not truncated. + add_special_tokens (`bool`): + Whether to add special tokens to the sequences. Typically used for encoder-decoder models. If `True`, + the prompt sequence will have a bos token prepended and an eos token appended. In any case, the + completion sequences will have an eos token appended. + + Returns: + `dict[str, list[int]]`: + Tokenized sequences with the keys `"prompt_input_ids"`, `"chosen_input_ids"`, and + `"rejected_input_ids". + + Example: + ```python + >>> from transformers import GPT2Tokenizer + + >>> tokenizer = GPT2Tokenizer.from_pretrained("gpt2") + >>> features = {"prompt": "The sky is", "chosen": " blue", "rejected": " green"} + >>> DPOTrainer.tokenize_row( + ... features, tokenizer, max_prompt_length=3, max_completion_length=3, add_special_tokens=False + ... ) + {'prompt_input_ids': [464, 6766, 318], 'chosen_input_ids': [4171, 50256], 'rejected_input_ids': [4077, 50256]} + ``` + """ + tokenizer = processing_class # the processing class is a tokenizer + prompt_input_ids = tokenizer(features["prompt"], add_special_tokens=False)["input_ids"] + chosen_input_ids = tokenizer(features["chosen"], add_special_tokens=False)["input_ids"] + rejected_input_ids = tokenizer(features["rejected"], add_special_tokens=False)["input_ids"] + + # Add special tokens (typically for encoder-decoder models) + if add_special_tokens: + if tokenizer.bos_token_id is not None: + prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids + if tokenizer.eos_token_id is not None: + prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id] + chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id] + rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id] + + # Truncate prompt and completion sequences + if max_prompt_length is not None: + prompt_input_ids = prompt_input_ids[-max_prompt_length:] + if max_completion_length is not None: + chosen_input_ids = chosen_input_ids[:max_completion_length] + rejected_input_ids = rejected_input_ids[:max_completion_length] + + return { + "prompt_input_ids": prompt_input_ids, + "chosen_input_ids": chosen_input_ids, + "rejected_input_ids": rejected_input_ids, + } + + @staticmethod + def process_row(features, processing_class, max_prompt_length, max_completion_length, add_special_tokens): + """ + Same as `tokenize_row` but for vision models. Please refer to `tokenize_row` for more information. + """ + processor, tokenizer = processing_class, processing_class.tokenizer # the processing class is a processor + processed_features = processor(images=features["images"], text=features["prompt"], add_special_tokens=False) + + prompt_input_ids = processed_features["input_ids"][0] + pixel_values = processed_features["pixel_values"][0] + chosen_input_ids = tokenizer(features["chosen"], add_special_tokens=False)["input_ids"] + rejected_input_ids = tokenizer(features["rejected"], add_special_tokens=False)["input_ids"] + + # Add special tokens (typically for encoder-decoder models) + if add_special_tokens: + if tokenizer.bos_token_id is not None: + prompt_input_ids = [tokenizer.bos_token_id] + prompt_input_ids + if tokenizer.eos_token_id is not None: + prompt_input_ids = prompt_input_ids + [tokenizer.eos_token_id] + chosen_input_ids = chosen_input_ids + [tokenizer.eos_token_id] + rejected_input_ids = rejected_input_ids + [tokenizer.eos_token_id] + + # Truncate prompt and completion sequences + if max_prompt_length is not None: + prompt_input_ids = prompt_input_ids[-max_prompt_length:] + if max_completion_length is not None: + chosen_input_ids = chosen_input_ids[:max_completion_length] + rejected_input_ids = rejected_input_ids[:max_completion_length] + + output = { + "prompt_input_ids": prompt_input_ids, + "pixel_values": pixel_values, + "chosen_input_ids": chosen_input_ids, + "rejected_input_ids": rejected_input_ids, + } + + if "pixel_attention_mask" in processed_features: + output["pixel_attention_mask"] = processed_features["pixel_attention_mask"][0] + if "image_sizes" in processed_features: + output["image_sizes"] = processed_features["image_sizes"][0] + + return output + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs. + # In DPOTrainer, we preprocess data, so using the model's signature columns doesn't work. + # Instead, we set them to the columns expected by `DataCollatorForPreference`, hence the override. + if self._signature_columns is None: + self._signature_columns = [ + "prompt_input_ids", + "chosen_input_ids", + "rejected_input_ids", + "image_sizes", + "ref_chosen_logps", + "ref_rejected_logps", + ] + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`. + """ + + if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs: + batch_size = self.args.precompute_ref_batch_size or self.args.per_device_train_batch_size + dataloader_params = { + "batch_size": batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params)) + + ref_chosen_logps = [] + ref_rejected_logps = [] + for padded_batch in tqdm(iterable=data_loader, desc="Train dataset reference log probs"): + ref_chosen_logp, ref_rejected_logp = self.compute_ref_log_probs(padded_batch) + ref_chosen_logp, ref_rejected_logp = self.accelerator.gather_for_metrics( + (ref_chosen_logp, ref_rejected_logp) + ) + ref_chosen_logps.append(ref_chosen_logp.cpu()) + ref_rejected_logps.append(ref_rejected_logp.cpu()) + + # Unnecessary cache clearing to avoid OOM + empty_cache() + self.accelerator.free_memory() + + all_ref_chosen_logps = torch.cat(ref_chosen_logps).float().numpy() + all_ref_rejected_logps = torch.cat(ref_rejected_logps).float().numpy() + + self.train_dataset = self.train_dataset.add_column(name="ref_chosen_logps", column=all_ref_chosen_logps) + self.train_dataset = self.train_dataset.add_column( + name="ref_rejected_logps", column=all_ref_rejected_logps + ) + + self._precomputed_train_ref_log_probs = True + + return super().get_train_dataloader() + + def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader: + """ + Returns the evaluation [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`. + + Args: + eval_dataset (`torch.utils.data.Dataset`, *optional*): + If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted + by the `model.forward()` method are automatically removed. It must implement `__len__`. + """ + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset + + if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs: + batch_size = self.args.precompute_ref_batch_size or self.args.per_device_eval_batch_size + dataloader_params = { + "batch_size": batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params)) + + ref_chosen_logps = [] + ref_rejected_logps = [] + for padded_batch in tqdm(iterable=data_loader, desc="Eval dataset reference log probs"): + ref_chosen_logp, ref_rejected_logp = self.compute_ref_log_probs(padded_batch) + ref_chosen_logp, ref_rejected_logp = self.accelerator.gather_for_metrics( + (ref_chosen_logp, ref_rejected_logp) + ) + ref_chosen_logps.append(ref_chosen_logp.cpu()) + ref_rejected_logps.append(ref_rejected_logp.cpu()) + + all_ref_chosen_logps = torch.cat(ref_chosen_logps).float().numpy() + all_ref_rejected_logps = torch.cat(ref_rejected_logps).float().numpy() + + eval_dataset = eval_dataset.add_column(name="ref_chosen_logps", column=all_ref_chosen_logps) + eval_dataset = eval_dataset.add_column(name="ref_rejected_logps", column=all_ref_rejected_logps) + + # Save calculated ref_chosen_logps and ref_rejected_logps to the eval_dataset for subsequent runs + if self.eval_dataset is not None: + self.eval_dataset = eval_dataset + self._precomputed_eval_ref_log_probs = True + + return super().get_eval_dataloader(eval_dataset=eval_dataset) + + @contextmanager + def null_ref_context(self): + """Context manager for handling null reference model (that is, peft adapter manipulation).""" + with ( + self.accelerator.unwrap_model(self.model).disable_adapter() + if self.is_peft_model and not self.ref_adapter_name + else nullcontext() + ): + if self.ref_adapter_name: + self.model.set_adapter(self.ref_adapter_name) + yield + if self.ref_adapter_name: + self.model.set_adapter(self.model_adapter_name or "default") + + def compute_ref_log_probs(self, batch: dict[str, torch.LongTensor]) -> dict: + """Computes log probabilities of the reference model for a single padded batch of a DPO specific dataset.""" + compte_ref_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + with torch.no_grad(), compte_ref_context_manager: + if self.ref_model is None: + with self.null_ref_context(): + ref_model_output = self.concatenated_forward(self.model, batch, is_ref_model=True) + else: + ref_model_output = self.concatenated_forward(self.ref_model, batch, is_ref_model=True) + return ref_model_output["chosen_logps"], ref_model_output["rejected_logps"] + + @staticmethod + def concatenated_inputs( + batch: dict[str, Union[list, torch.LongTensor]], padding_value: int + ) -> dict[str, torch.LongTensor]: + """ + Concatenate the `chosen` and `rejected` inputs from the batch into a single tensor for both the prompt and + completion sequences. + + Args: + batch (`dict[str, Union[list, torch.LongTensor]]`): + A batch of input data. The batch must contain the following keys: + + - `"prompt_input_ids"`: Tensor of shape `(batch_size, prompt_length)` representing the prompt input + IDs. + - `"chosen_input_ids"`: Tensor of shape `(batch_size, chosen_length)` representing the chosen + completion input IDs. + - `"rejected_input_ids"`: Tensor of shape `(batch_size, rejected_length)` representing the rejected + completion input IDs. + - `"prompt_pixel_values"` (optional): Tensor for pixel values, if available. + - `"prompt_pixel_attention_mask"` (optional): Tensor for pixel attention masks, if available. + + padding_value (`int`): + The padding value to use for the concatenated completion sequences (`chosen_input_ids` and + `rejected_input_ids`). + + Returns: + `dict[str, torch.LongTensor]`: A dictionary containing: + + - `"prompt_input_ids"`: Concatenated prompt input IDs of shape `(2 * batch_size, prompt_length)`. + - `"completion_input_ids"`: Concatenated chosen and rejected completion input IDs of shape `(2 * + batch_size, max_completion_length)`. + - `"prompt_attention_mask"`: Concatenated prompt attention masks of shape `(2 * batch_size, + prompt_length)`. + - `"completion_attention_mask"`: Concatenated chosen and rejected attention masks of shape `(2 * + batch_size, max_completion_length)`. + - `"pixel_values"` (optional): Concatenated pixel values if `"prompt_pixel_values"` are present. + - `"pixel_attention_mask"` (optional): Concatenated pixel attention masks if + `"prompt_pixel_attention_mask"` are present. + + Notes: + The completion input IDs and attention masks are padded to the maximum completion length of the chosen or + rejected sequences. + """ + output = {} + + # For the prompt, the input_ids are the same for both the chosen and rejected responses + output["prompt_input_ids"] = torch.cat([batch["prompt_input_ids"], batch["prompt_input_ids"]], dim=0) + output["prompt_attention_mask"] = torch.cat( + [batch["prompt_attention_mask"], batch["prompt_attention_mask"]], dim=0 + ) + if "pixel_values" in batch: + output["pixel_values"] = torch.cat([batch["pixel_values"], batch["pixel_values"]], dim=0) + + if "pixel_attention_mask" in batch: + output["pixel_attention_mask"] = torch.cat( + [batch["pixel_attention_mask"], batch["pixel_attention_mask"]], dim=0 + ) + if "image_sizes" in batch: + output["image_sizes"] = torch.cat([batch["image_sizes"], batch["image_sizes"]], dim=0) + + # Concatenate the chosen and rejected completions + max_completion_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) + output["completion_input_ids"] = torch.cat( + ( + pad_to_length(batch["chosen_input_ids"], max_completion_length, pad_value=padding_value), + pad_to_length(batch["rejected_input_ids"], max_completion_length, pad_value=padding_value), + ), + ) + output["completion_attention_mask"] = torch.cat( + ( + pad_to_length(batch["chosen_attention_mask"], max_completion_length, pad_value=0), + pad_to_length(batch["rejected_attention_mask"], max_completion_length, pad_value=0), + ), + ) + + return output + + def dpo_loss( + self, + chosen_logps: torch.FloatTensor, + rejected_logps: torch.FloatTensor, + ref_chosen_logps: torch.FloatTensor, + ref_rejected_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """ + Compute the DPO loss for a batch of policy and reference model log probabilities. + + Args: + chosen_logps (`torch.FloatTensor`): + Log probabilities of the model for the chosen responses. Shape: `(batch_size,)`. + rejected_logps (`torch.FloatTensor`): + Log probabilities of the model for the rejected responses. Shape: `(batch_size,)`. + ref_chosen_logps (`torch.FloatTensor`): + Log probabilities of the reference model for the chosen responses. Shape: `(batch_size,)`. + ref_rejected_logps (`torch.FloatTensor`): + Log probabilities of the reference model for the rejected responses. Shape: `(batch_size,)`. + + Returns: + A tuple of three tensors: `(losses, chosen_rewards, rejected_rewards)`. The losses tensor contains the DPO + loss for each example in the batch. The `chosen_rewards` and `rejected_rewards` tensors contain the rewards + for the chosen and rejected responses, respectively. + """ + device = self.accelerator.device + + # Get the log ratios for the chosen and rejected responses + chosen_logratios = chosen_logps.to(device) - (not self.reference_free) * ref_chosen_logps.to(device) + rejected_logratios = rejected_logps.to(device) - (not self.reference_free) * ref_rejected_logps.to(device) + + if self.f_divergence_type == FDivergenceType.ALPHA_DIVERGENCE.value: + # The alpha-divergence formula: (1 - u^-alpha) / alpha + # The divergence difference between the chosen and rejected sample is: + # (1 - u[w]^-alpha) / alpha - (1 - u[l]^-alpha) / alpha + # = (u[l]^-alpha - u[w]^-alpha) / alpha + # where u[w] and u[l] are the policy/reference probability ratios + # for the chosen and rejected samples, respectively. + alpha_coef = FDivergenceConstants.ALPHA_DIVERGENCE_COEF_DEFAULT + if self.f_divergence_params and FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY in self.f_divergence_params: + alpha_coef = float(self.f_divergence_params[FDivergenceConstants.ALPHA_DIVERGENCE_COEF_KEY]) + logits = (cap_exp(rejected_logratios * -alpha_coef) - cap_exp(chosen_logratios * -alpha_coef)) / alpha_coef + else: + logratios = chosen_logps - rejected_logps + if self.reference_free: + ref_logratios = torch.tensor([0], dtype=logratios.dtype, device=logratios.device) + else: + ref_logratios = ref_chosen_logps - ref_rejected_logps + + logratios = logratios.to(self.accelerator.device) + ref_logratios = ref_logratios.to(self.accelerator.device) + logits = logratios - ref_logratios + + if self.f_divergence_type == FDivergenceType.JS_DIVERGENCE.value: + # The js-divergence formula: log(2 * u / (1 + u)) + # The divergence difference between the chosen and rejected sample is: + # log(2 * u[w] / (1 + u[w])) - log(2 * u[l] / (1 + u[l])) + # = log(u[w]) - log(u[l]) - (log(1 + u[w]) - log(1 + u[l])) + # where u[w] and u[l] are the policy/reference probability ratios + # for the chosen and rejected samples, respectively. + logits -= F.softplus(chosen_logratios) - F.softplus(rejected_logratios) + + # The beta is a temperature parameter for the DPO loss, typically something in the range of 0.1 to 0.5. + # We ignore the reference model as beta -> 0. The label_smoothing parameter encodes our uncertainty about the + # labels and calculates a conservative DPO loss. + if self.loss_type == "sigmoid": + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) + + elif self.loss_type == "robust": + losses = ( + -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing) + + F.logsigmoid(-self.beta * logits) * self.label_smoothing + ) / (1 - 2 * self.label_smoothing) + + elif self.loss_type == "exo_pair": + # eqn (16) of the EXO paper: https://huggingface.co/papers/2402.00856 + import math + + if self.label_smoothing == 0: + self.label_smoothing = 1e-3 + losses = (self.beta * logits).sigmoid() * ( + F.logsigmoid(self.beta * logits) - math.log(1 - self.label_smoothing) + ) + (-self.beta * logits).sigmoid() * (F.logsigmoid(-self.beta * logits) - math.log(self.label_smoothing)) + + elif self.loss_type == "hinge": + losses = torch.relu(1 - self.beta * logits) + + elif self.loss_type == "ipo": + # eqn (17) of the paper where beta is the regularization parameter for the IPO loss, denoted by tau in the paper. + losses = (logits - 1 / (2 * self.beta)) ** 2 + + elif self.loss_type == "bco_pair": + chosen_logratios = chosen_logps - ref_chosen_logps + rejected_logratios = rejected_logps - ref_rejected_logps + chosen_rewards = self.beta * chosen_logratios + rejected_rewards = self.beta * rejected_logratios + rewards = torch.cat((chosen_rewards, rejected_rewards), 0).mean().detach() + self.running.update(rewards) + delta = self.running.mean + losses = -F.logsigmoid((self.beta * chosen_logratios) - delta) - F.logsigmoid( + -(self.beta * rejected_logratios - delta) + ) + + elif self.loss_type == "sppo_hard": + # In the paper (https://huggingface.co/papers/2405.00675), SPPO employs a soft probability approach, + # estimated using the PairRM score. The probability calculation is conducted outside of the trainer class. + # The version described here is the hard probability version, where P in Equation (4.7) of Algorithm 1 is + # set to 1 for the winner and 0 for the loser. + a = chosen_logps - ref_chosen_logps + b = rejected_logps - ref_rejected_logps + losses = (a - 0.5 / self.beta) ** 2 + (b + 0.5 / self.beta) ** 2 + + elif self.loss_type == "nca_pair": + chosen_rewards = (chosen_logps - ref_chosen_logps) * self.beta + rejected_rewards = (rejected_logps - ref_rejected_logps) * self.beta + losses = ( + -F.logsigmoid(chosen_rewards) + - 0.5 * F.logsigmoid(-chosen_rewards) + - 0.5 * F.logsigmoid(-rejected_rewards) + ) + + elif self.loss_type == "aot_pair": + chosen_logratios = chosen_logps - ref_chosen_logps + rejected_logratios = rejected_logps - ref_rejected_logps + chosen_logratios_sorted, _ = torch.sort(chosen_logratios, dim=0) + rejected_logratios_sorted, _ = torch.sort(rejected_logratios, dim=0) + delta = chosen_logratios_sorted - rejected_logratios_sorted + losses = ( + -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * delta) * self.label_smoothing + ) + + elif self.loss_type == "aot": + logratios = chosen_logps - rejected_logps + ref_logratios = ref_chosen_logps - ref_rejected_logps + logratios_sorted, _ = torch.sort(logratios, dim=0) + ref_logratios_sorted, _ = torch.sort(ref_logratios, dim=0) + delta = logratios_sorted - ref_logratios_sorted + losses = ( + -F.logsigmoid(self.beta * delta) * (1 - self.label_smoothing) + - F.logsigmoid(-self.beta * delta) * self.label_smoothing + ) + + elif self.loss_type == "apo_zero": + # Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are better than your model's default output + losses_chosen = 1 - F.sigmoid(self.beta * chosen_logratios) # Increase chosen likelihood + losses_rejected = F.sigmoid(self.beta * rejected_logratios) # Decrease rejected likelihood + losses = losses_chosen + losses_rejected + + elif self.loss_type == "apo_down": + # Eqn (8) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are worse than your model's default output. + # Decrease chosen likelihood and decrease rejected likelihood more + losses_chosen = F.sigmoid(self.beta * chosen_logratios) + losses_rejected = 1 - F.sigmoid(self.beta * (chosen_logratios - rejected_logratios)) + losses = losses_chosen + losses_rejected + + elif self.loss_type == "discopop": + # Eqn (5) of the DiscoPOP paper (https://huggingface.co/papers/2406.08414) + # This loss was discovered with LLM discovery + logratios = chosen_logps - rejected_logps + ref_logratios = ref_chosen_logps - ref_rejected_logps + logits = logratios - ref_logratios + logits = logits * self.beta + # Modulate the mixing coefficient based on the log ratio magnitudes + log_ratio_modulation = torch.sigmoid(logits / self.args.discopop_tau) + logistic_component = -F.logsigmoid(logits) + exp_component = torch.exp(-logits) + # Blend between logistic and exponential component based on log ratio modulation + losses = logistic_component * (1 - log_ratio_modulation) + exp_component * log_ratio_modulation + + else: + raise ValueError( + f"Unknown loss type: {self.loss_type}. Should be one of ['sigmoid', 'hinge', 'ipo', 'exo_pair', " + "'nca_pair', 'robust', 'bco_pair', 'sppo_hard', 'aot', 'aot_pair', 'discopop', 'apo_zero', 'apo_down']" + ) + + chosen_rewards = self.beta * (chosen_logps.to(device) - ref_chosen_logps.to(device)).detach() + rejected_rewards = self.beta * (rejected_logps.to(device) - ref_rejected_logps.to(device)).detach() + + return losses, chosen_rewards, rejected_rewards + + def _compute_loss_liger(self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]]): + unwrapped_model = self.accelerator.unwrap_model(model) + concatenated_batch = self.concatenated_inputs(batch, padding_value=self.padding_value) + + model_kwargs = {} + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + # Add the pixel values and attention masks for vision models + if "pixel_values" in concatenated_batch: + model_kwargs["pixel_values"] = concatenated_batch["pixel_values"] + if "pixel_attention_mask" in concatenated_batch: + model_kwargs["pixel_attention_mask"] = concatenated_batch["pixel_attention_mask"] + if "image_sizes" in concatenated_batch: + model_kwargs["image_sizes"] = concatenated_batch["image_sizes"] + + prompt_attention_mask = concatenated_batch["prompt_attention_mask"] + completion_attention_mask = concatenated_batch["completion_attention_mask"] + + if self.is_encoder_decoder: + # 1. Get encoder outputs + encoder_outputs = unwrapped_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + # 2. Prepare decoder inputs + decoder_input_ids = shift_tokens_right( + concatenated_batch["completion_input_ids"], + unwrapped_model.config.decoder_start_token_id, + ) + # 3. Get decoder outputs + decoder_outputs = unwrapped_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + hidden_states = decoder_outputs.last_hidden_state + + ref_hidden_states = None + if not self.reference_free and self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + ref_encoder_outputs = unwrapped_ref_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + ref_decoder_outputs = unwrapped_ref_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + ref_hidden_states = ref_decoder_outputs.last_hidden_state + elif not self.reference_free: + with self.null_ref_context(): + ref_encoder_outputs = unwrapped_model.get_encoder()( + concatenated_batch["prompt_input_ids"], + attention_mask=concatenated_batch["prompt_attention_mask"], + return_dict=True, + ) + ref_decoder_outputs = unwrapped_model.get_decoder()( + input_ids=decoder_input_ids, + attention_mask=concatenated_batch["completion_attention_mask"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + encoder_attention_mask=concatenated_batch["prompt_attention_mask"], + use_cache=False, + ) + ref_hidden_states = ref_decoder_outputs.last_hidden_state + + labels = concatenated_batch["completion_input_ids"] + loss_mask = completion_attention_mask.bool() + else: + # For decoder-only models + input_ids = torch.cat( + (concatenated_batch["prompt_input_ids"], concatenated_batch["completion_input_ids"]), dim=1 + ) + attention_mask = torch.cat( + (concatenated_batch["prompt_attention_mask"], concatenated_batch["completion_attention_mask"]), + dim=1, + ) + # Mask the prompt but not the completion for the loss + loss_mask = torch.cat( + (torch.zeros_like(prompt_attention_mask), completion_attention_mask), + dim=1, + ) + + # Flush and truncate + if self.max_length is not None and self.max_length < attention_mask.size(1): + if self.truncation_mode == "keep_start": + # Flush left to reduce the memory usage + # [[0, 0, x, x, x, x], -> [[x, x, x, x], + # [0, x, x, x, 0, 0]] [x, x, x, 0]] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + attention_mask = attention_mask[:, : self.max_length] + input_ids = input_ids[:, : self.max_length] + loss_mask = loss_mask[:, : self.max_length] + elif self.truncation_mode == "keep_end": + # Flush right before truncating left, then flush left + # [[0, 0, x, x, x, x], -> [[0, 0, x, x], + # [0, x, x, x, 0, 0]] [0, x, x, x]] + attention_mask, input_ids, loss_mask = flush_right(attention_mask, input_ids, loss_mask) + input_ids = input_ids[:, -self.max_length :] + attention_mask = attention_mask[:, -self.max_length :] + loss_mask = loss_mask[:, -self.max_length :] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + else: + raise ValueError( + f"Unknown truncation mode: '{self.truncation_mode}'. Should be one of ['keep_end', " + "'keep_start']." + ) + else: + # Flush left to reduce the memory usage + # [[0, 0, x, x, x, x], -> [[x, x, x, x], + # [0, x, x, x, 0, 0]] [x, x, x, 0]] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + + # Add logits_to_keep optimization + if self.use_logits_to_keep: + first_compute_index = loss_mask.nonzero(as_tuple=True)[1].min() + logits_to_keep = (loss_mask.shape[1] - first_compute_index).item() + 1 + model_kwargs["logits_to_keep"] = logits_to_keep + + model_kwargs["output_hidden_states"] = True + + # Add padding-free training support + if self.padding_free: + input_ids = input_ids[attention_mask.bool()].unsqueeze(0) + loss_mask = loss_mask[attention_mask.bool()].unsqueeze(0) + position_ids = attention_mask.cumsum(1)[attention_mask.bool()].unsqueeze(0) - 1 + model_kwargs["position_ids"] = position_ids + else: + model_kwargs["attention_mask"] = attention_mask + + # Get the base model outputs (before LM head) + if hasattr(unwrapped_model, "get_decoder"): + base_model = unwrapped_model.get_decoder() + else: + base_model = getattr(unwrapped_model, self.args.base_model_attribute_name, unwrapped_model) + + outputs = base_model( + input_ids, + use_cache=False, + **model_kwargs, + ) + hidden_states = outputs.last_hidden_state[:, :-1] + + # Get reference hidden states if needed + ref_hidden_states = None + if not self.reference_free and self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + if hasattr(unwrapped_ref_model, "get_decoder"): + ref_base_model = unwrapped_ref_model.get_decoder() + else: + ref_base_model = getattr( + unwrapped_ref_model, self.args.base_model_attribute_name, unwrapped_ref_model + ) + + ref_outputs = ref_base_model( + input_ids, + use_cache=False, + **model_kwargs, + ) + ref_hidden_states = ref_outputs.last_hidden_state[:, :-1] + elif not self.reference_free: + if hasattr(unwrapped_model, "get_decoder"): + ref_base_model = unwrapped_model.get_decoder() + else: + ref_base_model = getattr(unwrapped_model, self.args.base_model_attribute_name, unwrapped_model) + with self.null_ref_context(): + ref_outputs = ref_base_model( + input_ids, + attention_mask=attention_mask, + use_cache=False, + **model_kwargs, + ) + ref_hidden_states = ref_outputs.last_hidden_state[:, :-1] + + masked_input_ids = torch.where(loss_mask != 0, input_ids, self.label_pad_token_id) + labels = masked_input_ids[:, 1:] # Shift right for casual LM + + # Get the LM head + lm_head = unwrapped_model.get_output_embeddings() + + # Get reference model weights if needed + ref_weight = None + ref_bias = None + if not self.reference_free: + if self.ref_model is not None: + unwrapped_ref_model = self.accelerator.unwrap_model(self.ref_model) + ref_lm_head = unwrapped_ref_model.get_output_embeddings() + else: + with self.null_ref_context(): + ref_lm_head = unwrapped_model.get_output_embeddings() + ref_weight = ref_lm_head.weight + ref_bias = ref_lm_head.bias if hasattr(ref_lm_head, "bias") else None + + # Compute loss using Liger kernel + loss_output = self.dpo_loss_fn( + lm_head.weight, + hidden_states, + labels, + bias=lm_head.bias if hasattr(lm_head, "bias") else None, + ref_input=ref_hidden_states if not self.reference_free else None, + ref_weight=ref_weight if not self.reference_free else None, + ref_bias=ref_bias if not self.reference_free else None, + ) + ( + loss, + (chosen_logps, rejected_logps, chosen_logits_mean, rejected_logits_mean, nll_loss, *aux_outputs), + ) = loss_output + + output = { + "loss": loss, + "chosen_logps": chosen_logps, + "rejected_logps": rejected_logps, + "mean_chosen_logits": chosen_logits_mean, + "mean_rejected_logits": rejected_logits_mean, + "nll_loss": nll_loss, + "chosen_rewards": aux_outputs[0], + "rejected_rewards": aux_outputs[1], + } + if self.aux_loss_enabled: + output["aux_loss"] = outputs.aux_loss + + return output + + def concatenated_forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]], is_ref_model: bool = False + ): + """ + Runs the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. + + We do this to avoid doing two forward passes, because it's faster for FSDP. + + Args: + model: + Model to run the forward pass on. + batch: + Batch of input data. + is_ref_model: + Whether this method is being called for the reference model. If `True`, length desensitization is not + applied. + """ + num_examples = batch["prompt_input_ids"].shape[0] + + concatenated_batch = self.concatenated_inputs(batch, padding_value=self.padding_value) + + model_kwargs = {"use_cache": False} + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + # Add the pixel values and attention masks for vision models + if "pixel_values" in concatenated_batch: + model_kwargs["pixel_values"] = concatenated_batch["pixel_values"] + if "pixel_attention_mask" in concatenated_batch: + model_kwargs["pixel_attention_mask"] = concatenated_batch["pixel_attention_mask"] + if "image_sizes" in concatenated_batch: + model_kwargs["image_sizes"] = concatenated_batch["image_sizes"] + + prompt_input_ids = concatenated_batch["prompt_input_ids"] + prompt_attention_mask = concatenated_batch["prompt_attention_mask"] + completion_input_ids = concatenated_batch["completion_input_ids"] + completion_attention_mask = concatenated_batch["completion_attention_mask"] + if self.is_encoder_decoder: + labels = completion_input_ids + labels[completion_attention_mask == 0] = self.label_pad_token_id + outputs = model( + input_ids=prompt_input_ids, + attention_mask=prompt_attention_mask, + labels=labels, # we need the labels for the logits to be returned + **model_kwargs, + ) + logits = outputs.logits + loss_mask = completion_attention_mask.bool() + else: + # Concatenate the prompt and completion inputs + input_ids = torch.cat((prompt_input_ids, completion_input_ids), dim=1) + attention_mask = torch.cat((prompt_attention_mask, completion_attention_mask), dim=1) + # Mask the prompt but not the completion for the loss + loss_mask = torch.cat( + (torch.zeros_like(prompt_attention_mask), completion_attention_mask), + dim=1, + ) + + # Flush and truncate + if self.max_length is not None and self.max_length < attention_mask.size(1): + if self.truncation_mode == "keep_start": + # Flush left to reduce the memory usage + # [[0, 0, x, x, x, x], -> [[x, x, x, x], + # [0, x, x, x, 0, 0]] [x, x, x, 0]] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + attention_mask = attention_mask[:, : self.max_length] + input_ids = input_ids[:, : self.max_length] + loss_mask = loss_mask[:, : self.max_length] + elif self.truncation_mode == "keep_end": + # Flush right before truncating left, then flush left + # [[0, 0, x, x, x, x], -> [[0, 0, x, x], + # [0, x, x, x, 0, 0]] [0, x, x, x]] + attention_mask, input_ids, loss_mask = flush_right(attention_mask, input_ids, loss_mask) + input_ids = input_ids[:, -self.max_length :] + attention_mask = attention_mask[:, -self.max_length :] + loss_mask = loss_mask[:, -self.max_length :] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + else: + raise ValueError( + f"Unknown truncation mode: '{self.truncation_mode}'. Should be one of ['keep_end', " + "'keep_start']." + ) + else: + # Flush left to reduce the memory usage + # [[0, 0, x, x, x, x], -> [[x, x, x, x], + # [0, x, x, x, 0, 0]] [x, x, x, 0]] + attention_mask, input_ids, loss_mask = flush_left(attention_mask, input_ids, loss_mask) + + if self.use_logits_to_keep: + # Compute logits_to_keep based on loss_mask pattern: + # [[0, 0, 0, x, x, x, x], + # [0, 0, 0, x, x, x, 0]] + # ^ start computing logits from here ([:, -(7-3+1):]) + first_compute_index = loss_mask.nonzero(as_tuple=True)[1].min() + logits_to_keep = (loss_mask.shape[1] - first_compute_index).item() + 1 # +1 for the first label + model_kwargs["logits_to_keep"] = logits_to_keep + + model_kwargs["output_hidden_states"] = True + + if self.padding_free: + # Flatten the input_ids, position_ids, and loss_mask + # input_ids = [[a, b, c, 0], -> input_ids = [[a, b, c, d, e, f, g]] + # [d, e, f, g]] position_ids = [[0, 1, 2, 0, 1, 2, 3]] + input_ids = input_ids[attention_mask.bool()].unsqueeze(0) + loss_mask = loss_mask[attention_mask.bool()].unsqueeze(0) + position_ids = attention_mask.cumsum(1)[attention_mask.bool()].unsqueeze(0) - 1 + model_kwargs["position_ids"] = position_ids + else: + model_kwargs["attention_mask"] = attention_mask + + outputs = model(input_ids, **model_kwargs) + logits = outputs.logits + + # Offset the logits by one to align with the labels + labels = torch.roll(input_ids, shifts=-1, dims=1) + loss_mask = torch.roll(loss_mask, shifts=-1, dims=1).bool() + + if self.use_logits_to_keep: + # Align labels with logits + # logits: -, -, [x2, x3, x4, x5, x6] + # ^ --------- ^ after logits[:, :-1, :] + # labels: [y0, y1, y2, y3, y4, y5, y6] + # ^ --------- ^ with logits_to_keep=4, [:, -4:] + # loss_mask: [0, 0, 0, 1, 1, 1, 1] + labels = labels[:, -logits_to_keep:] + loss_mask = loss_mask[:, -logits_to_keep:] + + if logits.shape[:2] != labels.shape[:2]: + # for llava, the returned logits include the image tokens (placed before the text tokens) + seq_len = labels.shape[1] + logits = logits[:, -seq_len:] + + # Compute the log probabilities of the labels + labels[~loss_mask] = 0 # dummy token; we'll ignore the losses on these tokens later + per_token_logps = selective_log_softmax(logits, labels) + per_token_logps[~loss_mask] = 0 + per_token_logps = torch.roll(per_token_logps, shifts=1, dims=1) + + if self.padding_free: + # Unflatten the per_token_logps (shape: [1, sum_seq_len] -> [batch_size, seq_len]) + batch_size, seq_len = attention_mask.shape + per_token_logps_ = torch.zeros( + batch_size, seq_len, device=outputs.logits.device, dtype=outputs.logits.dtype + ) + per_token_logps_[attention_mask.bool()] = per_token_logps + per_token_logps = per_token_logps_ + + all_logps = per_token_logps[:, 1:].sum(-1) + + output = {} + + if self.use_weighting: + with torch.no_grad(): + # Eq (2) of the WPO paper: https://huggingface.co/papers/2406.11827 + logprobs = F.log_softmax(logits, dim=-1) + weights_adjustment_factor = torch.logsumexp(2 * logprobs, dim=-1) # same as sum(probs**2) in log space + per_token_logps_adjusted = per_token_logps - weights_adjustment_factor + all_weights = (per_token_logps_adjusted * loss_mask).sum(-1) / loss_mask.sum(-1) + chosen_weights = all_weights[:num_examples] + rejected_weights = all_weights[num_examples:] + output["policy_weights"] = torch.clamp(torch.exp(chosen_weights + rejected_weights), max=1) + + if self.args.rpo_alpha is not None: + # Only use the chosen logits for the RPO loss + chosen_logits = logits[:num_examples, :-1] if not self.is_encoder_decoder else logits[:num_examples] + chosen_labels = labels[:num_examples, :-1] if not self.is_encoder_decoder else labels[:num_examples] + + # Compute the log probabilities of the labels + output["nll_loss"] = F.cross_entropy( + torch.flatten(chosen_logits, end_dim=1), torch.flatten(chosen_labels, end_dim=1), ignore_index=0 + ) + + if self.loss_type == "ipo": + all_logps = all_logps / loss_mask.sum(-1) + + if self.args.ld_alpha is not None and not is_ref_model: + # Compute response lengths based on loss_mask + completion_lengths = loss_mask.sum(dim=1) + + chosen_lengths = completion_lengths[:num_examples] + rejected_lengths = completion_lengths[num_examples:] + public_lengths = torch.min(chosen_lengths, rejected_lengths) # l_p in the paper + public_lengths = torch.cat([public_lengths, public_lengths], dim=0) + + seq_len = per_token_logps.size(1) + position_ids = torch.arange(seq_len, device=per_token_logps.device).expand_as(per_token_logps) + + ld_mask = position_ids < public_lengths.unsqueeze(1) + mask = position_ids < completion_lengths.unsqueeze(1) + + front_mask = (ld_mask & mask).float() + rear_mask = (~ld_mask & mask).float() + front_logps = (per_token_logps * front_mask).sum(dim=1) + rear_logps = (per_token_logps * rear_mask).sum(dim=1) + + all_logps = front_logps + self.args.ld_alpha * rear_logps + + output["chosen_logps"] = all_logps[:num_examples] + output["rejected_logps"] = all_logps[num_examples:] + + # Compute the mean logits + if self.padding_free: + # position_ids contains a sequence of range identifiers (e.g., [[0, 1, 2, 0, 1, 2, 3, ...]]). + # There are 2*num_examples ranges in total: the first half corresponds to the chosen tokens, + # and the second half to the rejected tokens. + # To find the start of the rejected tokens, we look for the num_examples+1-th zero in pos_id. + split_idx = (position_ids == 0).nonzero(as_tuple=True)[1][num_examples] + mean_chosen_logits = logits[0, :split_idx][loss_mask[0, :split_idx]].mean() + mean_rejected_logits = logits[0, split_idx:][loss_mask[0, split_idx:]].mean() + else: + mean_chosen_logits = logits[:num_examples][loss_mask[:num_examples]].mean() + mean_rejected_logits = logits[num_examples:][loss_mask[num_examples:]].mean() + + output["mean_chosen_logits"] = mean_chosen_logits + output["mean_rejected_logits"] = mean_rejected_logits + + if self.aux_loss_enabled: + output["aux_loss"] = outputs.aux_loss + + return output + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, Union[list, torch.LongTensor]], + train_eval: Literal["train", "eval"] = "train", + ): + """Compute the DPO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + + if self.args.use_liger_loss: + model_output = self._compute_loss_liger(model, batch) + losses = model_output["loss"] + chosen_rewards = model_output["chosen_rewards"] + rejected_rewards = model_output["rejected_rewards"] + else: + model_output = self.concatenated_forward(model, batch) + + # if ref_chosen_logps and ref_rejected_logps in batch use them, otherwise use the reference model + if "ref_chosen_logps" in batch and "ref_rejected_logps" in batch: + ref_chosen_logps = batch["ref_chosen_logps"] + ref_rejected_logps = batch["ref_rejected_logps"] + else: + ref_chosen_logps, ref_rejected_logps = self.compute_ref_log_probs(batch) + + losses, chosen_rewards, rejected_rewards = self.dpo_loss( + model_output["chosen_logps"], model_output["rejected_logps"], ref_chosen_logps, ref_rejected_logps + ) + reward_accuracies = (chosen_rewards > rejected_rewards).float() + + if self.args.rpo_alpha is not None: + losses = losses + self.args.rpo_alpha * model_output["nll_loss"] # RPO loss from V3 of the paper + + if self.use_weighting: + losses = losses * model_output["policy_weights"] + + if self.aux_loss_enabled: + losses = losses + self.aux_loss_coef * model_output["aux_loss"] + + prefix = "eval_" if train_eval == "eval" else "" + metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean().item() + metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean().item() + metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean().item() + metrics[f"{prefix}rewards/margins"] = ( + self.accelerator.gather_for_metrics(chosen_rewards - rejected_rewards).mean().item() + ) + metrics[f"{prefix}logps/chosen"] = ( + self.accelerator.gather_for_metrics(model_output["chosen_logps"]).detach().mean().item() + ) + metrics[f"{prefix}logps/rejected"] = ( + self.accelerator.gather_for_metrics(model_output["rejected_logps"]).detach().mean().item() + ) + metrics[f"{prefix}logits/chosen"] = ( + self.accelerator.gather_for_metrics(model_output["mean_chosen_logits"]).detach().mean().item() + ) + metrics[f"{prefix}logits/rejected"] = ( + self.accelerator.gather_for_metrics(model_output["mean_rejected_logits"]).detach().mean().item() + ) + if self.args.rpo_alpha is not None: + metrics[f"{prefix}nll_loss"] = ( + self.accelerator.gather_for_metrics(model_output["nll_loss"]).detach().mean().item() + ) + if self.aux_loss_enabled: + metrics[f"{prefix}aux_loss"] = ( + self.accelerator.gather_for_metrics(model_output["aux_loss"]).detach().mean().item() + ) + + return losses.mean(), metrics + + def compute_loss( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + return_outputs=False, + num_items_in_batch=None, + ) -> Union[torch.Tensor, tuple[torch.Tensor, dict[str, torch.Tensor]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") + + # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: + loss = loss.to(self.args.device) + # force log the metrics + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return loss, metrics + + return loss + + def generate_from_model_and_ref(self, model, batch: dict[str, torch.LongTensor]) -> tuple[str, str]: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.padding_value, + ) + + # if ref_output in batch use that otherwise use the reference model + if "ref_output" in batch: + ref_output = batch["ref_output"] + else: + if self.ref_model is None: + with self.null_ref_context(): + ref_output = self.model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.padding_value, + ) + else: + ref_output = self.ref_model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.padding_value, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.padding_value) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + ref_output = pad_to_length(ref_output, self.max_length, self.padding_value) + ref_output_decoded = self.processing_class.batch_decode(ref_output, skip_special_tokens=True) + + return policy_output_decoded, ref_output_decoded + + def prediction_step( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[list[str]] = None, + ): + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval") + + # force log the metrics + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return loss.detach(), None, None + + # logits for the chosen and rejected samples from model + logits_dict = { + "eval_logits/chosen": metrics["eval_logits/chosen"], + "eval_logits/rejected": metrics["eval_logits/rejected"], + } + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: Optional[bool] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, random_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy", "Ref Model"], + data=[ + [prompt, pol[len(prompt) :], ref[len(prompt) :]] + for prompt, pol, ref in zip( + random_batch_dataset["prompt"], policy_output_decoded, ref_output_decoded + ) + ], + ) + if "wandb" in self.args.report_to and self.accelerator.is_main_process: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float` or `None`, *optional*, defaults to `None`): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[key] = torch.tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent( + """\ + @inproceedings{rafailov2023direct, + title = {{Direct Preference Optimization: Your Language Model is Secretly a Reward Model}}, + author = {Rafael Rafailov and Archit Sharma and Eric Mitchell and Christopher D. Manning and Stefano Ermon and Chelsea Finn}, + year = 2023, + booktitle = {Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 - 16, 2023}, + url = {http://papers.nips.cc/paper_files/paper/2023/hash/a85b405ed65c6477a4fe8302b5e06ce7-Abstract-Conference.html}, + editor = {Alice Oh and Tristan Naumann and Amir Globerson and Kate Saenko and Moritz Hardt and Sergey Levine}, + }""" + ) + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="DPO", + trainer_citation=citation, + paper_title="Direct Preference Optimization: Your Language Model is Secretly a Reward Model", + paper_id="2305.18290", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/gkd_config.py b/lib/python3.12/site-packages/trl/trainer/gkd_config.py new file mode 100644 index 0000000000000000000000000000000000000000..c299719e0b6558f789e44ad4e465ae60ff74356a --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/gkd_config.py @@ -0,0 +1,112 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Any, Optional + +from transformers import TrainingArguments + +from .sft_config import SFTConfig + + +@dataclass +class GKDConfig(SFTConfig): + """ + Configuration class for [`GKDTrainer`]. + + This class includes only the parameters that are specific to GKD training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`SFTConfig`] documentation. + + Args: + temperature (`float`, *optional*, defaults to `0.9`): + Temperature for sampling. The higher the temperature, the more random the completions. + lmbda (`float`, *optional*, defaults to `0.5`): + Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy + student-generated outputs). + beta (`float`, *optional*, defaults to `0.5`): + Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence loss. When + beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL Divergence. + max_new_tokens (`int`, *optional*, defaults to `128`): + Maximum number of tokens to generate per completion. + teacher_model_name_or_path (`str` or `None`, *optional*, defaults to `None`): + Model name or path of the teacher model. If `None`, the teacher model will be the same as the model being + trained. + teacher_model_init_kwargs (`dict[str, Any]]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model + from a string. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + seq_kd (`bool`, *optional*, defaults to `False`): + Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised FT on + teacher-generated output). + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["teacher_model_init_kwargs"] + + temperature: float = field( + default=0.9, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + lmbda: float = field( + default=0.5, + metadata={ + "help": "Lambda parameter that controls the student data fraction (i.e., the proportion of on-policy " + "student-generated outputs)." + }, + ) + beta: float = field( + default=0.5, + metadata={ + "help": "Interpolation coefficient between `0.0` and `1.0` of the Generalized Jensen-Shannon Divergence " + "loss. When beta is `0.0`, the loss is the KL divergence. When beta is `1.0`, the loss is the Inverse KL " + "Divergence." + }, + ) + max_new_tokens: int = field( + default=128, + metadata={"help": "Maximum number of tokens to generate per completion."}, + ) + teacher_model_name_or_path: Optional[str] = field( + default=None, + metadata={ + "help": "Model name or path of the teacher model. If `None`, the teacher model will be the same as the " + "model being trained." + }, + ) + teacher_model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " + "teacher model from a string." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropouts in `model`."}, + ) + seq_kd: bool = field( + default=False, + metadata={ + "help": "Seq_kd parameter that controls whether to perform Sequence-Level KD (can be viewed as supervised " + "FT on teacher-generated output)." + }, + ) + + def __post_init__(self): + super().__post_init__() + # check lmbda and beta are in the range [0, 1] + if self.lmbda < 0.0 or self.lmbda > 1.0: + raise ValueError("lmbda must be in the range [0.0, 1.0].") + if self.beta < 0.0 or self.beta > 1.0: + raise ValueError("beta must be in the range [0.0, 1.0].") diff --git a/lib/python3.12/site-packages/trl/trainer/gkd_trainer.py b/lib/python3.12/site-packages/trl/trainer/gkd_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..d1ce3ab5406e766b6ecc82a2d61edb922b1a9603 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/gkd_trainer.py @@ -0,0 +1,369 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import random +import textwrap +from typing import Any, Callable, Optional, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from datasets import Dataset +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + is_wandb_available, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_peft_available + +from ..models import prepare_deepspeed +from ..models.utils import unwrap_model_for_generation +from .gkd_config import GKDConfig +from .sft_trainer import SFTTrainer +from .utils import ( + DataCollatorForChatML, + disable_dropout_in_model, + empty_cache, + generate_model_card, + get_comet_experiment_url, +) + + +if is_peft_available(): + from peft import PeftConfig + +if is_wandb_available(): + import wandb + + +class GKDTrainer(SFTTrainer): + _tag_names = ["trl", "gkd"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + teacher_model: Union[PreTrainedModel, nn.Module, str] = None, + args: Optional[GKDConfig] = None, + data_collator: Optional[DataCollator] = None, # type: ignore + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional["PeftConfig"] = None, + formatting_func: Optional[Callable] = None, + ): + # add remove_unused_columns=False to the dataclass args + args.remove_unused_columns = False + data_collator = DataCollatorForChatML(tokenizer=processing_class, max_length=args.max_length) + + super().__init__( + model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + peft_config=peft_config, + formatting_func=formatting_func, + ) + + if args.teacher_model_init_kwargs is None: + teacher_model_init_kwargs = {} + elif not isinstance(teacher_model, str): + raise ValueError( + "You passed teacher_model_init_kwargs to the GKDConfig, but your teacher_model is already instantiated." + ) + else: + teacher_model_init_kwargs = args.teacher_model_init_kwargs + teacher_model_init_kwargs["torch_dtype"] = ( + teacher_model_init_kwargs["torch_dtype"] + if teacher_model_init_kwargs["torch_dtype"] in ["auto", None] + else getattr(torch, teacher_model_init_kwargs["torch_dtype"]) + ) + + if isinstance(teacher_model, str): + teacher_model = AutoModelForCausalLM.from_pretrained(teacher_model, **teacher_model_init_kwargs) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(self.model) + + if self.is_deepspeed_enabled: + self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator) + else: + self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True) + + self.lmbda = args.lmbda + self.beta = args.beta + self.temperature = args.temperature + self.seq_kd = args.seq_kd + + self.generation_config = GenerationConfig( + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + do_sample=True, + top_k=0, + use_cache=False if args.gradient_checkpointing else True, + pad_token_id=self.processing_class.pad_token_id, + ) + # Set custom EOS tokens if they are specified by the model's generation + # config. This is important for models with the Llama 3 chat template, + # which use special tokens <|eot_id|> and <|eom_id|> to mark the end of + # turns or messages. + if ( + hasattr(self.model.generation_config, "eos_token_id") + and self.model.generation_config.eos_token_id is not None + ): + self.generation_config.eos_token_id = self.model.generation_config.eos_token_id + + @staticmethod + def generalized_jsd_loss( + student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0, reduction="batchmean" + ): + """ + Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation using F.kl_div. See Eq. (1) + of https://huggingface.co/papers/2306.13649 for the definition. + + Args: + student_logits: + Tensor of shape (batch_size, sequence_length, vocab_size) + teacher_logits: + Tensor of shape (batch_size, sequence_length, vocab_size) + labels: + Tensor of shape (batch_size, sequence_length) with -100 for padding tokens to ignore when computing + loss + beta: + Interpolation coefficient between 0 and 1 (default: 0.5) + temperature: + Softmax temperature (default: 1.0) + reduction: + Specifies the reduction to apply to the output (default: 'batchmean') + + Returns: + loss: Scalar tensor with the generalized JSD loss + """ + + # Apply temperature scaling + student_logits = student_logits / temperature + teacher_logits = teacher_logits / temperature + + # Compute log probabilities for student and probabilities for teacher + student_log_probs = F.log_softmax(student_logits, dim=-1) + teacher_log_probs = F.log_softmax(teacher_logits, dim=-1) + + if beta == 0: + jsd = F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True) + elif beta == 1: + jsd = F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True) + else: + # Compute the log of the mixture distribution + # log(a + b) = log(exp(log(a)) + exp(log(b))) -> for mixture + beta = torch.tensor(beta, dtype=student_log_probs.dtype) + mixture_log_probs = torch.logsumexp( + torch.stack([student_log_probs + torch.log(1 - beta), teacher_log_probs + torch.log(beta)]), + dim=0, + ) + + # Compute KL divergences using F.kl_div + # PyTorch differs from the standard mathematical definition, so the order of the probability distributions is swapped compared to that defined in the paper. + kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True) + kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True) + + # Compute the Generalized Jensen-Shannon Divergence + jsd = beta * kl_teacher + (1 - beta) * kl_student + + # Masking + if labels is not None: + mask = labels != -100 + jsd = jsd[mask] + + # Apply reduction + if reduction == "batchmean": + return jsd.sum() / mask.sum() if labels is not None else jsd.sum() / (jsd.size(0) * jsd.size(1)) + elif reduction == "sum": + return jsd.sum() + elif reduction == "mean": + return jsd.mean() + else: + return jsd + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + # compute student output + outputs_student = model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + + # compute teacher output in eval mode + self.teacher_model.eval() + with torch.no_grad(): + outputs_teacher = self.teacher_model( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + ) + + # slice the logits for the generated tokens using the inputs["prompts"] lengths + prompt_lengths = inputs["prompts"].shape[1] + shifted_student_logits = outputs_student.logits[:, prompt_lengths - 1 : -1, :] + shifted_teacher_logits = outputs_teacher.logits[:, prompt_lengths - 1 : -1, :] + shifted_labels = inputs["labels"][:, prompt_lengths:] + + # compute loss + loss = self.generalized_jsd_loss( + student_logits=shifted_student_logits, + teacher_logits=shifted_teacher_logits, + labels=shifted_labels, + beta=self.beta, + ) + + # empty cache + empty_cache() + + # Return loss + return (loss, outputs_student) if return_outputs else loss + + @staticmethod + def generate_on_policy_outputs(model, inputs, generation_config, pad_token_id=None): + # Generate output with respect to the prompt only + generated_outputs = model.generate( + input_ids=inputs["prompts"], + attention_mask=inputs.get("prompt_attention_mask", None), + generation_config=generation_config, + return_dict_in_generate=True, + ) + + # Get the generated token IDs + generated_tokens = generated_outputs.sequences + # Calculate new attention mask + new_attention_mask = torch.ones_like(generated_tokens) + new_labels = generated_tokens.clone() + + # If there's pad_token_id, set attention mask to 0 for padding tokens + if pad_token_id is not None: + new_labels[new_labels == pad_token_id] = -100 + new_attention_mask[generated_tokens == pad_token_id] = 0 + + return generated_tokens, new_attention_mask, new_labels + + def training_step( + self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None + ) -> torch.Tensor: + """ + Perform a training step for the Generalized Knowledge Distillation (GKD) model. + + This method implements the on-policy learning approach described in the GKD paper. With probability + `self.lmbda`, it generates new responses using the student model, which are then used for training instead of + the original inputs. + """ + if self.seq_kd: + with unwrap_model_for_generation(self.teacher_model, self.accelerator) as unwrapped_model: + new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs( + unwrapped_model, inputs, self.generation_config, self.processing_class.pad_token_id + ) + inputs["input_ids"] = new_input_ids + inputs["attention_mask"] = new_attention_mask + inputs["labels"] = new_labels + if random.random() <= self.lmbda: + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model: + new_input_ids, new_attention_mask, new_labels = self.generate_on_policy_outputs( + unwrapped_model, inputs, self.generation_config, self.processing_class.pad_token_id + ) + inputs["input_ids"] = new_input_ids + inputs["attention_mask"] = new_attention_mask + inputs["labels"] = new_labels + + loss = super().training_step(model, inputs, num_items_in_batch) + return loss + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{agarwal2024on-policy, + title = {{On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes}}, + author = {Rishabh Agarwal and Nino Vieillard and Yongchao Zhou and Piotr Stanczyk and Sabela Ramos Garea and Matthieu Geist and Olivier Bachem}, + year = 2024, + booktitle = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=3zKtaqxLhW}, + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="GKD", + trainer_citation=citation, + paper_title="On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes", + paper_id="2306.13649", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/grpo_config.py b/lib/python3.12/site-packages/trl/trainer/grpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..b23e96efc2080dc9a65e724c53a7c78442f828e1 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/grpo_config.py @@ -0,0 +1,585 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Optional, Union + +import transformers +from packaging import version +from transformers import TrainingArguments + + +@dataclass +class GRPOConfig(TrainingArguments): + r""" + Configuration class for the [`GRPOTrainer`]. + + This class includes only the parameters that are specific to GRPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model and reference model + + model_init_kwargs (`str`, `dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`GRPOTrainer`] is provided as a string. + disable_dropout (`bool`, *optional*, defaults to `False`): + Whether to disable dropout in the model. This is useful for training with a reference model, as it prevents + the model from generating different logprobs for the same input. + + > Parameters that control the data preprocessing + + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to only keep the column `"prompt"` in the dataset. If you use a custom reward function that + requires any column other than `"prompts"` and `"completions"`, you should keep this to `False`. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left. + num_generations (`int` or `None`, *optional*, defaults to `8`): + Number of generations per prompt to sample. The effective batch size (num_processes * per_device_batch_size + * gradient_accumulation_steps) must be evenly divisible by this value. + max_completion_length (`int` or `None`, *optional*, defaults to `256`): + Maximum length of the generated completion. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. Disabling this option is not compatible + with vLLM generation. + shuffle_dataset (`bool`, *optional*, defaults to `True`): + Whether to shuffle the training dataset. + + > Parameters that control generation + + generation_batch_size: (`int` or `None`, *optional*, defaults to `None`): + Batch size to use for generation. If `None`, it defaults to the effective training batch size: + `per_device_train_batch_size * num_processes * gradient_accumulation_steps`. + steps_per_generations: (`int` or `None`, *optional*, defaults to `None`): + Number of optimization steps per generation. If `None`, it defaults to gradient_accumulation_steps. + temperature (`float`, defaults to `1.0`): + Temperature for sampling. The higher the temperature, the more random the completions. + top_p (`float`, *optional*, defaults to `1.0`): + Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. Set to + `1.0` to consider all tokens. + top_k (`int` or `None`, *optional*, defaults to `None`): + Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, top-k-filtering is + disabled and all tokens are considered. + min_p (`float` or `None`, *optional*, defaults to `None`): + Minimum token probability, which will be scaled by the probability of the most likely token. It must be a + value between `0.0` and `1.0`. Typical values are in the `0.01-0.2` range. + repetition_penalty (`float`, *optional*, defaults to `1.0`): + Float that penalizes new tokens based on whether they appear in the prompt and the generated text so far. + Values > `1.0` encourage the model to use new tokens, while values < `1.0` encourage the model to repeat + tokens. + cache_implementation (`str` or `None`, *optional*, defaults to `None`): + Implementation of the cache method for faster generation when use_vllm is set to False. + generation_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or `SamplingParams` (if + using vLLM) when sampling completions. This can be used to further customize the generation behavior, such + as setting `supress_tokens`, `num_beams`, etc. If it contains keys that conflict with the other generation + parameters (like `min_p`, `top_p`, etc.), they will override them. + + > Parameters that control generation acceleration powered by vLLM + + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for generation + instead of the default model.generate(). Requires `vllm` to be installed. + vllm_mode (`str`, *optional*, defaults to `"server"`): + Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `"server"` or + `"colocate"`. + + - `"server"`: The trainer will send generation requests to a separate vLLM server. Make sure a TRL vLLM + server is running (start with `trl vllm-serve`). + - `"colocate"`: vLLM will run in the same process and share the training GPUs. This avoids the need for a + separate server but may cause resource contention with training. + vllm_guided_decoding_regex (`str` or `None`, *optional*, defaults to `None`): + Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled. + + > Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + vllm_server_base_url (`str` or `None`, *optional*, defaults to `None`): + Base URL for the vLLM server (e.g., `"http://localhost:8000"`). If provided, `vllm_server_host` and + `vllm_server_port` are ignored. + vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`): + Host of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_port (`int`, *optional*, defaults to `8000`): + Port of the vLLM server to connect to. Ignored if `vllm_server_base_url` is provided. + vllm_server_timeout (`float`, *optional*, defaults to `240.0`): + Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up after the + timeout, a `ConnectionError` is raised. + + > Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + + vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`): + Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_gpu_memory_utilization` flag. + vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`): + Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set to + `"colocate"`. If you are using `vllm_mode="server"`, this parameter must be passed separately when + launching the vLLM server via the `--vllm_tensor_parallel_size` flag. + + > Parameters that control the training + + beta (`float`, *optional*, defaults to `0.0`): + KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and improving + training speed. + num_iterations (`int`, *optional*, defaults to `1`): + Number of iterations per batch (denoted as μ in the algorithm). + epsilon (`float`, *optional*, defaults to `0.2`): + Epsilon value for clipping. + delta: (`float` or `None`, *optional*, defaults to `None`): + Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` (default), standard + GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This method is introduced in + the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291). + epsilon_high (`float` or `None`, *optional*, defaults to `None`): + Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the lower-bound + specified in argument `epsilon`. Paper [DAPO](https://huggingface.co/papers/2503.14476) recommends `0.28`. + reward_weights (`list[float]` or `None`, *optional*, defaults to `None`): + Weights for each reward function. Must match the number of reward functions. If `None`, all rewards are + weighted equally with weight `1.0`. + scale_rewards (`bool`, *optional*, defaults to `True`): + Whether to scale the rewards by dividing them by their standard deviation. If `True` (default), the rewards + are normalized by the standard deviation, ensuring they have unit variance. If `False`, no scaling is + applied. The [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) recommends not scaling the rewards, + as scaling by the standard deviation introduces a question-level difficulty bias. + loss_type (`str`, *optional*, defaults to `"bnpo"`): + Specifies the loss formulation to use. Supported values are: + + - `"grpo"`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to + length bias—this approach tends to prefer shorter completions with positive advantages and longer ones + with negative advantages. + - `"bnpo"`: Aggregates token-level losses by normalizing number of active token in the local batch. + Note that normalization is performed over the local batch only, so results may slightly vary depending + on the local batch size, despite a constant effective batch size. When using + `per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss. + - `"dr_grpo"`: Aggregates token-level losses by normalizing with a global constant. This method was + introduced in the [Dr. GRPO paper](https://huggingface.co/papers/2503.20783) to eliminate length bias. + The value of the constant corresponds to `max_completion_length`. + mask_truncated_completions (`bool`, *optional*, defaults to `False`): + When enabled, truncated completions are excluded from the loss calculation, preventing them from being + incorrectly penalized and introducing noise during training. According to the + [DAPO](https://huggingface.co/papers/2503.14476) paper, this is a good practice for training stability. + sync_ref_model (`bool`, *optional*, defaults to `False`): + Whether to synchronize the reference model with the active model every `ref_model_sync_steps` steps, using + the `ref_model_mixup_alpha` parameter. This synchronization originates from the + [TR-DPO](https://huggingface.co/papers/2404.09656) paper. + ref_model_mixup_alpha (`float`, *optional*, defaults to `0.6`): + α parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which controls the mix + between the current policy and the previous reference policy during updates. The reference policy is + updated according to the equation: `π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you + must set `sync_ref_model=True`. + ref_model_sync_steps (`int`, *optional*, defaults to `512`): + τ parameter from the [TR-DPO](https://huggingface.co/papers/2404.09656) paper, which determines how + frequently the current policy is synchronized with the reference policy. To use this parameter, you must + set `sync_ref_model=True`. + use_liger_loss (`bool`, *optional*, defaults to `False`): + Whether to use the Liger GRPO loss. + + > Parameters that control the logging + + log_completions (`bool`, *optional*, defaults to `False`): + Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is installed, + it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`. + num_completions_to_print (`int` or `None`, *optional*, defaults to `None`): + Number of completions to print with `rich`. If `None`, all completions are logged. + wandb_log_unique_prompts (`bool`, *optional*, defaults to `False`): + Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, all prompts + are logged. + """ + + if version.parse(transformers.__version__) >= version.parse("4.51.0"): + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model and reference model + model_init_kwargs: Optional[Union[dict, str]] = field( + default=None, + metadata={ + "help": "Keyword arguments for `transformers.AutoModelForCausalLM.from_pretrained`, used when the `model` " + "argument of the `GRPOTrainer` is provided as a string." + }, + ) + disable_dropout: bool = field( + default=False, + metadata={ + "help": "Whether to disable dropout in the model. This is useful for training with a reference model, as " + "it prevents the model from generating different logprobs for the same input." + }, + ) + + # Parameters that control the data preprocessing + # The default value remove_unused_columns is overwritten from the parent class, because in GRPO we usually rely on + # additional columns to compute the reward + remove_unused_columns: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether to only keep the column 'prompt' in the dataset. If you use a custom reward function " + "that requires any column other than 'prompts' and 'completions', you should keep this to `False`." + }, + ) + max_prompt_length: Optional[int] = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. If the prompt is longer than this value, it will be truncated left." + }, + ) + num_generations: Optional[int] = field( + default=8, + metadata={ + "help": "Number of generations to sample. The effective batch size (num_processes * per_device_batch_size " + "* gradient_accumulation_steps) must be evenly divisible by this value." + }, + ) + max_completion_length: Optional[int] = field( + default=256, + metadata={"help": "Maximum length of the generated completion."}, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation. Disabling this option " + "is not compatible with vLLM generation." + }, + ) + shuffle_dataset: Optional[bool] = field( + default=True, + metadata={"help": "Whether to shuffle the training dataset."}, + ) + + # Parameters that control generation + generation_batch_size: Optional[int] = field( + default=None, + metadata={ + "help": "Batch size to use for generation. If `None`, it defaults to the effective training batch size: " + "`per_device_train_batch_size * num_processes * gradient_accumulation_steps`." + }, + ) + steps_per_generation: Optional[int] = field( + default=None, + metadata={ + "help": "Number of optimization steps per generation. If `None`, it defaults to gradient_accumulation_steps." + }, + ) + temperature: float = field( + default=1.0, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + top_p: float = field( + default=1.0, + metadata={ + "help": "Float that controls the cumulative probability of the top tokens to consider. Must be in (0, 1]. " + "Set to 1.0 to consider all tokens." + }, + ) + top_k: Optional[int] = field( + default=None, + metadata={ + "help": "Number of highest probability vocabulary tokens to keep for top-k-filtering. If `None`, " + "top-k-filtering is disabled and all tokens are considered." + }, + ) + min_p: Optional[float] = field( + default=None, + metadata={ + "help": "Minimum token probability, which will be scaled by the probability of the most likely token. It " + "must be a value between 0.0 and 1.0. Typical values are in the 0.01-0.2 range." + }, + ) + generation_kwargs: Optional[dict] = field( + default=None, + metadata={ + "help": "Additional keyword arguments to pass to `GenerationConfig` (if using transformers) or " + "`SamplingParams` (if using vLLM) when sampling completions. This can be used to further customize the " + "generation behavior, such as setting `supress_tokens`, `num_beams`, etc. If it contains keys that " + "conflict with the other generation parameters (like `min_p`, `top_p`, etc.), they will override them." + }, + ) + repetition_penalty: float = field( + default=1.0, + metadata={ + "help": "Float that penalizes new tokens based on whether they appear in the prompt and the generated " + "text so far. Values > 1.0 encourage the model to use new tokens, while values < 1.0 encourage the model " + "to repeat tokens." + }, + ) + cache_implementation: Optional[str] = field( + default=None, + metadata={"help": "Implementation of the cache method for faster generation when use_vllm is set to False."}, + ) + + # Parameters that control generation acceleration powered by vLLM + use_vllm: bool = field( + default=False, + metadata={ + "help": "Whether to use vLLM for generating completions. If set to `True`, the trainer will use vLLM for " + "generation instead of the default model.generate(). Requires `vllm` to be installed." + }, + ) + vllm_server_base_url: Optional[str] = field( + default=None, + metadata={ + "help": "Base URL for the vLLM server (e.g., 'http://localhost:8000'). If provided, `vllm_server_host` " + "and `vllm_server_port` are ignored." + }, + ) + vllm_mode: str = field( + default="server", + metadata={ + "help": "Mode to use for vLLM integration when `use_vllm` is set to `True`. Must be one of `server` or " + "`'colocate'`. `'server'`: The trainer will send generation requests to a separate vLLM server. Make sure a " + "TRL vLLM server is running (start with `trl vllm-serve`). `'colocate'`: vLLM will run in the same " + "process and share the training GPUs. This avoids the need for a separate server but may cause resource " + "contention with training." + }, + ) + vllm_guided_decoding_regex: Optional[str] = field( + default=None, + metadata={"help": "Regex for vLLM guided decoding. If `None` (default), guided decoding is disabled."}, + ) + + # Parameters that control the vLLM server (only used when `vllm_mode` is `"server"`) + vllm_server_host: str = field( + default="0.0.0.0", + metadata={"help": "Host of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_port: int = field( + default=8000, + metadata={"help": "Port of the vLLM server to connect to. Ignored if vllm_server_base_url is provided."}, + ) + vllm_server_timeout: float = field( + default=240.0, + metadata={ + "help": "Total timeout duration in seconds to wait for the vLLM server to be up. If the server is not up " + "after the timeout, a `ConnectionError` is raised." + }, + ) + + # Parameters that control colocated vLLM execution (only used when `vllm_mode` is `"colocate"`) + vllm_gpu_memory_utilization: float = field( + default=0.3, + metadata={ + "help": "Control the GPU memory utilization for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_gpu_memory_utilization` flag." + }, + ) + vllm_tensor_parallel_size: int = field( + default=1, + metadata={ + "help": "Control the tensor parallel size for vLLM. This setting only applies when `vllm_mode` is set " + "to `'colocate'`. If you are using `vllm_mode='server'`, this parameter must be passed separately when " + "launching the vLLM server via the `--vllm_tensor_parallel_size` flag." + }, + ) + + # Parameters that control the training + beta: float = field( + default=0.0, + metadata={ + "help": "KL coefficient. If `0.0` (default), the reference model is not loaded, reducing memory usage and " + "improving training speed." + }, + ) + num_iterations: int = field( + default=1, + metadata={"help": "Number of iterations per batch (denoted as μ in the algorithm)."}, + ) + epsilon: float = field( + default=0.2, + metadata={"help": "Epsilon value for clipping."}, + ) + delta: Optional[float] = field( + default=None, + metadata={ + "help": "Enables the upper clipping bound in two-sided GRPO loss when set to a float. If `None` " + "(default), standard GRPO clipping is used. Recommended to be greater than `1 + ε` when enabled. This " + "method is introduced in the [INTELLECT-2 tech report](https://huggingface.co/papers/2505.07291)." + }, + ) + epsilon_high: Optional[float] = field( + default=None, + metadata={ + "help": "Upper-bound epsilon value for clipping. If not specified, it defaults to the same value as the " + "lower-bound specified in argument `epsilon`. Paper DAPO recommends `0.28`." + }, + ) + reward_weights: Optional[list[float]] = field( + default=None, + metadata={ + "help": "Weights for each reward function. Must match the number of reward functions. If `None`, all " + "rewards are weighted equally with weight `1.0`." + }, + ) + scale_rewards: bool = field( + default=True, + metadata={ + "help": "Whether to scale the rewards by dividing them by their standard deviation. If `True` (default), " + "the rewards are normalized by the standard deviation, ensuring they have unit variance. If `False`, no " + "scaling is applied. The Dr. GRPO paper recommends not scaling the rewards, as scaling by the standard " + "deviation introduces a question-level difficulty bias." + }, + ) + loss_type: str = field( + default="bnpo", + metadata={ + "help": "Specifies the loss formulation to use. Supported values are `grpo`, `bnpo`, and `dr_grpo`. " + "`'grpo'`: Aggregates token-level losses by normalizing over sequence length. Not recommended due to " + "length bias—this approach tends to prefer shorter completions with positive advantages and longer ones " + "with negative advantages. " + "`'bnpo'`: Aggregates token-level losses by normalizing number of active token in the local batch. " + "Note that normalization is performed over the local batch only, so results may slightly vary depending " + "on the local batch size, despite a constant effective batch size. When using " + "`per_device_train_batch_size==1`, the loss is equivalent to the GRPO loss. " + "`'dr_grpo'`: Aggregates token-level losses by normalizing with a global constant. This method was " + "introduced in the Dr. GRPO paper to eliminate length bias. The value of the constant corresponds to " + "`max_completion_length`." + }, + ) + mask_truncated_completions: bool = field( + default=False, + metadata={ + "help": "When enabled, truncated completions are excluded from the loss calculation, preventing them from " + "being incorrectly penalized and introducing noise during training. According to the DAPO paper, this is " + "a good practice for training stability." + }, + ) + sync_ref_model: bool = field( + default=False, + metadata={ + "help": "Whether to synchronize the reference model with the active model every `ref_model_sync_steps` " + "steps, using the `ref_model_mixup_alpha` parameter." + }, + ) + ref_model_mixup_alpha: float = field( + default=0.6, + metadata={ + "help": "α parameter from the TR-DPO paper, which controls the mix between the current policy and the " + "previous reference policy during updates. The reference policy is updated according to the equation: " + "`π_ref = α * π_θ + (1 - α) * π_ref_prev`. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + ref_model_sync_steps: int = field( + default=512, + metadata={ + "help": "τ parameter from the TR-DPO paper, which determines how frequently the current policy is " + "synchronized with the reference policy. To use this parameter, you must set `sync_ref_model=True`." + }, + ) + use_liger_loss: bool = field( + default=False, + metadata={"help": "Whether to use the Liger GRPO loss."}, + ) + + # Parameters that control the logging + log_completions: bool = field( + default=False, + metadata={ + "help": "Whether to log a sample of (prompt, completion) pairs every `logging_steps` steps. If `rich` is " + "installed, it prints the sample. If `wandb` logging is enabled, it logs it to `wandb`." + }, + ) + num_completions_to_print: Optional[int] = field( + default=None, + metadata={"help": "Number of completions to print with `rich`. If `None`, all completions are logged."}, + ) + wandb_log_unique_prompts: Optional[bool] = field( + default=False, + metadata={ + "help": "Whether to log unique prompts in wandb. If `True`, only unique prompts are logged. If `False`, " + "all prompts are logged." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + num_processes = self.world_size + # The current default effective batch size + if self.generation_batch_size is not None and self.steps_per_generation is not None: + raise ValueError( + "'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time" + ) + + if self.steps_per_generation is None: + self.steps_per_generation = self.gradient_accumulation_steps + + if self.generation_batch_size is None: + self.generation_batch_size = self.per_device_train_batch_size * num_processes * self.steps_per_generation + + if self.generation_batch_size % self.per_device_train_batch_size * num_processes != 0: + raise ValueError( + f"generation_batch_size ({self.generation_batch_size}) must be divisible by the global batch size " + f"({self.per_device_train_batch_size * num_processes})." + ) + + self.steps_per_generation = self.generation_batch_size // (self.per_device_train_batch_size * num_processes) + + # Check if the effective batch size can be divided by the number of generations + if self.num_generations < 2: + raise ValueError( + "GRPO requires at least 2 generations per prompt to calculate the advantages. You provided " + f"{self.num_generations}, which is less than the minimum required." + ) + possible_values = [ + n_gen for n_gen in range(2, self.generation_batch_size + 1) if (self.generation_batch_size) % n_gen == 0 + ] + + if self.num_generations not in possible_values: + raise ValueError( + f"The effective train batch size ({num_processes} x {self.per_device_train_batch_size} x " + f"{self.steps_per_generation}) must be evenly divisible by the number of generations per " + f"prompt ({self.num_generations}). Given the current effective train batch size, the valid values for " + f"the number of generations are: {possible_values}." + ) + if self.eval_strategy != "no": + global_eval_batch_size = self.per_device_eval_batch_size * num_processes + possible_values = [ + n_gen for n_gen in range(2, global_eval_batch_size + 1) if (global_eval_batch_size) % n_gen == 0 + ] + if self.num_generations not in possible_values: + raise ValueError( + f"The global eval batch size ({num_processes} x {self.per_device_eval_batch_size}) must be " + f"evenly divisible by the number of generations per prompt ({self.num_generations}). Given the " + "current global eval batch size, the valid values for the number of generations are: " + f"{possible_values}." + ) + if self.delta is not None and self.use_liger_loss: + raise ValueError("Liger loss does not support two-sided GRPO loss yet.") diff --git a/lib/python3.12/site-packages/trl/trainer/grpo_trainer.py b/lib/python3.12/site-packages/trl/trainer/grpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..558b77e9b2ed51b0e31734986f02c10c8ce15e7f --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/grpo_trainer.py @@ -0,0 +1,1550 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import textwrap +import warnings +from collections import defaultdict, deque +from collections.abc import Sized +from contextlib import nullcontext +from functools import partial +from pathlib import Path +from typing import Any, Callable, Optional, Union + +import datasets +import torch +import torch.utils.data +import transformers +from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed +from datasets import Dataset, IterableDataset +from packaging import version +from torch import nn +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.utils.data import DataLoader, Sampler +from transformers import ( + AutoModelForCausalLM, + AutoModelForSequenceClassification, + AutoTokenizer, + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + Trainer, + TrainerCallback, + is_wandb_available, +) +from transformers.trainer_utils import seed_worker +from transformers.utils import is_datasets_available, is_peft_available, is_rich_available + +from ..data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template +from ..extras.profiling import profiling_context, profiling_decorator +from ..extras.vllm_client import VLLMClient +from ..import_utils import is_liger_kernel_available, is_vllm_available +from ..models import prepare_deepspeed, prepare_fsdp, unwrap_model_for_generation +from ..models.utils import _ForwardRedirection +from .callbacks import SyncRefModelCallback +from .grpo_config import GRPOConfig +from .utils import ( + disable_dropout_in_model, + generate_model_card, + get_comet_experiment_url, + pad, + print_prompt_completions_sample, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import PeftConfig, get_peft_model + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss + +if is_vllm_available(): + from vllm import LLM, SamplingParams + from vllm.sampling_params import GuidedDecodingParams + +if is_wandb_available(): + import wandb + +# What we call a reward function is a callable that takes a list of prompts and completions and returns a list of +# rewards. When it's a string, it's a model ID, so it's loaded as a pretrained model. +RewardFunc = Union[str, PreTrainedModel, Callable[[list, list], list[float]]] + + +class RepeatSampler(Sampler): + """ + Sampler that repeats the indices of a dataset in a structured manner. + + Args: + data_source (`Sized`): + Dataset to sample from. + mini_repeat_count (`int`): + Number of times to repeat each index per batch. + batch_size (`int`, *optional*, defaults to `1`): + Number of unique indices per batch. + repeat_count (`int`, *optional*, defaults to `1`): + Number of times to repeat the full sampling process. + shuffle (`bool`, *optional*, defaults to `True`): + Whether to shuffle the dataset. + seed (`int` or `None`, *optional*, defaults to `None`): + Random seed for reproducibility (only affects this sampler). + + Example: + ```python + >>> sampler = RepeatRandomSampler( + ... ["a", "b", "c", "d", "e", "f", "g"], mini_repeat_count=2, batch_size=3, repeat_count=4 + ... ) + >>> list(sampler) + [4, 4, 3, 3, 0, 0, + 4, 4, 3, 3, 0, 0, + 4, 4, 3, 3, 0, 0, + 4, 4, 3, 3, 0, 0, + 1, 1, 2, 2, 6, 6, + 1, 1, 2, 2, 6, 6, + 1, 1, 2, 2, 6, 6, + 1, 1, 2, 2, 6, 6] + ``` + + ```txt + mini_repeat_count = 3 + - - - + [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, | + 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, | + 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 11, 11, | + repeat_count = 2 + 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, | + 4, 4, 4, 5, 5, 5, 6, 6, 6, 7, 7, 7, | + 8, 8, 8, 9, 9, 9, 10, 10, 10, 11, 11, 11, ...] | + --------- --------- --------- --------- + --------- --------- --------- --------- + --------- --------- --------- --------- + batch_size = 12 + ``` + """ + + def __init__( + self, + data_source: Sized, + mini_repeat_count: int, + batch_size: int = 1, + repeat_count: int = 1, + shuffle: bool = True, + seed: Optional[int] = None, + ): + self.data_source = data_source + self.mini_repeat_count = mini_repeat_count + self.batch_size = batch_size + self.repeat_count = repeat_count + self.num_samples = len(data_source) + self.shuffle = shuffle + self.seed = seed + + if shuffle: + self.generator = torch.Generator() # Create a local random generator + if seed is not None: + self.generator.manual_seed(seed) + + def __iter__(self): + if self.shuffle: + # E.g., [2, 4, 3, 1, 0, 6, 5] (num_samples = 7) + indexes = torch.randperm(self.num_samples, generator=self.generator).tolist() + else: + indexes = list(range(self.num_samples)) + + # [2, 4, 3, 1, 0, 6, 5] + # -> [[2, 4, 3], [1, 0, 6], [5]] (batch_size = 3) + indexes = [indexes[i : i + self.batch_size] for i in range(0, len(indexes), self.batch_size)] + + # [[2, 4, 3], [1, 0, 6], [5]] + # -> [[2, 4, 3], [1, 0, 6]] + indexes = [chunk for chunk in indexes if len(chunk) == self.batch_size] + + for chunk in indexes: + for _ in range(self.repeat_count): + for index in chunk: + for _ in range(self.mini_repeat_count): + yield index + + def __len__(self) -> int: + return self.num_samples * self.mini_repeat_count * self.repeat_count + + +# torch.nanstd doesn't exist, so we define it here +def nanstd(tensor: torch.Tensor) -> torch.Tensor: + """ + Compute the standard deviation of a tensor, ignoring NaNs. This function only supports 1D tensors. + + Args: + tensor (`torch.Tensor`): + Input tensor of shape `(N,)`. + + Returns: + `torch.Tensor`: + Standard deviation of the tensor, ignoring NaNs. + """ + variance = torch.nanmean((tensor - torch.nanmean(tensor, keepdim=True)) ** 2) # Compute variance ignoring NaNs + count = torch.sum(~torch.isnan(tensor)) # Count of non-NaN values + variance *= count / (count - 1) # Bessel's correction + return torch.sqrt(variance) + + +def split_tensor_dict( + tensor_dict: dict[str, Optional[torch.Tensor]], num_chunks: int +) -> list[dict[str, Optional[torch.Tensor]]]: + """ + Splits a dictionary of tensors along the first dimension into `num_chunks` equal parts. + + Example: + ```python + >>> x = torch.arange(12).reshape(6, 2) + >>> y = torch.arange(6).reshape(6, 1) + >>> tensor_dict = {"x": x, "y": y} + >>> split_tensor_dict(tensor_dict, 3) + [ + {"x": tensor([[0, 1], [2, 3]]), "y": tensor([[0], [1]])}, + {"x": tensor([[4, 5], [6, 7]]), "y": tensor([[2], [3]])}, + {"x": tensor([[ 8, 9], [10, 11]]), "y": tensor([[4], [5]])} + ] + ``` + """ + first_tensor = next(tensor for tensor in tensor_dict.values() if tensor is not None) + chunk_size = first_tensor.shape[0] // num_chunks + return [ + { + key: tensor[i * chunk_size : (i + 1) * chunk_size] if tensor is not None else None + for key, tensor in tensor_dict.items() + } + for i in range(num_chunks) + ] + + +def shuffle_tensor_dict(tensor_dict: dict[str, Optional[torch.Tensor]]) -> dict[str, Optional[torch.Tensor]]: + """ + Shuffles a dictionary of tensors along the first dimension in unison. + + Example: + ```python + >>> x = torch.arange(6).reshape(3, 2) + >>> y = torch.arange(3).reshape(3, 1) + >>> tensor_dict = {"x": x, "y": y} + >>> shuffle_tensor_dict(tensor_dict) + {'x': tensor([[2, 3], + [0, 1], + [4, 5]]), + 'y': tensor([[1], + [0], + [2]])} + ``` + """ + first_tensor = next(tensor for tensor in tensor_dict.values() if tensor is not None) + batch_size = first_tensor.shape[0] + permutation = torch.randperm(batch_size) + return {key: tensor[permutation] if tensor is not None else None for key, tensor in tensor_dict.items()} + + +def nanmin(tensor: torch.Tensor) -> torch.Tensor: + """ + Compute the minimum value of a tensor, ignoring NaNs. This function only supports 1D tensors. + + Args: + tensor (`torch.Tensor`): Input tensor of shape `(N,)`. + + Returns: + `torch.Tensor`: Minimum value of the tensor, ignoring NaNs. Returns NaN if all values are NaN. + """ + if torch.isnan(tensor).all(): + return torch.tensor(float("nan"), dtype=tensor.dtype, device=tensor.device) + return torch.min(tensor[~torch.isnan(tensor)]) + + +def nanmax(tensor: torch.Tensor) -> torch.Tensor: + """ + Compute the maximum value of a tensor, ignoring NaNs. This function only supports 1D tensors. + + Args: + tensor (`torch.Tensor`): Input tensor of shape `(N,)`. + + Returns: + `torch.Tensor`: Maximum value of the tensor, ignoring NaNs. Returns NaN if all values are NaN. + """ + if torch.isnan(tensor).all(): + return torch.tensor(float("nan"), dtype=tensor.dtype, device=tensor.device) + return torch.max(tensor[~torch.isnan(tensor)]) + + +def identity(x): + """Do we really need docs for this?""" + return x + + +class GRPOTrainer(Trainer): + """ + Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the + paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language + Models](https://huggingface.co/papers/2402.03300). + + Example: + + ```python + from datasets import load_dataset + from trl import GRPOTrainer + + dataset = load_dataset("trl-lib/tldr", split="train") + + + def reward_func(completions, **kwargs): + # Dummy reward function that rewards completions with more unique letters. + return [float(len(set(completion))) for completion in completions] + + + trainer = GRPOTrainer( + model="Qwen/Qwen2-0.5B-Instruct", + reward_funcs=reward_func, + train_dataset=dataset, + ) + + trainer.train() + ``` + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + reward_funcs (`Union[RewardFunc, list[RewardFunc]]`): + Reward functions to be used for computing the rewards. To compute the rewards, we call all the reward + functions with the prompts and completions and sum the rewards. Can be either: + + - A single reward function, such as: + - A string: The *model ID* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForSequenceClassification.from_pretrained`] with `num_labels=1` and the + keyword arguments in `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported. + - A custom reward function: The function is provided with the prompts and the generated completions, + plus any additional columns in the dataset. It should return a list of rewards. Custom reward + functions can also return None when the reward is not applicable to those samples. This is useful for + multi-task training where different reward functions apply to different types of samples. When a + reward function returns None for a sample, that reward function is excluded from the reward + calculation for that sample. For more details, see [Using a custom reward + function](#using-a-custom-reward-function). + - A list of reward functions, where each item can independently be any of the above types. Mixing different + types within the list (e.g., a string model ID and a custom reward function) is allowed. + args ([`GRPOConfig`], *optional*, defaults to `None`): + Configuration for this trainer. If `None`, a default configuration is used. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. It must include a column `"prompt"`. Any additional columns in the dataset is + ignored. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`): + Processing class used to process the data. The padding side must be set to "left". If `None`, the + processing class is loaded from the model's name with [`~transformers.AutoTokenizer.from_pretrained`]. A + padding token, `processing_class.pad_token`, must be set. If the processing class has not set a padding + token, `processing_class.eos_token` will be used as the default. + reward_processing_classes (`Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]`, *optional*, defaults to `None`): + Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either: + + - A single processing class: Used when `reward_funcs` contains only one reward function. + - A list of processing classes: Must match the order and length of the reward functions in `reward_funcs`. + If set to `None`, or if an element of the list corresponding to a [`~transformers.PreTrainedModel`] is + `None`, the tokenizer for the model is automatically loaded using + [`~transformers.AutoTokenizer.from_pretrained`]. For elements in `reward_funcs` that are custom reward + functions (not [`~transformers.PreTrainedModel`]), the corresponding entries in `reward_processing_classes` + are ignored. + callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + """ + + _tag_names = ["trl", "grpo"] + + def __init__( + self, + model: Union[str, PreTrainedModel], + reward_funcs: Union[RewardFunc, list[RewardFunc]], + args: Optional[GRPOConfig] = None, + train_dataset: Optional[Union[Dataset, IterableDataset]] = None, + eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None, + processing_class: Optional[PreTrainedTokenizerBase] = None, + reward_processing_classes: Optional[Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + peft_config: Optional["PeftConfig"] = None, + ): + # Args + if args is None: + model_name = model if isinstance(model, str) else model.config._name_or_path + model_name = model_name.split("/")[-1] + args = GRPOConfig(f"{model_name}-GRPO") + + # Models + # Trained model + model_init_kwargs = args.model_init_kwargs or {} + if isinstance(model, str): + model_id = model + torch_dtype = model_init_kwargs.get("torch_dtype") + if isinstance(torch_dtype, torch.dtype) or torch_dtype == "auto" or torch_dtype is None: + pass # torch_dtype is already a torch.dtype or "auto" or None + elif isinstance(torch_dtype, str): # it's a str, but not "auto" + torch_dtype = getattr(torch, torch_dtype) + model_init_kwargs["torch_dtype"] = torch_dtype + else: + raise ValueError( + "Invalid `torch_dtype` passed to `GRPOConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}." + ) + # Disable caching if gradient checkpointing is enabled (not supported) + model_init_kwargs["use_cache"] = ( + False if args.gradient_checkpointing else model_init_kwargs.get("use_cache") + ) + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + else: + model_id = model.config._name_or_path + if args.model_init_kwargs is not None: + raise ValueError( + "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. " + "This argument can only be used when the `model` argument is a string." + ) + + if peft_config is not None: + if not is_peft_available(): + raise ImportError("PEFT is required to use `peft_config`. Run `pip install peft`.") + model = get_peft_model(model, peft_config) + + # Enable gradient checkpointing if requested + if args.gradient_checkpointing: + model = self._enable_gradient_checkpointing(model, args) + + # Processing class + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model.config._name_or_path, padding_side="left") + if processing_class.pad_token is None: + processing_class.pad_token = processing_class.eos_token + + # Reward functions + if not isinstance(reward_funcs, list): + reward_funcs = [reward_funcs] + self.reward_func_names = [] + for i, reward_func in enumerate(reward_funcs): + if isinstance(reward_func, str): + reward_funcs[i] = AutoModelForSequenceClassification.from_pretrained( + reward_func, num_labels=1, **model_init_kwargs + ) + if isinstance(reward_funcs[i], nn.Module): # Use Module over PretrainedModel for compat w/ compiled models + self.reward_func_names.append(reward_funcs[i].config._name_or_path.split("/")[-1]) + else: + self.reward_func_names.append(reward_funcs[i].__name__) + self.reward_funcs = reward_funcs + + # Reward weights + if args.reward_weights is not None: + if len(args.reward_weights) != len(reward_funcs): + raise ValueError( + f"Number of reward weights ({len(args.reward_weights)}) must match number of reward " + f"functions ({len(reward_funcs)})" + ) + self.reward_weights = torch.tensor(args.reward_weights, dtype=torch.float32) + else: + self.reward_weights = torch.ones(len(reward_funcs), dtype=torch.float32) + + # Reward processing class + if reward_processing_classes is None: + reward_processing_classes = [None] * len(reward_funcs) + elif not isinstance(reward_processing_classes, list): + reward_processing_classes = [reward_processing_classes] + else: + if len(reward_processing_classes) != len(reward_funcs): + raise ValueError("The number of reward processing classes must match the number of reward functions.") + + for i, (reward_processing_class, reward_func) in enumerate(zip(reward_processing_classes, reward_funcs)): + if isinstance(reward_func, PreTrainedModel): + if reward_processing_class is None: + reward_processing_class = AutoTokenizer.from_pretrained(reward_func.config._name_or_path) + if reward_processing_class.pad_token_id is None: + reward_processing_class.pad_token = reward_processing_class.eos_token + # The reward model computes the reward for the latest non-padded token in the input sequence. + # So it's important to set the pad token ID to the padding token ID of the processing class. + reward_func.config.pad_token_id = reward_processing_class.pad_token_id + reward_processing_classes[i] = reward_processing_class + self.reward_processing_classes = reward_processing_classes + + # Training arguments + self.max_prompt_length = args.max_prompt_length + self.max_completion_length = args.max_completion_length # = |o_i| in the GRPO paper + self.num_generations = args.num_generations # = G in the GRPO paper + self.temperature = args.temperature + self.top_p = args.top_p + self.top_k = args.top_k + self.min_p = args.min_p + self.repetition_penalty = args.repetition_penalty + self.use_vllm = args.use_vllm + self.vllm_mode = args.vllm_mode + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization # only applies to colocation mode + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size # only applies to colocation mode + self.use_liger_loss = args.use_liger_loss + self.loss_type = args.loss_type + self.scale_rewards = args.scale_rewards + self.mask_truncated_completions = args.mask_truncated_completions + + # Datasets + self.shuffle_dataset = args.shuffle_dataset + + if ( + isinstance(train_dataset, IterableDataset) + or isinstance(eval_dataset, IterableDataset) + or ( + isinstance(eval_dataset, dict) and any(isinstance(ds, IterableDataset) for ds in eval_dataset.values()) + ) + ): + # See https://github.com/huggingface/trl/issues/3213 + raise NotImplementedError( + "Iterable datasets are not yet supported in GRPOTrainer. Please use a standard dataset instead." + ) + + # Multi-step + self.num_iterations = args.num_iterations # = 𝜇 in the GRPO paper + self.epsilon_low = args.epsilon + self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon + # Tracks the number of iterations (forward + backward passes), including those within a grad accum cycle + self._step = 0 + # Buffer the batch to reuse generated outputs across multiple updates. For more details, see + # `_get_train_sampler` and `_prepare_inputs`. + self._buffered_inputs = None + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in GRPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys is "prompt". As a result, the trainer issues the warning: + # "Could not estimate the number of tokens of the input, floating-point operations will not be computed." To + # suppress this warning, we set the "estimate_tokens" key in the model's "warnings_issued" dictionary to True. + # This acts as a flag to indicate that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + super().__init__( + model=model, + args=args, + data_collator=identity, # No data collation is needed in GRPO + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + callbacks=callbacks, + optimizers=optimizers, + ) + + # Reference model + self.beta = args.beta + if self.beta == 0.0: + # If beta is 0.0, the reference model is not needed + self.ref_model = None + elif is_peft_model(model): + # If PEFT is used, the reference model is not needed since the adapter can be disabled + # to revert to the initial model. + self.ref_model = None + else: + # For deepspeed, fsdp or non-distributed models, create a reference model from scratch + self.ref_model = AutoModelForCausalLM.from_pretrained(model_id, **model_init_kwargs) + + # Disable dropout in the models + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Liger loss + if self.use_liger_loss: + if not is_liger_kernel_available(): + raise ImportError( + "Liger is required to use `liger_loss` as the GRPO loss. Run `pip install liger-kernel`." + ) + # redirect the model.module forward to the model forward to ensure pre-forward hooks are called + self._forward_redirection = _ForwardRedirection() + + self.liger_grpo_loss = LigerFusedLinearGRPOLoss( + beta=self.beta, + epsilon_low=self.epsilon_low, + epsilon_high=self.epsilon_high, + temperature=self.temperature, + use_ref_model=self.beta != 0.0, + loss_type=self.loss_type, + max_completion_length=self.max_completion_length, + ) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + self.log_completions = args.log_completions + self.wandb_log_unique_prompts = args.wandb_log_unique_prompts + self.num_completions_to_print = args.num_completions_to_print + # maxlen is set to the total number of forward passes per step. This value of `maxlen` ensures we log only the + # final optimization step. + maxlen = self.accelerator.num_processes * args.per_device_train_batch_size * args.steps_per_generation + self._textual_logs = { + "prompt": deque(maxlen=maxlen), + "completion": deque(maxlen=maxlen), + "rewards": defaultdict(lambda: deque(maxlen=maxlen)), + "advantages": deque(maxlen=maxlen), + } + + # Ensure each process receives a unique seed to prevent duplicate completions when generating with + # transformers if num_generations exceeds per_device_train_batch_size. We could skip it if we use vLLM, but + # it's safer to set it in all cases. + set_seed(args.seed, device_specific=True) + + if self.use_vllm: + if not is_vllm_available(): + raise ImportError( + "vLLM is not available and `use_vllm` is set to True. Please install vLLM with " + "`pip install vllm` to use it." + ) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + if args.vllm_server_base_url is not None: + base_url = args.vllm_server_base_url + else: + base_url = f"http://{args.vllm_server_host}:{args.vllm_server_port}" + self.vllm_client = VLLMClient(base_url=base_url, connection_timeout=args.vllm_server_timeout) + self.vllm_client.init_communicator() + + elif self.vllm_mode == "colocate": + # Make sure vllm_tensor_parallel_size group size evenly divides the world size - each group should have + # the same number of ranks + if not self.accelerator.num_processes % self.vllm_tensor_parallel_size == 0: + raise ValueError( + f"vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size " + f"({self.accelerator.num_processes}) evenly." + ) + + if self.vllm_tensor_parallel_size > 1: + # Create subgroups of ranks for TP, each group with `vllm_tensor_parallel_size` ranks. + # For example, if world_size=8 and vllm_tensor_parallel_size=2 → groups: [0,1], [2,3], [4,5], [6,7] + self.tp_group, _ = torch.distributed.new_subgroups_by_enumeration( + [ + list(range(i * self.vllm_tensor_parallel_size, (i + 1) * self.vllm_tensor_parallel_size)) + for i in range(self.accelerator.num_processes // self.vllm_tensor_parallel_size) + ] + ) + + self.llm = LLM( + model=model.name_or_path, + tensor_parallel_size=args.vllm_tensor_parallel_size, + gpu_memory_utilization=self.vllm_gpu_memory_utilization, + max_num_seqs=self.args.per_device_train_batch_size + * self.vllm_tensor_parallel_size + * self.args.gradient_accumulation_steps, + max_model_len=self.max_prompt_length + self.max_completion_length, + distributed_executor_backend="external_launcher", + # Feed identical seed for tp groups to ensure sampling results are the same across workers + seed=self.accelerator.process_index // self.vllm_tensor_parallel_size, + # Latest vLLM v1 memory profiler is misled by the high default value (i.e., 32768) - thinking there's not enough memory + max_num_batched_tokens=4096, + ) + + # vLLM specific sampling arguments + self.guided_decoding_regex = args.vllm_guided_decoding_regex + + self._last_loaded_step = -1 # tag to avoid useless loading during grad accumulation + + # When using vLLM, the main process is responsible for loading the model weights. This can cause process + # desynchronization and seems to lead to DeepSpeed hanging during initialization. To prevent this, we + # synchronize all processes after vLLM has been fully initialized. + self.accelerator.wait_for_everyone() + else: + generation_kwargs = { + "max_new_tokens": self.max_completion_length, + "do_sample": True, + "pad_token_id": processing_class.pad_token_id, + "bos_token_id": processing_class.bos_token_id, + "eos_token_id": processing_class.eos_token_id, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "repetition_penalty": self.repetition_penalty, + "cache_implementation": args.cache_implementation, + } + if args.generation_kwargs is not None: + generation_kwargs.update(args.generation_kwargs) + self.generation_config = GenerationConfig(**generation_kwargs) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags to the model + self.model.add_model_tags(self._tag_names) + + if self.ref_model is not None: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + elif self.is_fsdp_enabled: + self.ref_model = prepare_fsdp(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + if args.sync_ref_model: + self.add_callback(SyncRefModelCallback(ref_model=self.ref_model, accelerator=self.accelerator)) + + for i, reward_func in enumerate(self.reward_funcs): + if isinstance(reward_func, PreTrainedModel): + if self.is_deepspeed_enabled: + self.reward_funcs[i] = prepare_deepspeed(reward_func, self.accelerator) + else: + # set device placement to True to make `prepare_model` move `reward_func` to device when using fsdp + self.reward_funcs[i] = self.accelerator.prepare_model( + reward_func, evaluation_mode=True, device_placement=True + ) + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs. + # In GRPOTrainer, we preprocess data, so using the model's signature columns doesn't work. + # Instead, we set them to the columns expected by the `training_step` method, hence the override. + if self._signature_columns is None: + self._signature_columns = ["prompt"] + + # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy. + # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an + # *generation* batch (i.e., `per_device_batch_size × steps_per_generation`). This allows us to generate completions + # once every steps_per_generation step—rather than once per accumulation step—which is significantly more + # efficient. The only change from the original implementation is multiplying the batch size by + # `steps_per_generation`. Thus, `_prepare_inputs` is called with this *generation* batch, and it handles the + # splitting internally. + # Maintenance note: This method is a copy-paste of the original `Trainer.get_train_dataloader` with only one line + # modification. As a result, some parts of the method aren't relevant to GRPO, but we keep them to stay one line + # apart from the super method, ensuring easier maintenance in the future. + def get_train_dataloader(self): + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + if is_datasets_available() and isinstance(train_dataset, datasets.Dataset): + train_dataset = self._remove_unused_columns(train_dataset, description="training") + else: + data_collator = self._get_collator_with_removed_columns(data_collator, description="training") + + dataloader_params = { + "batch_size": self._train_batch_size * self.args.steps_per_generation, # < this is the change + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + if version.parse(transformers.__version__) >= version.parse("4.52.0"): + # from transformers 4.52.0, the `seed_worker` requires the `num_workers` and `rank` arguments + dataloader_params["worker_init_fn"] = partial( + seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index + ) + else: + dataloader_params["worker_init_fn"] = seed_worker + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + def _get_train_sampler(self, dataset: Optional[Dataset] = None) -> Sampler: + # Returns a sampler that + # 1. ensures each prompt is repeated across multiple processes. This guarantees that identical prompts are + # distributed to different GPUs, allowing rewards to be computed and normalized correctly within each prompt + # group. Using the same seed across processes ensures consistent prompt assignment, preventing discrepancies + # in group formation. + # 2. repeats the batch multiple times to allow reusing generations across multiple updates. Refer to + # _prepare_inputs to see how the generations are stored and reused. + + # In the following figure, the values are the prompt indices. The first row shows the first sampled batch, the + # second row shows the second sampled batch, and so on. + # + # | GPU 0 | GPU 1 | + # + # global_step step <-───> num_generations=2 + # <-───────> per_device_train_batch_size=3 + # grad_accum ▲ ▲ 0 0 0 0 1 1 2 2 <- Generate for the first `steps_per_generation` (prompts 0 to 11); store the completions; use the first slice to compute the loss + # =2 ▼ | 0 1 3 3 4 4 5 5 <- Take the stored generations and use the second slice to compute the loss + # | + # | 1 2 6 6 7 7 8 8 <- Take the stored generations and use the third slice to compute the loss + # steps_per_gen=4 ▼ 1 3 9 9 10 10 11 11 <- Take the stored generations and use the fourth slice to compute the loss + # + # 2 4 12 12 13 13 14 14 <- Generate for the second `steps_per_generation` (prompts 12 to 23); store the completions; use the first slice to compute the loss + # 2 5 15 15 16 16 17 17 <- Take the stored generations and use the second slice to compute the loss + # ... + if dataset is None: + dataset = self.train_dataset + return RepeatSampler( + data_source=dataset, + mini_repeat_count=self.num_generations, + batch_size=self.args.generation_batch_size // self.num_generations, + repeat_count=self.num_iterations * self.args.steps_per_generation, + shuffle=self.shuffle_dataset, + seed=self.args.seed, + ) + + def _get_eval_sampler(self, eval_dataset) -> Sampler: + # See _get_train_sampler for an explanation of the sampler. + return RepeatSampler( + data_source=eval_dataset, + mini_repeat_count=self.num_generations, + seed=self.args.seed, + ) + + def _enable_gradient_checkpointing(self, model: PreTrainedModel, args: GRPOConfig) -> PreTrainedModel: + """Enables gradient checkpointing for the model.""" + # Ensure use_cache is disabled + model.config.use_cache = False + + # Enable gradient checkpointing on the base model for PEFT + if is_peft_model(model): + model.base_model.gradient_checkpointing_enable() + # Enable gradient checkpointing for non-PEFT models + else: + model.gradient_checkpointing_enable() + + gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + use_reentrant = ( + "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"] + ) + + if use_reentrant: + model.enable_input_require_grads() + + return model + + @profiling_decorator + def _get_last_hidden_state(self, unwrapped_model, input_ids, attention_mask, logits_to_keep=None): + if is_peft_model(unwrapped_model): + unwrapped_model = unwrapped_model.base_model.model + last_hidden_state = unwrapped_model.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state + last_hidden_state = last_hidden_state[:, :-1, :] # (B, L-1, H) + if logits_to_keep is not None: + last_hidden_state = last_hidden_state[:, -logits_to_keep:, :] # (B, logits_to_keep, H) + return last_hidden_state + + # Get the per-token log probabilities for the completions for the model and the reference model + @profiling_decorator + def _get_per_token_logps(self, model, input_ids, attention_mask, logits_to_keep, batch_size=None) -> torch.Tensor: + batch_size = batch_size or input_ids.size(0) # Chunk inputs into smaller batches to reduce memory peak + all_logps = [] + for i in range(0, input_ids.size(0), batch_size): + input_ids_batch = input_ids[i : i + batch_size] + attention_mask_batch = attention_mask[i : i + batch_size] + + # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded + logits = model( + input_ids=input_ids_batch, attention_mask=attention_mask_batch, logits_to_keep=logits_to_keep + 1 + ).logits + logits = logits[:, :-1, :] # (B, L-1, V), exclude the last logit: it corresponds to the next token pred + input_ids_batch = input_ids_batch[:, -logits_to_keep:] + # Divide logits by sampling temperature. + # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details + logits = logits / self.temperature + logps = selective_log_softmax(logits, input_ids_batch) # compute logprobs for the input tokens + all_logps.append(logps) + return torch.cat(all_logps, dim=0) + + def _sync_fsdp_params_to_vllm(self, module: nn.Module, prefix: str = "", visited=None): + """Memory-efficient post-order traversal of FSDP modules to extract full parameters and sync with vLLM.""" + if visited is None: + visited = set() + + for child_name, child_module in module.named_children(): + child_prefix = f"{prefix}.{child_name}" if prefix else child_name + self._sync_fsdp_params_to_vllm( + child_module, prefix=child_prefix, visited=visited + ) # recurse into the child + + if isinstance(module, FSDP): + with FSDP.summon_full_params(module, recurse=False, writeback=False): + for param_name, param in module.named_parameters(): + full_name = f"{prefix}.{param_name}" if prefix else param_name + for extra in ("_fsdp_wrapped_module.", "_checkpoint_wrapped_module."): + full_name = full_name.replace(extra, "") + + if full_name in visited: + continue # skip FSDP subtrees already traversed + visited.add(full_name) + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(full_name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(full_name, param.data)]) + + @profiling_decorator + def _move_model_to_vllm(self): + # For DeepSpeed ZeRO-3 and FSDP, we need to gather all parameters before operations + deepspeed_plugin = self.accelerator.state.deepspeed_plugin + zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3 + if zero_stage_3: + import deepspeed + + gather_if_zero3 = deepspeed.zero.GatheredParameters + else: + gather_if_zero3 = nullcontext + + if is_peft_model(self.model): + # With PEFT and FSDP/DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as + # merging adapters in a sharded manner is not supported. + # TODO: does this work with FSDP? + with gather_if_zero3(list(self.model.parameters())): + self.model.merge_adapter() + + # Update vLLM weights while parameters are gathered + if self.is_fsdp_enabled: # note if using FSDP, gather_if_zero3 is nullcontext + # Update vLLM weights while parameters are gathered + # For PEFT with FSDP we need to use the memory efficient post-order traversal + self._sync_fsdp_params_to_vllm(self.model) + else: + # DeepSpeed ZeRO-3 with PEFT + for name, param in self.model.named_parameters(): + # When using PEFT, we need to recover the original parameter name and discard some parameters + name = name.removeprefix("base_model.model.").replace(".base_layer", "") + if self.model.prefix in name: + continue + # When module to save, remove its prefix and discard the original module + if "original_module" in name: + continue + name = name.replace("modules_to_save.default.", "") + + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + # Unmerge adapters while parameters are still gathered + self.model.unmerge_adapter() + # Parameters will automatically be repartitioned when exiting the context + else: + # For non-PEFT models, simply gather (if needed) and update each parameter individually. + if self.is_fsdp_enabled: + self._sync_fsdp_params_to_vllm(self.model) # use memory-efficient post-order traversal for FSDP + else: + for name, param in self.model.named_parameters(): + with gather_if_zero3([param]): + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.update_named_param(name, param.data) + elif self.vllm_mode == "colocate": + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights([(name, param.data)]) + + # Reset cache on vLLM + if self.vllm_mode == "server" and self.accelerator.is_main_process: + self.vllm_client.reset_prefix_cache() + elif self.vllm_mode == "colocate": + self.llm.reset_prefix_cache() + + @profiling_decorator + def _prepare_inputs( + self, generation_batch: dict[str, Union[torch.Tensor, Any]] + ) -> dict[str, Union[torch.Tensor, Any]]: + # Prepares inputs for model training/evaluation by managing completion generation and batch handling. + # During training: + # - Receives the local generation batch (Per-GPU batch size × steps per generation) + # from the modified training dataloader instead of the standard local batch + # - Generates completions once for the entire generation batch and splits it into batches of size + # `per_device_train_batch_size` + # - Buffers these completions and returns the appropriate slice for the current accumulation step + # - Optimizes by regenerating completions only periodically (every steps_per_generation * num_iterations) + # During evaluation: + # - The input is treated as a standard local batch (no accumulation, no multiple iterations) + # - Completions are generated for each batch without buffering or reuse + # Returns a single local batch in both cases. + + mode = "train" if self.model.training else "eval" + if mode == "train": + generate_every = self.args.steps_per_generation * self.num_iterations + if self._step % generate_every == 0 or self._buffered_inputs is None: + # self._buffered_inputs=None can occur when resuming from a checkpoint + generation_batch = self._generate_and_score_completions(generation_batch) + generation_batch = shuffle_tensor_dict(generation_batch) + self._buffered_inputs = split_tensor_dict(generation_batch, self.args.steps_per_generation) + inputs = self._buffered_inputs[self._step % self.args.steps_per_generation] + self._step += 1 + else: + # In evaluation, there is neither batch grouping for generation, nor multiple iterations, hence + # local generation batch == local eval batch + inputs = self._generate_and_score_completions(generation_batch) + return inputs + + @profiling_decorator + def _calculate_rewards(self, inputs, prompts, completions, completion_ids_list): + device = self.accelerator.device + rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device) + + # Repeat all input columns (but "prompt", "completion", and "completion_ids") to match the num of generations + keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]] + reward_kwargs = {key: [example[key] for example in inputs] for key in keys} + + for i, (reward_func, reward_processing_class, reward_func_name) in enumerate( + zip(self.reward_funcs, self.reward_processing_classes, self.reward_func_names) + ): + with profiling_context(self, reward_func_name): + if isinstance(reward_func, nn.Module): # Module (no PretrainedModel) for compat with compiled models + if is_conversational(inputs[0]): + messages = [{"messages": p + c} for p, c in zip(prompts, completions)] + texts = [apply_chat_template(x, reward_processing_class)["text"] for x in messages] + else: + texts = [p + c for p, c in zip(prompts, completions)] + reward_inputs = reward_processing_class( + text=texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False + ) + reward_inputs = super()._prepare_inputs(reward_inputs) + with torch.inference_mode(): + rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0] # Shape (B*G,) + else: + output_reward_func = reward_func( + prompts=prompts, completions=completions, completion_ids=completion_ids_list, **reward_kwargs + ) + # Convert None values to NaN + output_reward_func = [reward if reward is not None else torch.nan for reward in output_reward_func] + + rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device) + + # If all reward functions return None for a given row, issue a detailed warning + if torch.isnan(rewards_per_func).all(dim=1).any(): + nan_row_idx = torch.isnan(rewards_per_func).all(dim=1).nonzero(as_tuple=True)[0][0] + row_reward_kwargs = {key: value[nan_row_idx] for key, value in reward_kwargs.items()} + row_reward_kwargs["prompt"] = prompts[nan_row_idx] + row_reward_kwargs["completion"] = completions[nan_row_idx] + warnings.warn( + f"All reward functions returned None for the following kwargs: {row_reward_kwargs}. " + "Please ensure that at least one reward function returns a valid reward." + ) + + # Gather the reward per function: this part is crucial, because the rewards are normalized per group and the + # completions may be distributed across processes + rewards_per_func = gather(rewards_per_func) + return rewards_per_func + + def _generate_and_score_completions( + self, inputs: list[dict[str, Union[torch.Tensor, Any]]] + ) -> dict[str, Union[torch.Tensor, Any]]: + device = self.accelerator.device + mode = "train" if self.model.training else "eval" + + prompts = [x["prompt"] for x in inputs] + prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs] + prompt_inputs = self.processing_class( + text=prompts_text, return_tensors="pt", padding=True, padding_side="left", add_special_tokens=False + ) + prompt_inputs = super()._prepare_inputs(prompt_inputs) + prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"] + + if self.max_prompt_length is not None: + prompt_ids = prompt_ids[:, -self.max_prompt_length :] + prompt_mask = prompt_mask[:, -self.max_prompt_length :] + prompts_text = self.processing_class.batch_decode( + prompt_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False + ) + + # Generate completions using either vLLM or regular generation + if self.use_vllm: + # First, update the vLLM weights if needed + if self.state.global_step != self._last_loaded_step: + self._move_model_to_vllm() + self._last_loaded_step = self.state.global_step + + # Generate completions using vLLM: gather all prompts and use them in a single call in the main process + if self.vllm_mode == "server": + all_prompts_text = gather_object(prompts_text) + if self.accelerator.is_main_process: + # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate + # num_generations outputs for each one. This is faster than generating outputs for each duplicate + # prompt individually. + ordered_set_of_prompts = all_prompts_text[:: self.num_generations] + with profiling_context(self, "vLLM.generate"): + completion_ids = self.vllm_client.generate( + prompts=ordered_set_of_prompts, + n=self.num_generations, + repetition_penalty=self.repetition_penalty, + temperature=self.temperature, + top_p=self.top_p, + top_k=-1 if self.top_k is None else self.top_k, + min_p=0.0 if self.min_p is None else self.min_p, + max_tokens=self.max_completion_length, + guided_decoding_regex=self.guided_decoding_regex, + generation_kwargs=self.args.generation_kwargs, + ) + else: + completion_ids = [None] * len(all_prompts_text) + # Broadcast the completions from the main process to all processes, ensuring each process receives its + # corresponding slice. + completion_ids = broadcast_object_list(completion_ids, from_process=0) + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + completion_ids = completion_ids[process_slice] + + # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts + elif self.vllm_mode == "colocate": + if self.guided_decoding_regex: + guided_decoding = GuidedDecodingParams(backend="outlines", regex=self.guided_decoding_regex) + else: + guided_decoding = None + + generation_kwargs = { + "n": 1, # vLLM on each GPU generates only 1 in colocate mode + "repetition_penalty": self.repetition_penalty, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": -1 if self.top_k is None else self.top_k, + "min_p": 0.0 if self.min_p is None else self.min_p, + "max_tokens": self.max_completion_length, + "guided_decoding": guided_decoding, + } + generation_kwargs.update(self.args.generation_kwargs) + sampling_params = SamplingParams(**generation_kwargs) + + if self.vllm_tensor_parallel_size > 1: + # Gather prompts from all ranks in the TP group and flatten. + # Each rank starts with its own prompts; after gathering, all ranks see the full group set. + orig_size = len(prompts_text) + gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)] + torch.distributed.all_gather_object(gathered_prompts, prompts_text, group=self.tp_group) + all_prompts_text = [p for sublist in gathered_prompts for p in sublist] + else: + all_prompts_text = prompts_text + + with profiling_context(self, "vLLM.generate"): + all_outputs = self.llm.generate(all_prompts_text, sampling_params=sampling_params, use_tqdm=False) + + completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs] + + if self.vllm_tensor_parallel_size > 1: + # Slice completions for this rank within its TP group. + # Each rank generates all outputs — we keep only our share. + local_rank_in_group = torch.distributed.get_rank(group=self.tp_group) + tp_slice = slice(local_rank_in_group * orig_size, (local_rank_in_group + 1) * orig_size) + completion_ids = completion_ids[tp_slice] + + # Pad the completions, and concatenate them with the prompts + completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids] + completion_ids = pad(completion_ids, padding_value=self.processing_class.pad_token_id) + prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1) + else: + # Regular generation path + with unwrap_model_for_generation( + self.model_wrapped, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + with ( + FSDP.summon_full_params(self.model_wrapped, recurse=False) + if self.is_fsdp_enabled + else nullcontext() + ): + prompt_completion_ids = unwrapped_model.generate( + prompt_ids, attention_mask=prompt_mask, generation_config=self.generation_config + ) + + # Compute prompt length and extract completion ids + prompt_length = prompt_ids.size(1) + prompt_ids = prompt_completion_ids[:, :prompt_length] + completion_ids = prompt_completion_ids[:, prompt_length:] + + # Mask everything after the first EOS token + is_eos = completion_ids == self.processing_class.eos_token_id + eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device) + eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)] + sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1) + completion_mask = (sequence_indices <= eos_idx.unsqueeze(1)).int() + + # Convert tensor to a list of lists of token IDs. This will be passed to the reward function, avoiding the need + # to re-tokenize completions if the reward is computed from tokens. + completion_ids_list = [ + [id.item() for id, m in zip(row, mask_row) if m] for row, mask_row in zip(completion_ids, completion_mask) + ] + + # Sum along sequence dimension (dim=1) to get completion length per sequence, used for logging + completion_lengths = completion_mask.sum(1) + + # If mask_truncated_completions is enabled, zero out truncated completions in completion_mask + if self.mask_truncated_completions: + truncated_completions = ~is_eos.any(dim=1) + completion_mask = completion_mask * (~truncated_completions).unsqueeze(1).int() + + # Concatenate prompt_mask with completion_mask for logit computation + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) # (B, P+C) + + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size + + with torch.no_grad(): + # When using num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps + # old_per_token_logps == per_token_logps, so we can skip it's computation here, and use + # per_token_logps.detach() instead. + if self.num_iterations > 1 or self.args.steps_per_generation > self.args.gradient_accumulation_steps: + old_per_token_logps = self._get_per_token_logps( + self.model, prompt_completion_ids, attention_mask, logits_to_keep, batch_size + ) + else: + old_per_token_logps = None + + # Compute the per-token log probabilities for the reference model + if self.beta != 0.0: + if self.ref_model is not None: + ref_per_token_logps = self._get_per_token_logps( + self.ref_model, prompt_completion_ids, attention_mask, logits_to_keep + ) + else: + with self.accelerator.unwrap_model(self.model).disable_adapter(): + ref_per_token_logps = self._get_per_token_logps( + self.model, prompt_completion_ids, attention_mask, logits_to_keep + ) + else: + ref_per_token_logps = None + + # Decode the generated completions + completions_text = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational(inputs[0]): + completions = [] + for prompt, completion in zip(prompts, completions_text): + bootstrap = prompt.pop()["content"] if prompt[-1]["role"] == "assistant" else "" + completions.append([{"role": "assistant", "content": bootstrap + completion}]) + else: + completions = completions_text + + # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is + # important because rewards will be normalized per group, and completions are distributed. We will later slice + # rewards_per_func to extract each process's subset. + rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list) + + # Apply weights to each reward function's output and sum + rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1) + + # Compute grouped-wise rewards + mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1) + std_grouped_rewards = rewards.view(-1, self.num_generations).std(dim=1) + is_std_zero = torch.isclose(std_grouped_rewards, torch.zeros_like(std_grouped_rewards)) + + # Normalize the rewards to compute the advantages + mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0) + std_grouped_rewards = std_grouped_rewards.repeat_interleave(self.num_generations, dim=0) + advantages = rewards - mean_grouped_rewards + if self.scale_rewards: + advantages = advantages / (std_grouped_rewards + 1e-4) + + # Slice to keep only the local part of the data + process_slice = slice( + self.accelerator.process_index * len(prompts), + (self.accelerator.process_index + 1) * len(prompts), + ) + all_process_advantages = advantages.clone() # keep the aggregated advantages for logging + advantages = advantages[process_slice] + + # Log the metrics + if mode == "train": + self.state.num_input_tokens_seen += self.accelerator.gather(attention_mask.sum()).sum().item() + self._metrics[mode]["num_tokens"] = [self.state.num_input_tokens_seen] + + # Log completion lengths, mean, min, max + agg_completion_lengths = self.accelerator.gather(completion_lengths) + self._metrics[mode]["completions/mean_length"].append(agg_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_length"].append(agg_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_length"].append(agg_completion_lengths.float().max().item()) + + # Identify sequences that terminated with EOS and log their lengths + agg_terminated_with_eos = self.accelerator.gather(is_eos.any(dim=1)) + term_completion_lengths = agg_completion_lengths[agg_terminated_with_eos] + clipped_completions_ratio = 1 - len(term_completion_lengths) / len(agg_completion_lengths) + self._metrics[mode]["completions/clipped_ratio"].append(clipped_completions_ratio) + if len(term_completion_lengths) == 0: # edge case where no terminated sequences are found + term_completion_lengths = torch.zeros(1, device=device) + self._metrics[mode]["completions/mean_terminated_length"].append(term_completion_lengths.float().mean().item()) + self._metrics[mode]["completions/min_terminated_length"].append(term_completion_lengths.float().min().item()) + self._metrics[mode]["completions/max_terminated_length"].append(term_completion_lengths.float().max().item()) + + # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values) + for i, reward_func_name in enumerate(self.reward_func_names): + mean_rewards = torch.nanmean(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/mean"].append(mean_rewards) + std_rewards = nanstd(rewards_per_func[:, i]).item() + self._metrics[mode][f"rewards/{reward_func_name}/std"].append(std_rewards) + self._metrics[mode]["reward"].append(mean_grouped_rewards.mean().item()) + self._metrics[mode]["reward_std"].append(std_grouped_rewards.mean().item()) + self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item()) + + # Log prompt and completion texts + self._textual_logs["prompt"].extend(gather_object(prompts_text)) + self._textual_logs["completion"].extend(gather_object(completions_text)) + for i, name in enumerate(self.reward_func_names): + self._textual_logs["rewards"][name].extend(rewards_per_func[:, i].tolist()) + self._textual_logs["advantages"].extend(all_process_advantages.tolist()) + + return { + "prompt_ids": prompt_ids, + "prompt_mask": prompt_mask, + "completion_ids": completion_ids, + "completion_mask": completion_mask, + "advantages": advantages, + "old_per_token_logps": old_per_token_logps, + "ref_per_token_logps": ref_per_token_logps, + } + + def compute_liger_loss(self, unwrapped_model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + # get the last hidden state of the model + last_hidden_state = self._get_last_hidden_state(unwrapped_model, input_ids, attention_mask, logits_to_keep) + + # compute loss and metrics using liger grpo loss + loss, metrics = self.liger_grpo_loss( + _input=last_hidden_state, + lin_weight=unwrapped_model.lm_head.weight, + selected_token_ids=completion_ids, + attention_mask=completion_mask, + advantages=inputs["advantages"], + bias=unwrapped_model.lm_head.bias, + old_per_token_logps=inputs["old_per_token_logps"], + ref_per_token_logps=inputs["ref_per_token_logps"], + ) + # Extract metrics from the liger_grpo_loss output + # KL divergence is the first metric when beta is non-zero + mean_kl = metrics[0] if self.beta != 0.0 else None + clip_ratio = metrics[-1] + + mode = "train" if self.model.training else "eval" + if self.beta != 0.0: + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).mean().item()) + self._metrics[mode]["clip_ratio"].append(self.accelerator.gather(clip_ratio).mean().item()) + return loss + + @profiling_decorator + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + if return_outputs: + raise ValueError("The GRPOTrainer does not support returning outputs") + if self.use_liger_loss: + # Compute the loss using the liger grpo loss + unwrapped_model = self.accelerator.unwrap_model(model) + return self._forward_redirection(model, unwrapped_model, self.compute_liger_loss, unwrapped_model, inputs) + else: + return self._compute_loss(model, inputs) + + def _compute_loss(self, model, inputs): + # Compute the per-token log probabilities for the model + prompt_ids, prompt_mask = inputs["prompt_ids"], inputs["prompt_mask"] + completion_ids, completion_mask = inputs["completion_ids"], inputs["completion_mask"] + input_ids = torch.cat([prompt_ids, completion_ids], dim=1) + attention_mask = torch.cat([prompt_mask, completion_mask], dim=1) + logits_to_keep = completion_ids.size(1) # we only need to compute the logits for the completion tokens + + per_token_logps = self._get_per_token_logps(model, input_ids, attention_mask, logits_to_keep) + + # Compute the KL divergence between the model and the reference model + if self.beta != 0.0: + ref_per_token_logps = inputs["ref_per_token_logps"] + per_token_kl = ( + torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1 + ) + + # Compute the loss + advantages = inputs["advantages"] + # When using num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps + # old_per_token_logps == per_token_logps, so we can skip it's computation + # (see _generate_and_score_completions) and use per_token_logps.detach() instead. + old_per_token_logps = ( + per_token_logps.detach() if inputs["old_per_token_logps"] is None else inputs["old_per_token_logps"] + ) + coef_1 = torch.exp(per_token_logps - old_per_token_logps) + coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) + + # Two-sided clipping + if self.args.delta is not None: + coef_1 = torch.clamp(coef_1, max=self.args.delta) + + per_token_loss1 = coef_1 * advantages.unsqueeze(1) + per_token_loss2 = coef_2 * advantages.unsqueeze(1) + per_token_loss = -torch.min(per_token_loss1, per_token_loss2) + if self.beta != 0.0: + per_token_loss = per_token_loss + self.beta * per_token_kl + + if self.loss_type == "grpo": + loss = ((per_token_loss * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)).mean() + elif self.loss_type == "bnpo": + loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + elif self.loss_type == "dr_grpo": + loss = (per_token_loss * completion_mask).sum() / (per_token_loss.size(0) * self.max_completion_length) + else: + raise ValueError(f"Unknown loss type: {self.loss_type}") + + # Log the metrics + mode = "train" if self.model.training else "eval" + + if self.beta != 0.0: + mean_kl = (per_token_kl * completion_mask).sum() / completion_mask.sum() + self._metrics[mode]["kl"].append(self.accelerator.gather(mean_kl).nanmean().item()) + + # Compute the clipped probability ratios + is_low_clipped = (coef_1 < 1 - self.epsilon_low) & (advantages.unsqueeze(1) < 0) + is_high_clipped = (coef_1 > 1 + self.epsilon_high) & (advantages.unsqueeze(1) > 0) + is_region_clipped = is_low_clipped | is_high_clipped + + low_clip = (is_low_clipped * completion_mask).sum() / completion_mask.sum() + high_clip = (is_high_clipped * completion_mask).sum() / completion_mask.sum() + clip_ratio = (is_region_clipped * completion_mask).sum() / completion_mask.sum() + + gathered_low_clip = self.accelerator.gather(low_clip) + self._metrics[mode]["clip_ratio/low_mean"].append(gathered_low_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/low_min"].append(nanmin(gathered_low_clip).item()) + gathered_high_clip = self.accelerator.gather(high_clip) + self._metrics[mode]["clip_ratio/high_mean"].append(gathered_high_clip.nanmean().item()) + self._metrics[mode]["clip_ratio/high_max"].append(nanmax(gathered_high_clip).item()) + gathered_clip_ratio = self.accelerator.gather(clip_ratio) + self._metrics[mode]["clip_ratio/region_mean"].append(gathered_clip_ratio.nanmean().item()) + return loss + + def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys: Optional[list[str]] = None): + inputs = self._prepare_inputs(inputs) + with torch.no_grad(): + with self.compute_loss_context_manager(): + loss = self.compute_loss(model, inputs) + loss = loss.mean().detach() + return loss, None, None + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs = {**logs, **metrics} + super().log(logs, start_time) + self._metrics[mode].clear() + + if self.accelerator.is_main_process and self.log_completions: + if is_rich_available(): + print_prompt_completions_sample( + self._textual_logs["prompt"], + self._textual_logs["completion"], + self._textual_logs["rewards"], + self._textual_logs["advantages"], + self.state.global_step, + self.num_completions_to_print, + ) + + if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None: + import pandas as pd + + table = { + "step": [str(self.state.global_step)] * len(self._textual_logs["prompt"]), + "prompt": self._textual_logs["prompt"], + "completion": self._textual_logs["completion"], + **self._textual_logs["rewards"], + "advantage": self._textual_logs["advantages"], + } + df = pd.DataFrame(table) + if self.wandb_log_unique_prompts: + df = df.drop_duplicates(subset=["prompt"]) + wandb.log({"completions": wandb.Table(dataframe=df)}) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent( + """\ + @article{zhihong2024deepseekmath, + title = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}}, + author = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and Mingchuan Zhang and Y. K. Li and Y. Wu and Daya Guo}, + year = 2024, + eprint = {arXiv:2402.03300}, + } + """ + ) + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="GRPO", + trainer_citation=citation, + paper_title="DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models", + paper_id="2402.03300", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/iterative_sft_config.py b/lib/python3.12/site-packages/trl/trainer/iterative_sft_config.py new file mode 100644 index 0000000000000000000000000000000000000000..9695c7e063ad7740402cfa779353ad42f6d382bc --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/iterative_sft_config.py @@ -0,0 +1,101 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Any, Optional + +from transformers import TrainingArguments + + +@dataclass +class IterativeSFTConfig(TrainingArguments): + r""" + Configuration class for the [`IterativeSFTTrainer`]. + + This class includes only the parameters that are specific to Iterative SFT training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model + + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`IterativeSFTTrainer`] is provided as a string. + + > Parameters that control the data preprocessing + + max_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + The truncation mode to use, either `"keep_end"` or `"keep_start"`. + optimize_device_cache (`bool`, *optional*, defaults to `False`): + Whether to optimize accelerator cache for slightly more memory-efficient training. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + # Parameters that control the model + model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of " + "the `IterativeSFTTrainer` is provided as a string." + }, + ) + + # Parameters that control the data preprocessing + max_length: Optional[int] = field( + default=None, + metadata={ + "help": "Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated." + }, + ) + truncation_mode: str = field( + default="keep_end", + metadata={"help": "The truncation mode to use, either 'keep_end' or 'keep_start'."}, + ) + optimize_device_cache: bool = field( + default=False, + metadata={"help": "Whether to optimize accelerator cache for slightly more memory-efficient training."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + if self.truncation_mode not in ["keep_end", "keep_start"]: + raise ValueError(f"truncation_mode must be either 'keep_end' or 'keep_start', got {self.truncation_mode}") diff --git a/lib/python3.12/site-packages/trl/trainer/iterative_sft_trainer.py b/lib/python3.12/site-packages/trl/trainer/iterative_sft_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..aa5c734b0262d0ba878e710d4b67cc778f8a79ce --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/iterative_sft_trainer.py @@ -0,0 +1,524 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import warnings +from pathlib import Path +from typing import Callable, Optional, Union + +import torch +from datasets import Dataset +from torch.utils.data import DataLoader +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BaseImageProcessor, + DataCollator, + DataCollatorForLanguageModeling, + DataCollatorForSeq2Seq, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + TrainingArguments, + is_wandb_available, +) +from transformers.trainer_utils import EvalLoopOutput +from transformers.utils import is_peft_available + +from ..core import PPODecorators +from .iterative_sft_config import IterativeSFTConfig +from .utils import generate_model_card, get_comet_experiment_url + + +if is_peft_available(): + from peft import PeftModel + + +if is_wandb_available(): + import wandb + + +class IterativeSFTTrainer(Trainer): + """ + The IterativeSFTTrainer can be used to finetune models with methods that requires some steps between optimization. + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + args ([`IterativeSFTConfig`], *optional*, defaults to `None`): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator (`DataCollator`, *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to [`~transformers.default_data_collator`] if no `processing_class` is provided, an instance + of [`~transformers.DataCollatorWithPadding`] otherwise if the processing_class is a feature extractor or + tokenizer. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoTokenizer.from_pretrained`]. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + max_length (`int`, *optional*, deprecated): + Maximum length of the tokenized sequence. Use `args.max_length` instead. + truncation_mode (`str`, *optional*, deprecated): + The truncation mode to use. Use `args.truncation_mode` instead. + optimize_device_cache (`bool`, *optional*, deprecated): + Whether to optimize accelerator cache. Use `args.optimize_device_cache` instead. + """ + + _tag_names = ["trl", "iterative-sft"] + + def __init__( + self, + model: Union[str, PreTrainedModel], + args: Optional[Union[IterativeSFTConfig, TrainingArguments]] = None, + data_collator: Optional[DataCollator] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = ( + None, + None, + ), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + # Deprecated parameters + max_length: Optional[int] = None, + truncation_mode: Optional[str] = None, + optimize_device_cache: Optional[bool] = None, + ): + # Handle deprecated parameters + deprecated_params = {} + if max_length is not None: + deprecated_params["max_length"] = max_length + warnings.warn( + "The `max_length` parameter is deprecated and will be removed in version 0.20. " + "Pass it through the `args` parameter using `IterativeSFTConfig(max_length=...)` instead.", + DeprecationWarning, + ) + if truncation_mode is not None: + deprecated_params["truncation_mode"] = truncation_mode + warnings.warn( + "The `truncation_mode` parameter is deprecated and will be removed in version 0.20. " + "Pass it through the `args` parameter using `IterativeSFTConfig(truncation_mode=...)` instead.", + DeprecationWarning, + ) + if optimize_device_cache is not None: + deprecated_params["optimize_device_cache"] = optimize_device_cache + warnings.warn( + "The `optimize_device_cache` parameter is deprecated and will be removed in version 0.20 " + "Pass it through the `args` parameter using `IterativeSFTConfig(optimize_device_cache=...)` instead.", + DeprecationWarning, + ) + + # Args + model_id = model if isinstance(model, str) else model.config._name_or_path + if args is None: + model_name = model_id.split("/")[-1] + args = IterativeSFTConfig(f"{model_name}-IterativeSFT") + elif isinstance(args, TrainingArguments) and not isinstance(args, IterativeSFTConfig): + dict_args = args.to_dict() + dict_args["hub_token"] = args.hub_token # to_dict hides the hub_token + dict_args.pop("push_to_hub_token") + args = IterativeSFTConfig(**dict_args) + + # Update args with deprecated parameters if provided + if deprecated_params: + for key, value in deprecated_params.items(): + setattr(args, key, value) + + # Handle the tokenizer + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model_id) + + # Model + if args.model_init_kwargs is not None and not isinstance(model, str): + warnings.warn( + "You passed model_init_kwargs to the `IterativeSFTConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + if isinstance(model, str): + model = self._create_model_from_path(model, args) + + # PEFT configuration and model wrapping + if is_peft_available() and isinstance(model, PeftModel): + self.is_peft_model = True + else: + self.is_peft_model = False + + self.processing_class = processing_class + self.is_encoder_decoder = getattr(model.config, "is_encoder_decoder", False) + + if data_collator is None: + if self.is_encoder_decoder: + self.data_collator = DataCollatorForSeq2Seq( + processing_class, label_pad_token_id=-100, pad_to_multiple_of=8 + ) + else: + self.data_collator = DataCollatorForLanguageModeling(self.processing_class, mlm=False) + else: + self.data_collator = data_collator + + self.max_length = args.max_length + self.truncation_mode = args.truncation_mode + self.optimize_device_cache = args.optimize_device_cache + + super().__init__( + model=model, + args=args, + data_collator=self.data_collator, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + self.create_optimizer_and_scheduler(self.args.max_steps) + + # prepare model, optimizer and lr_scheduler + self.model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( + self.model, self.optimizer, self.lr_scheduler + ) + + self.processing_class.truncation_side = "left" if self.truncation_mode == "keep_end" else "right" + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + PPODecorators.optimize_device_cache = self.optimize_device_cache + + def _create_model_from_path(self, model_path: str, args: IterativeSFTConfig) -> PreTrainedModel: + """Creates a model from a path or model identifier.""" + model_init_kwargs = args.model_init_kwargs or {} + return AutoModelForCausalLM.from_pretrained(model_path, **model_init_kwargs) + + def prepare_model_inputs(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, labels: torch.Tensor): + if attention_mask is None: + attention_mask = [torch.ones_like(ids) for ids in input_ids] + + if self.is_encoder_decoder: + input_data = self.data_collator( + [ + {"input_ids": ids, "attention_mask": att, "labels": lab} + for ids, att, lab in zip(input_ids, attention_mask, labels) + ] + ).to(self.model.device) + + input_data.pop("decoder_input_ids", None) # This is directly computed inside the model + + input_data["labels"][input_data["labels"] == self.processing_class.pad_token_id] = -100 + + else: + input_data = self.data_collator( + [{"input_ids": ids, "attention_mask": att} for ids, att in zip(input_ids, attention_mask)] + ).to(self.model.device) + + # truncate in case the user has provided input_ids, attention_mask and labels + if self.max_length is not None: + if self.truncation_mode == "keep_start": + input_data = {k: v[: self.max_length] for k, v in input_data.items()} + elif self.truncation_mode == "keep_end": + input_data = {k: v[-self.max_length :] for k, v in input_data.items()} + else: + raise ValueError(f"Unknown truncation mode: {self.truncation_mode}") + + return input_data + + @staticmethod + def _step_safety_checker( + input_ids: list[torch.LongTensor], + attention_mask: list[torch.LongTensor], + labels: list[torch.LongTensor], + texts: list[str], + texts_labels: list[str], + ): + """ + Check if the input data is valid for training. + + Args: + input_ids (list[`torch.LongTensor`]): + List of tensors containing the input_ids + attention_mask (list[`torch.LongTensor`]): + List of tensors containing the attention_mask + labels (list[`torch.FloatTensor`]): + List of tensors containing the labels + texts (list[`str`]): + List of string containing the text input. + texts_labels (list[`str`]): + List of string containing the text labels. + + Returns: + `tuple`: The input data. + """ + if texts is None: + if attention_mask is None: + for name, tensor_list in zip(["input_ids", "labels"], [input_ids, labels]): + if not isinstance(tensor_list, list): + raise ValueError(f"{name} must be a list of tensors - got {type(tensor_list)}") + if not isinstance(tensor_list[0], torch.Tensor): + raise ValueError(f"Elements in {name} must be tensors - got {type(tensor_list[0])}") + else: + for name, tensor_list in zip( + ["input_ids", "attention_mask", "labels"], [input_ids, attention_mask, labels] + ): + if not isinstance(tensor_list, list): + raise ValueError(f"{name} must be a list of tensors - got {type(tensor_list)}") + if not isinstance(tensor_list[0], torch.Tensor): + raise ValueError(f"Elements in {name} must be tensors - got {type(tensor_list[0])}") + else: + if not isinstance(texts, list): + raise ValueError(f"'text' must be a list of strings - got {type(texts)}") + if not isinstance(texts[0], str): + raise ValueError(f"Elements in 'text' must be strings - got {type(texts[0])}") + if texts_labels is not None: + if not isinstance(texts_labels, list): + raise ValueError(f"'text_labels' must be a list of strings - got {type(texts_labels)}") + if not isinstance(texts_labels[0], str): + raise ValueError(f"Elements in 'text_labels' must be strings - got {type(texts_labels[0])}") + + return input_ids, attention_mask, labels, texts, texts_labels + + @PPODecorators.empty_device_cache() + def step( + self, + input_ids: Optional[list[torch.LongTensor]] = None, + attention_mask: Optional[list[torch.LongTensor]] = None, + labels: Optional[list[torch.LongTensor]] = None, + texts: Optional[list[str]] = None, + texts_labels: Optional[list[str]] = None, + ): + """ + Run an optimisation step given a list of input_ids, attention_mask, and labels or a list of text and + text_labels. + + Args: + input_ids (list[`torch.LongTensor`]): + List of tensors containing the input_ids (if not provided, text will be used) + attention_mask (list[`torch.LongTensor`], , *optional*): + List of tensors containing the attention_mask + labels (list[`torch.FloatTensor`], *optional*): + List of tensors containing the labels (if set to None, will default to input_ids) + texts (list[`str`], *optional*): + List of strings containing the text input (if not provided, input_ids will directly be used) + texts_labels (list[`str`], *optional*): + List of strings containing the text labels (if set to None, will default to text) + + Returns: + `dict[str, Any]`: A summary of the training statistics + """ + self.model.train() + + if self.state.global_step == 0: + self.tr_loss = torch.tensor(0.0).to(self.args.device) + self._globalstep_last_logged = self.state.global_step + + if input_ids is None and texts is None: + raise ValueError("Step should include `input_ids` or `texts` as keyword arguments.") + elif input_ids is not None and texts is not None: + warnings.warn( + "Both `input_ids` and `texts` argument are provided. `input_ids` will be ignored. " + "Please provide only one of the two.", + UserWarning, + ) + + if labels is None and texts_labels is None and self.is_encoder_decoder: + raise ValueError( + "No 'labels' or 'text_labels' are provided. When using an encoder-decoder architecture, 'labels' or 'text_labels' must be passed." + ) + + # Convert Column to list if not already + input_ids = input_ids[:] if input_ids is not None else None + attention_mask = attention_mask[:] if attention_mask is not None else None + labels = labels[:] if labels is not None else None + texts = texts[:] if texts is not None else None + texts_labels = texts_labels[:] if texts_labels is not None else None + + input_ids, attention_mask, labels, texts, texts_labels = self._step_safety_checker( + input_ids, attention_mask, labels, texts, texts_labels + ) + + if texts is not None: + model_inputs = self.processing_class( + texts, max_length=self.max_length, truncation=True, padding=True, return_tensors="pt" + ) + + input_ids, attention_mask = model_inputs["input_ids"], model_inputs["attention_mask"] + + if texts_labels is not None: + labels = self.processing_class( + texts, max_length=self.max_length, truncation=True, padding=True, return_tensors="pt" + )["input_ids"] + + if labels is None: + labels = input_ids + + model_inputs = self.prepare_model_inputs(input_ids, attention_mask, labels) + + model_inputs_names = list(model_inputs.keys()) + + batch_dict = {} + batch_dict.update(model_inputs) + + def collator(data): + return_dict = dict() + for key in data[0]: + if key in ["input_ids", "attention_mask", "labels"]: + return_dict[key] = torch.stack([d[key] for d in data]).to(self.model.device) + return return_dict + + batch_data = Dataset.from_dict(batch_dict) + batch_data.set_format("torch") + + step_dataloader = DataLoader( + batch_data, + batch_size=self.args.per_device_train_batch_size, + shuffle=True, + collate_fn=collator, + ) + + for _, batch in enumerate(step_dataloader): + with self.accelerator.accumulate(self.model): + model_inputs = {k: batch[k] for k in model_inputs_names} + loss = self.compute_loss(self.model, model_inputs) + + if self.args.n_gpu > 1: + loss = loss.mean() + + tr_loss_step = loss.detach() + + self.accelerator.backward(loss) + + if self.accelerator.sync_gradients and self.args.max_grad_norm is not None: + self.accelerator.clip_grad_norm_( + self.model.parameters(), + self.args.max_grad_norm, + ) + + self.optimizer.step() + self.optimizer.zero_grad() + if self.lr_scheduler is not None: + self.lr_scheduler.step() + + self.state.global_step += 1 + + # update stats etc + self.tr_loss += tr_loss_step + + self._maybe_log_save_evaluate() + + def _maybe_log_save_evaluate(self): + # check if eval is required + if self.args.eval_steps is not None: + if self.state.global_step % self.args.eval_steps == 0 and self.state.global_step != 0: + self.evaluate(self.eval_dataset) + + # check if logging is required + if self.args.logging_steps is not None: + if self.state.global_step % self.args.logging_steps == 0 and self.state.global_step != 0: + logs: dict[str, float] = {} + + tr_loss_scalar = self._nested_gather(self.tr_loss).mean().item() + + # reset tr_loss to zero + self.tr_loss -= self.tr_loss + + logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4) + logs["learning_rate"] = self._get_learning_rate() + + self._globalstep_last_logged = self.state.global_step + + self.log(logs) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="Iterative SFT", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/judges.py b/lib/python3.12/site-packages/trl/trainer/judges.py new file mode 100644 index 0000000000000000000000000000000000000000..b186494998115217c8a5ad171bcc4f04adfabf62 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/judges.py @@ -0,0 +1,459 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import concurrent.futures +import logging +from abc import ABC, abstractmethod +from typing import Optional, Union + +import numpy as np +from accelerate import Accelerator +from huggingface_hub import InferenceClient +from transformers.utils import is_openai_available + +from ..import_utils import is_llm_blender_available + + +if is_llm_blender_available(): + import llm_blender + +if is_openai_available(): + from openai import OpenAI + + +DEFAULT_PAIRWISE_SYSTEM_PROMPT = '''I require a leaderboard for various large language models. I'll provide you with prompts given to these models and their corresponding outputs. Your task is to assess these responses, and select the model that produces the best output from a human perspective. + +## Instruction + +{{ + "instruction": """{prompt}""", +}} + +## Model Outputs + +Here are the unordered outputs from the models. Each output is associated with a specific model, identified by a unique model identifier. + +{{ + {{ + "model_identifier": "0", + "output": """{response0}""" + }}, + {{ + "model_identifier": "1", + "output": """{response1}""" + }} +}} + +## Task + +Evaluate the models on the basis of the quality and relevance of their results, and select the model that generated the best result. Reply with the identifier of the best model. Our evaluation will only take into account the first character of your answer, so make sure it contains only one of the identifiers and nothing else (no quotation marks, no spaces, no new lines, ...). +''' + + +class BaseJudge(ABC): + """ + Base class for judges. The subclasses of this class should implement the `judge` method. + """ + + @abstractmethod + def judge(self, prompts: list[str], completions: list[str], shuffle_order: bool = True) -> list: + raise NotImplementedError("Judge subclasses must implement the `judge` method.") + + +class BaseRankJudge(ABC): + """ + Base class for LLM ranking judges. + + **Example**: + ```python + class MyRankJudge(BaseRankJudge): + def judge(self, prompts, completions, shuffle_order=True): + return ... # Your ranking logic here + + + judge = MyRankJudge() + judge.judge( + prompts=["The capital of France is", "The capital of Germany is"], + completions=[[" Paris", " Marseille", "Lyon"], [" Munich", " Berlin"]], + ) # [[0, 1, 2], [1, 0]] + ``` + """ + + @abstractmethod + def judge(self, prompts: list[str], completions: list[list[str]], shuffle_order: bool = True) -> list[list[int]]: + """ + Judge the completion for the given prompts and return the ranks of each completion. + + Args: + prompts (`list[str]`): + List of prompts. + completions (`list[list[str]]`): + List of completions list, where each element is a list of completions for the corresponding prompt. + shuffle_order (`bool`, *optional*, defaults to `True`): + Whether to shuffle the order of the completions to avoid positional bias. + + Returns: + `list[list[int]]`: + List of lists of idxs, where each list contains the ranks of the completions for the corresponding + prompt. E.g., `[1, 2, 0]` means that the second completion (`idx=1`) is the best, followed by the + third, and then the first. + """ + raise NotImplementedError("Judge subclasses must implement the `judge` method.") + + +class BasePairwiseJudge(BaseJudge): + """ + Base class for pairwise judges. + """ + + @abstractmethod + def judge(self, prompts: list[str], completions: list[list[str]], shuffle_order: bool = True) -> list[int]: + """ + Judge the completion pairs for the given prompts. + + Args: + prompts (`list[str]`): + List of prompts. + completions (`list[list[str]]`): + List of completions pairs, where each element is a pair of completions for the corresponding prompt. + shuffle_order (`bool`, *optional*, defaults to `True`): + Whether to shuffle the order of the completions to avoid positional bias. + + Returns: + `list[int]`: + List of idxs, where each idx is the rank of the best completion for the corresponding prompt. E.g., `1` + means that the second completion (`idx=1`) is the best. + + Note: + If the judge returns `-1` for any prompt, it indicates that the inner process used to compute the + preference has failed. For instance, this could occur if the underlying language model returned an invalid + answer. In such cases, the caller should handle these invalid indices appropriately, possibly by + implementing fallback logic or error handling. + """ + raise NotImplementedError("Judge subclasses must implement the `judge` method.") + + +class BaseBinaryJudge(BaseJudge): + """ + Base class for binary judges. + """ + + @abstractmethod + def judge( + self, + prompts: list[str], + completions: list[str], + gold_completions: Optional[list[str]] = None, + shuffle_order: bool = True, + ) -> list[int]: + """ + Judge the completion for a given prompt. Used to assess if a completion satisfies a constraint. + + This base class should be used to implement binary evaluations as done in section 4.1.4 of the [CGPO + paper](https://huggingface.co/papers/2409.20370). It is relevant for assessing whether a prompt completion pair + satisfies a specific contraint. + + Args: + prompts (`list[str]`): List of prompts. + completions (`list[str]`): List of completions. + gold_completions (`list[str]`, `optional`): List of gold completions if it exists. + shuffle_order (`bool`): Whether to shuffle the order of the completions to avoid positional bias. + + Returns: + list[int]: A list of binary labels: + - 1 indicates that the completion satisfies the evaluated constraint. + - 0 indicates that the completion does not satisfy the evaluated constraint. + + Note: + If the judge returns -1 for any prompt, it indicates that the inner process used to compute the preference + has failed. For instance, this could occur if the underlying language model or rule based contraint + returned an invalid answer. In such cases, the caller should handle these invalid indices appropriately, + possibly by implementing fallback logic or error handling. + """ + raise NotImplementedError("Judge subclasses must implement the `judge` method.") + + +class PairRMJudge(BasePairwiseJudge): + """ + LLM judge based on the PairRM model from AllenAI. + + This judge uses the PairRM model to rank pairs of completions for given prompts. It's designed for pairwise + comparison of language model outputs. The PairRM model is loaded using the llm-blender library and runs on the + default Accelerator device. + + **Attributes**: + + blender (`llm_blender.Blender`): + An instance of the Blender class from llm-blender. + + **Example**: + ```python + >>> pairrm_judge = PairRMJudge() + >>> prompts = ["Translate 'hello' to French", "What's the capital of Japan?"] + >>> completions = [["Bonjour", "Salut"], ["Kyoto", "Tokyo"]] + >>> results = pairrm_judge.judge(prompts, completions) + >>> print(results) # [0, 1] (indicating the first completion is preferred for the first prompt and the second) + ``` + + + + This class requires the llm-blender library to be installed. Install it with: `pip install llm-blender`. + + + """ + + def __init__(self): + if not is_llm_blender_available(): + raise ValueError("llm-blender is not installed. Please install it with `pip install llm-blender`.") + self.blender = llm_blender.Blender() + self.blender.loadranker("llm-blender/PairRM", device=Accelerator().device) + + def judge( + self, + prompts: list[str], + completions: list[list[str]], + shuffle_order: bool = True, + return_scores: bool = False, + temperature: float = 1.0, + ) -> list[Union[int, float]]: + """ + Judge the completion pairs for the given prompts using the PairRM model. + + Args: + prompts (`list[str]`): + List of prompts to judge. + completions (`list[list[str]]`): + List of completion pairs for each prompt. + shuffle_order (`bool`, *optional*, defaults to `True`): + Whether to shuffle the order of the completions to avoid positional bias. + return_scores (`bool`, *optional*, defaults to `False`): + If `True`, return probability scores of the first completion instead of ranks (i.e. a *soft-judge*). + temperature (`float`, *optional*, defaults to `1.0`): + Temperature for scaling logits if `return_scores` is True. + + Returns: + `Union[list[int, float]]`: + If `return_scores` is `False`, returns a list of ranks (`0` or `1`) for each prompt, indicating which + completion is preferred. If `return_scores` is `True`, returns softmax probabilities for the first + completion. + + Raises: + `ValueError`: + If the number of completions per prompt is not exactly 2. + + Note: + Unlike llm-blender, ranks are 0-indexed (`0` means the first completion is preferred). + """ + + if len(completions[0]) != 2: + raise ValueError("PairRM judge requires exactly 2 completions per prompt.") + + # Shuffle the order of the completions to avoid positional bias + if shuffle_order: + flip_mask = np.random.choice([True, False], size=len(prompts)) + completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions)] + + # Rank the completions + ranks = self.blender.rank(prompts, completions, return_scores=return_scores, disable_tqdm=True) + if not return_scores: + ranks -= 1 # PairRM rank is 1-indexed, so we subtract 1 to make it 0-indexed + else: + # scale the logits by temperature + ranks /= temperature + + # Flip back the ranks or scores to the original order if needed + if shuffle_order: + ranks[flip_mask] = ranks[flip_mask][:, ::-1] + + # Return the ranks or score probability + if return_scores: + logit_max = np.amax(ranks, axis=-1, keepdims=True) + exp_logit_shifted = np.exp(ranks - logit_max) + probs = exp_logit_shifted / np.sum(exp_logit_shifted, axis=-1, keepdims=True) + return probs[:, 0].tolist() + else: + return ranks[:, 0].tolist() + + +class HfPairwiseJudge(BasePairwiseJudge): + """ + Pairwise judge based on the Hugging Face API with chat completion. + + This judge is relevant for assessing the quality chat models, where the completion is a response to a given prompt. + + Args: + model (`str`, *optional*, defaults to `"meta-llama/Meta-Llama-3-70B-Instruct"`): + Model to use for the judge. + token (`str`, *optional*): + Hugging Face API token to use for the [`huggingface_hub.InferenceClient`]. + system_prompt (`str` or `None`, *optional*, defaults to `None`): + The system prompt to be used for the judge. If not provided, a default prompt is used. Note that the system + prompt should contain the following placeholders: `{prompt}`, `{response0}`, and `{response1}`. Also, the + inference is called with `max_tokens=1`, consequently the system prompt should ask for a single token + response. + """ + + def __init__( + self, + model="meta-llama/Meta-Llama-3-70B-Instruct", + token: Optional[str] = None, + system_prompt: Optional[str] = None, + ): + self.client = InferenceClient(model=model, token=token) + self.system_prompt = system_prompt or DEFAULT_PAIRWISE_SYSTEM_PROMPT + + def judge(self, prompts: list[str], completions: list[list[str]], shuffle_order: bool = True) -> list[int]: + # Shuffle the order of the completions to avoid positional bias + if shuffle_order: + flip_mask = np.random.choice([True, False], size=len(prompts)) + completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions)] + + # Define a function to get the rank for a single prompt, will be called concurrently + def get_rank(prompt, candidates): + content = self.system_prompt.format(prompt=prompt, response0=candidates[0], response1=candidates[1]) + completion = self.client.chat_completion(messages=[{"role": "user", "content": content}], max_tokens=1) + response = completion.choices[0].message.content + if response in ["0", "1"]: + return int(response) + else: + logging.debug(f"Invalid response from the judge model: '{response}'. Returning -1.") + return -1 + + # Call the completions concurrently + with concurrent.futures.ThreadPoolExecutor() as executor: + ranks = list(executor.map(get_rank, prompts, completions)) + + # Flip back the ranks to the original order if needed + if shuffle_order: + ranks = [ranks[i] if not flip else 1 - ranks[i] for i, flip in enumerate(flip_mask)] + + # Return the ranks + return ranks + + +class OpenAIPairwiseJudge(BasePairwiseJudge): + """ + Judge based on the OpenAI API. + + This judge is relevant for assessing the quality chat models, where the completion is a response to a given prompt. + + Args: + model (`str`, *optional*, defaults to `"gpt-4-turbo-preview"`): + Model to use for the judge. + system_prompt (`str` or `None`, *optional*, defaults to `None`): + System prompt to be used for the judge. If not provided, a default prompt is used. Note that the system + prompt should contain the following placeholders: `{prompt}`, `{response0}`, and `{response1}`. Also, the + inference is called with `max_tokens=1`, consequently the system prompt should ask for a single token + response. + max_requests (`int` or `None`, *optional*, defaults to `1000`): + Maximum number of requests to make to the OpenAI API. If set to `None`, there is no limit. + """ + + def __init__( + self, model="gpt-4-turbo-preview", system_prompt: Optional[str] = None, max_requests: Union[int, None] = 1_000 + ): + if not is_openai_available(): + raise ValueError("OpenAI client is not installed. Please install it with 'pip install openai'.") + self.client = OpenAI() + self.model = model + self.system_prompt = system_prompt or DEFAULT_PAIRWISE_SYSTEM_PROMPT + self.max_requests = max_requests + self.num_requests = 0 + self._warned = False + + def judge(self, prompts: list[str], completions: list[list[str]], shuffle_order: bool = True) -> list[int]: + # Check if the limit of requests is reached, if so, use random choice instead + if self.max_requests is not None and self.num_requests >= self.max_requests: + if not self._warned: # Print the warning only once + logging.warning( + f"Reached the maximum number of requests ({self.max_requests}). From now on, returning -1 instead. " + " To increase the limit, set `max_requests` to a higher value, or to `None` for no limit." + ) + self._warned = True + return [-1] * len(prompts) + + # Shuffle the order of the completions to avoid positional bias + if shuffle_order: + flip_mask = np.random.choice([True, False], size=len(prompts)) + completions = [pair[::-1] if flip else pair for flip, pair in zip(flip_mask, completions)] + + # Define a function to get the rank for a single prompt, will be called concurrently + def get_rank(prompt, candidates): + content = self.system_prompt.format(prompt=prompt, response0=candidates[0], response1=candidates[1]) + messages = [{"role": "user", "content": content}] + completion = self.client.chat.completions.create(model=self.model, messages=messages, max_tokens=1) + response = completion.choices[0].message.content + if response in ["0", "1"]: + return int(response) + else: + logging.debug(f"Invalid response from the judge model: '{response}'. Returning -1.") + return -1 + + # Call the completions concurrently + with concurrent.futures.ThreadPoolExecutor() as executor: + ranks = list(executor.map(get_rank, prompts, completions)) + + # Flip back the ranks to the original order if needed + if shuffle_order: + ranks = [ranks[i] if not flip else 1 - ranks[i] for i, flip in enumerate(flip_mask)] + + # Update the number of requests + self.num_requests += len(prompts) + + # Return the ranks + return ranks + + +class AllTrueJudge(BaseBinaryJudge): + """ + Unify the decision of multiple [`BaseBinaryJudge`] instances. + + Returns `1` only if all inner binary judges return `1`. If any judge returns `0`, it returns `0`. If any judge + returns `-1`, indicating a failure in its process, this judge will also return `-1`. + + Implements the Mixture of Judges as described in the [CGPO paper](https://huggingface.co/papers/2409.20370). + + Args: + judges (`list[BaseBinaryJudge]`): A list of [`BaseBinaryJudge`] instances whose decisions will be unified. + """ + + def __init__(self, judges: list[BaseBinaryJudge]): + self.judges = judges + + def judge( + self, + prompts: list[str], + completions: list[str], + gold_completions: Optional[list[str]] = None, + shuffle_order: bool = True, + ) -> list[int]: + all_binary_judgments = [ + judge.judge(prompts, completions, gold_completions, shuffle_order) for judge in self.judges + ] + output = [] + for binary_judgments in zip(*all_binary_judgments): + # Check that all values are in {0, 1, -1} + if any(binary_judgment not in {0, 1, -1} for binary_judgment in binary_judgments): + raise ValueError( + f"Invalid binary judgment: {binary_judgments}, expected list of values in {{0, 1, -1}}." + ) + + # Unify the decision + if -1 in binary_judgments: + output.append(-1) + elif all(binary_judgment == 1 for binary_judgment in binary_judgments): + output.append(1) + else: + output.append(0) + return output diff --git a/lib/python3.12/site-packages/trl/trainer/kto_config.py b/lib/python3.12/site-packages/trl/trainer/kto_config.py new file mode 100644 index 0000000000000000000000000000000000000000..bb823125d60b4543febbe4a7d68a79097cab2440 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/kto_config.py @@ -0,0 +1,235 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Any, Optional + +from transformers import TrainingArguments + + +@dataclass +class KTOConfig(TrainingArguments): + r""" + Configuration class for the [`KTOTrainer`]. + + This class includes only the parameters that are specific to KTO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. + loss_type (`str`, *optional*, defaults to `"kto"`): + Type of loss to use. Possible values are: + + - `"kto"`: KTO loss from the [KTO](https://huggingface.co/papers/2402.01306) paper. + - `"apo_zero_unpaired"`: Unpaired variant of APO-zero loss from the + [APO](https://huggingface.co/papers/2408.06266) paper. + + desirable_weight (`float`, *optional*, defaults to `1.0`): + Desirable losses are weighed by this factor to counter unequal number of desirable and undesirable paris. + undesirable_weight (`float`, *optional*, defaults to `1.0`): + Undesirable losses are weighed by this factor to counter unequal number of desirable and undesirable pairs. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from both the model and the reference model to W&B or Comet + during evaluation. + is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + precompute_ref_log_probs (`bool`, *optional*, defaults to `False`): + Whether to precompute reference model log probabilities for training and evaluation datasets. This is + useful when training without the reference model to reduce the total GPU memory needed. + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + ref_model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the reference model + from a string. + dataset_num_proc: (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + use_liger_loss (`bool`, *optional*, defaults to `False`): + Whether to use Liger loss. It requires liger-kernel to be installed. + base_model_attribute_name (`str`, *optional*, defaults to `"model"`): + Name of the attribute in the model that contains the base model. This is used to get the base model from + the model when the model does not have a `get_decoder` method in the case when `use_liger_loss` is `True`. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs", "ref_model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + max_length: Optional[int] = field( + default=1024, + metadata={"help": "Maximum length of the sequences (prompt + completion) in the batch."}, + ) + max_prompt_length: Optional[int] = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + max_completion_length: Optional[int] = field( + default=None, + metadata={ + "help": "Maximum length of the completion. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the deviation from the reference model. Higher β means less deviation from " + "the reference model." + }, + ) + loss_type: str = field( + default="kto", + metadata={ + "help": "Type of loss to use.", + "choices": ["kto", "apo_zero_unpaired"], + }, + ) + desirable_weight: float = field( + default=1.0, + metadata={ + "help": "Desirable losses are weighed by this factor to counter unequal number of desirable and " + "undesirable pairs.", + }, + ) + undesirable_weight: float = field( + default=1.0, + metadata={ + "help": "Undesirable losses are weighed by this factor to counter unequal number of desirable and " + "undesirable pairs.", + }, + ) + label_pad_token_id: int = field( + default=-100, + metadata={ + "help": "Label pad token id. This argument is required if you want to use the default data collator." + }, + ) + padding_value: Optional[int] = field( + default=None, + metadata={"help": "Padding value to use. If `None`, the padding value of the tokenizer is used."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the prompt is too long.", + "choices": ["keep_end", "keep_start"], + }, + ) + generate_during_eval: bool = field( + default=False, + metadata={ + "help": "If `True`, generates and logs completions from both the model and the reference model to W&B " + "during evaluation." + }, + ) + is_encoder_decoder: Optional[bool] = field( + default=None, + metadata={ + "help": "When using the `model_init` argument (callable) to instantiate the model instead of the `model` " + "argument, you need to specify if the model returned by the callable is an encoder-decoder model." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + precompute_ref_log_probs: bool = field( + default=False, + metadata={ + "help": "Whether to precompute reference model log probabilities for training and evaluation datasets. " + "This is useful when training without the reference model to reduce the total GPU memory needed." + }, + ) + model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model " + "from a string." + }, + ) + ref_model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the " + "reference model from a string." + }, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + use_liger_loss: bool = field( + default=False, + metadata={"help": "Whether to use Liger loss. It requires liger-kernel to be installed."}, + ) + base_model_attribute_name: str = field( + default="model", + metadata={ + "help": "Name of the attribute in the model that contains the base model. This is used to get the base " + "model from the model when the model does not have a `get_decoder` method in the case when " + "`use_liger_loss` is `True`." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/lib/python3.12/site-packages/trl/trainer/kto_trainer.py b/lib/python3.12/site-packages/trl/trainer/kto_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..aaf23633f504bc77b37fc7ce8defe6f88c010511 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/kto_trainer.py @@ -0,0 +1,1736 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import random +import textwrap +import warnings +from collections import defaultdict +from contextlib import contextmanager, nullcontext +from operator import itemgetter +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState +from accelerate.utils import tqdm +from datasets import Dataset, concatenate_datasets +from torch import autocast +from torch.utils.data import DataLoader, SequentialSampler +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + TrainerCallback, + TrainingArguments, + is_comet_available, + is_wandb_available, +) +from transformers.trainer_utils import EvalLoopOutput, has_length +from transformers.utils import is_peft_available + +from ..data_utils import maybe_apply_chat_template, maybe_extract_prompt, maybe_unpair_preference_dataset +from ..import_utils import is_liger_kernel_available +from ..models import create_reference_model, prepare_deepspeed +from .kto_config import KTOConfig +from .utils import ( + DPODataCollatorWithPadding, + disable_dropout_in_model, + generate_model_card, + get_comet_experiment_url, + log_table_to_comet_experiment, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_liger_kernel_available(): + from liger_kernel.chunked_loss import LigerFusedLinearKTOLoss + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + +if is_wandb_available(): + import wandb + + +if TYPE_CHECKING: + from transformers import PreTrainedModel, PreTrainedTokenizer + +RUNNING_NAME = "running.pt" + + +def _get_kl_dataset(batch: dict[str, list[Any]]) -> dict[str, list[Any]]: + """ + Creates mismatched pairs of prompts and completions for the KL dataset by adding a +1 offset to the order of + completions. For best results, the mismatched outputs y' used to estimate the KL term for a batch should be the + same set as the matched outputs y used to estimate the rewards in that batch, just paired with different x. + """ + batch["answer_input_ids"] = [batch["answer_input_ids"][-1]] + batch["answer_input_ids"][:-1] + batch["answer_attention_mask"] = [batch["answer_attention_mask"][-1]] + batch["answer_attention_mask"][:-1] + return batch + + +def _tokenize( + batch: dict[str, list[Any]], + tokenizer: "PreTrainedTokenizer", +) -> dict[str, list[Any]]: + """Tokenize a batch from a KTO specific dataset.""" + prompt_tokenized = tokenizer(batch["prompt"], add_special_tokens=False) + prompt_input_ids = prompt_tokenized["input_ids"] + prompt_attention_mask = prompt_tokenized["attention_mask"] + prompt_and_completion = [prompt + completion for prompt, completion in zip(batch["prompt"], batch["completion"])] + full_tokenized = tokenizer(prompt_and_completion, add_special_tokens=False) + full_input_ids = full_tokenized["input_ids"] + full_attention_mask = full_tokenized["attention_mask"] + + answer_input_ids = [f[len(p) :] for f, p in zip(full_input_ids, prompt_input_ids)] + answer_attention_mask = [f[len(p) :] for f, p in zip(full_attention_mask, prompt_attention_mask)] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = [np.concatenate([p, a]) for p, a in zip(prompt_input_ids, answer_input_ids)] + # Prepare input tokens for token by token comparison + full_input_ids = [np.array(f) for f in full_input_ids] + for full, concat in zip(full_input_ids, full_concat_input_ids): + if len(full) != len(concat): + raise ValueError( + "The elements in 'full_input_ids' and 'full_concat_input_ids' must have the same pairwise length." + ) + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = [len(p) for p in prompt_input_ids] + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + for idx, (p, f, r) in enumerate(zip(prompt_input_ids, full_input_ids, response_token_ids_start_idx)): + if not np.array_equal(p, f[:r]): + response_token_ids_start_idx[idx] -= 1 + + prompt_input_ids = [f[:r] for f, r in zip(full_input_ids, response_token_ids_start_idx)] + prompt_attention_mask = [f[:r] for f, r in zip(full_attention_mask, response_token_ids_start_idx)] + + for p, m in zip(prompt_input_ids, prompt_attention_mask): + if len(p) != len(m): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = [f[r:] for f, r in zip(full_input_ids, response_token_ids_start_idx)] + answer_attention_mask = [f[r:] for f, r in zip(full_attention_mask, response_token_ids_start_idx)] + + output = dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + answer_input_ids=answer_input_ids, + answer_attention_mask=answer_attention_mask, + ) + + return output + + +def _process_tokens(example: dict[str, Any], model: "PreTrainedModel" = None, **kwargs) -> dict: + """Process tokens of a KTO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + completion responses is/are too long. First we truncate the prompt; if we're still too long, we truncate the + completion. + + We also create the labels for the completion responses, which are of length equal to the sum of the length of the + prompt and the completion response, with label_pad_token_id for the prompt tokens. + """ + prompt = example["prompt"] + completion = example["completion"] + + batch = { + f"{kwargs['prefix']}prompt": prompt, + f"{kwargs['prefix']}completion": completion, + f"{kwargs['prefix']}label": example["label"], + } + + if not kwargs["is_encoder_decoder"]: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + + if not isinstance(completion, str): + raise ValueError(f"completion should be an str but got {type(completion)}") + + # keys of format prompt_* refers to just the prompt and answer_* refers to just the answer + all_tokens = { + "prompt_input_ids": example["prompt_input_ids"], + "prompt_attention_mask": example["prompt_attention_mask"], + "answer_input_ids": example["answer_input_ids"], + "answer_attention_mask": example["answer_attention_mask"], + } + + # calculate max length by checking if BOS/EOS is already there + max_length = kwargs["max_length"] + bos_token_id = kwargs["tokenizer"].bos_token_id + eos_token_id = kwargs["tokenizer"].eos_token_id + if len(all_tokens["prompt_input_ids"]) > 0 and bos_token_id != all_tokens["prompt_input_ids"][0]: + max_length -= 1 + if len(all_tokens["answer_input_ids"]) > 0 and eos_token_id != all_tokens["answer_input_ids"][-1]: + max_length -= 1 + + # if combined sequence is too long (> max_length - 1 for BOS token - 1 for EOS), truncate the prompt + if len(all_tokens["prompt_input_ids"]) + len(all_tokens["answer_input_ids"]) > max_length: + for k in ["prompt_input_ids", "prompt_attention_mask"]: + if kwargs["truncation_mode"] == "keep_start": + all_tokens[k] = all_tokens[k][: kwargs["max_prompt_length"]] + elif kwargs["truncation_mode"] == "keep_end": + all_tokens[k] = all_tokens[k][-kwargs["max_prompt_length"] :] + else: + raise ValueError(f"Unknown truncation mode: {kwargs['truncation_mode']}") + + # if that's still too long, truncate the response + if len(all_tokens["prompt_input_ids"]) + len(all_tokens["answer_input_ids"]) > max_length: + for k in ["answer_input_ids", "answer_attention_mask"]: + all_tokens[k] = all_tokens[k][: max_length - kwargs["max_prompt_length"]] + + # all input_ids and attention mask as is. We then check if we need to add BOS/EOS tokens + batch[f"{kwargs['prefix']}prompt_input_ids"] = all_tokens["prompt_input_ids"] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = all_tokens["prompt_attention_mask"] + batch[f"{kwargs['prefix']}completion_input_ids"] = ( + all_tokens["prompt_input_ids"] + all_tokens["answer_input_ids"] + ) + batch[f"{kwargs['prefix']}completion_attention_mask"] = ( + all_tokens["prompt_attention_mask"] + all_tokens["answer_attention_mask"] + ) + + # add BOS, which affects both prompt and the full completion + if bos_token_id is not None: + if len(all_tokens["prompt_input_ids"]) == 0 or bos_token_id != all_tokens["prompt_input_ids"][0]: + batch[f"{kwargs['prefix']}prompt_input_ids"] = [bos_token_id] + batch[ + f"{kwargs['prefix']}prompt_input_ids" + ] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = [1] + batch[ + f"{kwargs['prefix']}prompt_attention_mask" + ] + batch[f"{kwargs['prefix']}completion_input_ids"] = [bos_token_id] + batch[ + f"{kwargs['prefix']}completion_input_ids" + ] + batch[f"{kwargs['prefix']}completion_attention_mask"] = [1] + batch[ + f"{kwargs['prefix']}completion_attention_mask" + ] + # add EOS, which affects only the full completion + if len(all_tokens["answer_input_ids"]) == 0 or eos_token_id != all_tokens["answer_input_ids"][-1]: + batch[f"{kwargs['prefix']}completion_input_ids"] = batch[f"{kwargs['prefix']}completion_input_ids"] + [ + eos_token_id + ] + batch[f"{kwargs['prefix']}completion_attention_mask"] = batch[ + f"{kwargs['prefix']}completion_attention_mask" + ] + [1] + + batch[f"{kwargs['prefix']}completion_labels"] = batch[f"{kwargs['prefix']}completion_input_ids"][:] + batch[f"{kwargs['prefix']}completion_labels"][: len(batch[f"{kwargs['prefix']}prompt_input_ids"])] = [ + kwargs["label_pad_token_id"] + ] * len(batch[f"{kwargs['prefix']}prompt_input_ids"]) + else: + completion_tokens = kwargs["tokenizer"]( + completion, truncation=True, max_length=kwargs["max_completion_length"], add_special_tokens=True + ) + prompt_tokens = kwargs["tokenizer"]( + prompt, truncation=True, max_length=kwargs["max_prompt_length"], add_special_tokens=True + ) + + batch[f"{kwargs['prefix']}prompt_input_ids"] = prompt_tokens["input_ids"] + batch[f"{kwargs['prefix']}prompt_attention_mask"] = prompt_tokens["attention_mask"] + + batch[f"{kwargs['prefix']}completion_labels"] = completion_tokens["input_ids"] + batch[f"{kwargs['prefix']}completion_attention_mask"] = completion_tokens["attention_mask"] + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch[f"{kwargs['prefix']}completion_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["completion_labels"]) + ) + + return batch + + +class KTOTrainer(Trainer): + r""" + Initialize KTOTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForSequenceClassification`. + ref_model (`PreTrainedModelWrapper`): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + args (`KTOConfig`): + The arguments to use for training. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + data_collator (`transformers.DataCollator`, *optional*, defaults to `None`): + The data collator to use for training. If None is specified, the default data collator + (`DPODataCollatorWithPadding`) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + model_adapter_name (`str`, defaults to `None`): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str`, defaults to `None`): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + """ + + _tag_names = ["trl", "kto"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module, str] = None, + ref_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: KTOConfig = None, + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + data_collator: Optional[DataCollator] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional[dict] = None, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + model_adapter_name: Optional[str] = None, + ref_adapter_name: Optional[str] = None, + ): + if type(args) is TrainingArguments: + raise ValueError("Please use `KTOConfig` instead TrainingArguments.") + + if not isinstance(model, str) and ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, you must mass a copy of it, or `None` if you use peft." + ) + + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the KTOTrainer. But your model is already instantiated.") + else: + model_init_kwargs = args.model_init_kwargs + torch_dtype = model_init_kwargs.get("torch_dtype") + if torch_dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(torch_dtype, str) and torch_dtype != "auto": + torch_dtype = getattr(torch, torch_dtype) + if torch_dtype != "auto" and not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"Invalid `torch_dtype` passed to the KTOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}." + ) + model_init_kwargs["torch_dtype"] = torch_dtype + + if args.ref_model_init_kwargs is None: + ref_model_init_kwargs = {} + elif not isinstance(ref_model, str): + raise ValueError( + "You passed ref_model_kwargs to the KTOTrainer. But your ref_model is already instantiated." + ) + else: + ref_model_init_kwargs = args.ref_model_init_kwargs + torch_dtype = ref_model_init_kwargs.get("torch_dtype") + if torch_dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(torch_dtype, str) and torch_dtype != "auto": + torch_dtype = getattr(torch, torch_dtype) + if torch_dtype != "auto" and not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"Invalid `torch_dtype` passed to the KTOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}." + ) + ref_model_init_kwargs["torch_dtype"] = torch_dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + + if isinstance(ref_model, str): + ref_model = AutoModelForCausalLM.from_pretrained(ref_model, **ref_model_init_kwargs) + + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it with `pip install peft` to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." + " Please install `wandb` or `comet-ml` to resolve." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + self.is_peft_model = is_peft_available() and isinstance(model, PeftModel) + self.model_adapter_name = model_adapter_name + self.ref_adapter_name = ref_adapter_name + + if ref_model: + self.ref_model = ref_model + elif self.is_peft_model or args.precompute_ref_log_probs: + # The `model` with adapters turned off will be used as the reference model + self.ref_model = None + else: + self.ref_model = create_reference_model(model) + + if processing_class is None: + raise ValueError( + "max_length or a processing_class must be specified when using the default DPODataCollatorWithPadding" + ) + if args.max_length is None: + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `max_length` in the KTOTrainer's init" + " it will be set to `512` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_length = 512 + if args.max_length is not None: + max_length = args.max_length + + if args.max_prompt_length is None: + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `max_prompt_length` in the KTOTrainer's init" + " it will be set to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_prompt_length = 128 + if args.max_prompt_length is not None: + max_prompt_length = args.max_prompt_length + + max_completion_length = None + if args.max_completion_length is None and self.is_encoder_decoder: + warnings.warn( + "When using DPODataCollatorWithPadding with an encoder decoder architecture, you should set `max_completion_length` in the KTOTrainer's init" + " it will be set to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_completion_length = 128 + if args.max_completion_length is not None and self.is_encoder_decoder: + max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your KTOConfig" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + self.loss_type = args.loss_type + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.max_completion_length = max_completion_length + self.processing_class = processing_class + self.precompute_ref_log_probs = args.precompute_ref_log_probs + + # Not all losses require a KL calculation + self.calculate_KL = True + if self.loss_type in ["apo_zero_unpaired"]: + self.calculate_KL = False + + # Since ref_logs are precomputed on the first call to get_train/eval_dataloader + # keep track of first called to avoid computation of future calls + self._precomputed_train_ref_log_probs = False + self._precomputed_eval_ref_log_probs = False + + # metric + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # KTO parameter + self.beta = args.beta + self.desirable_weight = args.desirable_weight + self.undesirable_weight = args.undesirable_weight + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + warnings.warn( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + UserWarning, + ) + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in KTO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids" and "completion_input_ids". As a result, + # the trainer issues the warning: "Could not estimate the number of tokens of the input, floating-point + # operations will not be computed." To suppress this warning, we set the "estimate_tokens" key in the model's + # "warnings_issued" dictionary to True. This acts as a flag to indicate that the warning has already been + # issued. + model.warnings_issued["estimate_tokens"] = True + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed + train_dataset = train_dataset.map( + maybe_extract_prompt, num_proc=args.dataset_num_proc, desc="Extracting prompt from train dataset" + ) + # Unpair the dataset if needed + train_dataset = maybe_unpair_preference_dataset( + train_dataset, args.dataset_num_proc, desc="Unpairing train dataset" + ) + # Apply the chat template if needed + train_dataset = train_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + desc="Applying chat template to train dataset", + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + maybe_extract_prompt, num_proc=args.dataset_num_proc, desc="Extracting prompt from eval dataset" + ) + eval_dataset = maybe_unpair_preference_dataset( + eval_dataset, args.dataset_num_proc, desc="Unpairing eval dataset" + ) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + desc="Applying chat template to eval dataset", + ) + + # Tokenize and prepare the training datasets + train_dataset = train_dataset.map( + _tokenize, + batched=True, + fn_kwargs={"tokenizer": self.processing_class}, + num_proc=args.dataset_num_proc, + desc="Tokenizing train dataset", + ) + + fn_kwargs = { + "prefix": "", + "is_encoder_decoder": self.is_encoder_decoder, + "tokenizer": self.processing_class, + "max_length": self.max_length, + "truncation_mode": self.truncation_mode, + "label_pad_token_id": self.label_pad_token_id, + "max_prompt_length": self.max_prompt_length, + "max_completion_length": self.max_completion_length, + } + + train_dataset = train_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized train dataset", + ) + + # Tokenize and prepare the eval datasets + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + _tokenize, + fn_kwargs={"tokenizer": self.processing_class}, + batched=True, + num_proc=args.dataset_num_proc, + desc="Tokenizing eval dataset", + ) + + eval_dataset = eval_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + desc="Processing tokenized eval dataset", + ) + + # Get KL datasets if needed + if self.calculate_KL: + if args.per_device_train_batch_size <= 1: + raise ValueError( + "Actual (not effective) batch size must be > 1. KTO will not work properly because the KL term will be equivalent to the implied reward." + ) + + # create pairs for estimating the KL term by flipping the matched pairs in each batch of size total_batch_size + # i.e., (x_1, y_1), ..., (x_n, y_n) --> (x_1, y_n), ..., (x_n, y_1) = (x'_1, y'_1), ..., (x'_n, y'_n) + train_kl_dataset = train_dataset.map( + _get_kl_dataset, + batched=True, + batch_size=args.per_device_train_batch_size, + num_proc=args.dataset_num_proc, + desc="Extracting KL train dataset", + ) + + fn_kwargs["prefix"] = "KL_" + train_kl_dataset = train_kl_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=[c for c in train_kl_dataset.column_names if c in train_dataset.column_names], + desc="Processing tokenized train KL dataset", + ) + + # merge the datasets + train_dataset = concatenate_datasets([train_dataset, train_kl_dataset], axis=1) + + if eval_dataset is not None: + # Get KL dataset + eval_kl_dataset = eval_dataset.map( + _get_kl_dataset, + batched=True, + batch_size=args.per_device_train_batch_size, + num_proc=args.dataset_num_proc, + desc="Extracting eval KL dataset", + ) + + eval_kl_dataset = eval_kl_dataset.map( + _process_tokens, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=[c for c in eval_kl_dataset.column_names if c in eval_dataset.column_names], + desc="Processing tokenized eval KL dataset", + ) + + # merge the datasets + eval_dataset = concatenate_datasets([eval_dataset, eval_kl_dataset], axis=1) + + # calculate dataset desirability balance + num_desirable = max(sum(train_dataset["label"]), 1) + num_undesirable = max(len(train_dataset["label"]) - num_desirable, 1) # "label" is binary + + if num_desirable != num_undesirable: + # The lower and upper bounds come from Eq. (8) of https://huggingface.co/papers/2402.01306 + des_weight_lower_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1, 2) + des_weight_upper_bound = round((num_undesirable * self.undesirable_weight / num_desirable) * 1.33, 2) + und_weight_lower_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1.33, 2) + und_weight_upper_bound = round((num_desirable * self.desirable_weight / num_undesirable) / 1, 2) + + des_weight_in_range = des_weight_lower_bound <= self.desirable_weight <= des_weight_upper_bound + und_weight_in_range = und_weight_lower_bound <= self.undesirable_weight <= und_weight_upper_bound + + if not (des_weight_in_range or und_weight_in_range): + warnings.warn( + "You have different amounts of desirable/positive and undesirable/negative examples but the " + "weights on the desirable and undesirable losses don't seem to be in an ideal range. Based " + f"on your data, we recommend EITHER " + f"desirable_weight in [{des_weight_lower_bound}, {des_weight_upper_bound}] or " + f"undesirable_weight in [{und_weight_lower_bound}, {und_weight_upper_bound}] (but NOT BOTH). " + "See the documentation on how to optimally set these weights.", + UserWarning, + ) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the + # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set + # self.model_accepts_loss_kwargs to False to enable scaling. + self.model_accepts_loss_kwargs = False + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + # Deepspeed Zero-3 does not support precompute_ref_log_probs + if self.is_deepspeed_enabled: + if self.accelerator.state.deepspeed_plugin.zero_stage == 3 and self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with Deepspeed ZeRO-3. Please set `precompute_ref_log_probs=False`." + ) + + if self.ref_model is None: + if not (self.is_peft_model or self.precompute_ref_log_probs): + raise ValueError( + "No reference model and model is not a Peft model. Try setting `precompute_ref_log_probs=True`" + ) + else: + if self.is_deepspeed_enabled: + self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator) + else: + self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True) + + # Import Liger loss if enabled + if self.args.use_liger_loss: + if not is_liger_kernel_available(): + raise ImportError( + "You set `use_liger_loss=True` but the liger kernel is not available. " + "Please install liger-kernel first: `pip install liger-kernel`" + ) + if self.loss_type in ["apo_zero_unpaired"]: + raise ValueError( + "You cannot set `loss_type='apo_zero_unpaired'` with liger-kernel." + "Only KTO loss is supported with liger-kernel." + ) + if self.precompute_ref_log_probs: + raise ValueError( + "You cannot use `precompute_ref_log_probs=True` with liger kernel. Please set " + "`precompute_ref_log_probs=False`." + ) + if self.is_peft_model or self.ref_adapter_name is not None: + raise ValueError( + "You cannot use `use_liger_loss=True` with Peft models. Please set `use_liger_loss=False`." + ) + self.kto_loss_fn = LigerFusedLinearKTOLoss( + ignore_index=self.label_pad_token_id, beta=self.beta, use_ref_model=(self.ref_model is not None) + ) + + @contextmanager + def null_ref_context(self): + """Context manager for handling null reference model (that is, peft adapter manipulation).""" + with ( + self.accelerator.unwrap_model(self.model).disable_adapter() + if self.is_peft_model and not self.ref_adapter_name + else nullcontext() + ): + if self.ref_adapter_name: + self.model.set_adapter(self.ref_adapter_name) + yield + if self.ref_adapter_name: + self.model.set_adapter(self.model_adapter_name or "default") + + def get_train_dataloader(self) -> DataLoader: + """ + Returns the training [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_train_dataloader to precompute `ref_log_probs`. + """ + + if self.precompute_ref_log_probs and not self._precomputed_train_ref_log_probs: + dataloader_params = { + "batch_size": self.args.per_device_train_batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(self.train_dataset, **dataloader_params)) + reference_completion_logps = [] + reference_KL_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Train dataset reference log probs"): + reference_completion_logp, reference_KL_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + if self.calculate_KL: + reference_KL_logp = self.accelerator.gather_for_metrics(reference_KL_logp) + reference_KL_logps.append(reference_KL_logp.cpu()) + + self.train_dataset = self.train_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + + if self.calculate_KL: + self.train_dataset = self.train_dataset.add_column( + name="reference_KL_logps", column=torch.cat(reference_KL_logps).float().numpy() + ) + + self._precomputed_train_ref_log_probs = True + + return super().get_train_dataloader() + + def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader: + """ + Returns the evaluation [`~torch.utils.data.DataLoader`]. + + Subclass of transformers.src.transformers.trainer.get_eval_dataloader to precompute `ref_log_probs`. + + Args: + eval_dataset (`torch.utils.data.Dataset`, *optional*): + If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted + by the `model.forward()` method are automatically removed. It must implement `__len__`. + """ + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset + + if self.precompute_ref_log_probs and not self._precomputed_eval_ref_log_probs: + dataloader_params = { + "batch_size": self.args.per_device_eval_batch_size, + "collate_fn": self.data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "shuffle": False, + } + + # prepare dataloader + data_loader = self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params)) + + reference_completion_logps = [] + reference_KL_logps = [] + + for padded_batch in tqdm(iterable=data_loader, desc="Eval dataset reference log probs"): + reference_completion_logp, reference_KL_logp = self.compute_reference_log_probs(padded_batch) + + reference_completion_logp = self.accelerator.gather_for_metrics(reference_completion_logp) + reference_completion_logps.append(reference_completion_logp.cpu()) + + if self.calculate_KL: + reference_KL_logp = self.accelerator.gather_for_metrics(reference_KL_logp) + reference_KL_logps.append(reference_KL_logp.cpu()) + + eval_dataset = eval_dataset.add_column( + name="reference_logps", column=torch.cat(reference_completion_logps).float().numpy() + ) + if self.calculate_KL: + eval_dataset = eval_dataset.add_column( + name="reference_KL_logps", column=torch.cat(reference_KL_logps).float().numpy() + ) + + # Save calculated reference_chosen_logps and reference_rejected_logps to the eval_dataset for subsequent runs + if self.eval_dataset is not None: + self.eval_dataset = eval_dataset + self._precomputed_eval_ref_log_probs = True + + return super().get_eval_dataloader(eval_dataset=eval_dataset) + + def compute_reference_log_probs(self, padded_batch: dict) -> dict: + """Computes log probabilities of the reference model for a single padded batch of a KTO specific dataset.""" + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + if self.is_encoder_decoder: + completion_logits = self.model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + if self.calculate_KL: + KL_logits = self.model( + padded_batch["KL_prompt_input_ids"], + attention_mask=padded_batch["KL_prompt_attention_mask"], + decoder_input_ids=padded_batch.get("KL_completion_decoder_input_ids"), + labels=padded_batch["KL_completion_labels"], + ).logits + else: + completion_logits = self.model( + padded_batch["completion_input_ids"], + attention_mask=padded_batch["completion_attention_mask"], + ).logits + + if self.calculate_KL: + KL_logits = self.model( + padded_batch["KL_completion_input_ids"], + attention_mask=padded_batch["KL_completion_attention_mask"], + ).logits + else: + if self.is_encoder_decoder: + completion_logits = self.ref_model( + padded_batch["prompt_input_ids"], + attention_mask=padded_batch["prompt_attention_mask"], + decoder_input_ids=padded_batch.get("completion_decoder_input_ids"), + labels=padded_batch["completion_labels"], + ).logits + + if self.calculate_KL: + KL_logits = self.ref_model( + padded_batch["KL_prompt_input_ids"], + attention_mask=padded_batch["KL_prompt_attention_mask"], + decoder_input_ids=padded_batch.get("KL_completion_decoder_input_ids"), + labels=padded_batch["KL_completion_labels"], + ).logits + else: + completion_logits = self.ref_model( + padded_batch["completion_input_ids"], attention_mask=padded_batch["completion_attention_mask"] + ).logits + + if self.calculate_KL: + KL_logits = self.ref_model( + padded_batch["KL_completion_input_ids"], + attention_mask=padded_batch["KL_completion_attention_mask"], + ).logits + + completion_logps = self.get_batch_logps( + completion_logits, + padded_batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + if self.calculate_KL: + KL_logps = self.get_batch_logps( + KL_logits, + padded_batch["KL_completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + else: + KL_logps = None + + return completion_logps, KL_logps + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: + Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + else: + # Fixes end-dec RuntimeError + labels = labels.clone() + + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels[labels == label_pad_token_id] = 0 + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + KL_logps = self._compute_kl_logps(model, batch) + + model_kwargs = ( + { + "labels": batch["completion_labels"], + "decoder_input_ids": batch.get("completion_decoder_input_ids"), + } + if self.is_encoder_decoder + else {} + ) + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + **model_kwargs, + ) + completion_logits = outputs.logits + + completion_logps = self.get_batch_logps( + completion_logits, + batch["completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + if completion_logps.shape[0] != len(batch["label"]): + raise ValueError( + "There is a mismatch between the number of examples in this batch and the number of " + "examples for which an output sequence was predicted." + ) + + chosen_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(completion_logps.shape[0]) if batch["label"][i] is False] + + chosen_logps = completion_logps[chosen_idx, ...] + rejected_logps = completion_logps[rejected_idx, ...] + + chosen_logits = completion_logits[chosen_idx, ...] + rejected_logits = completion_logits[rejected_idx, ...] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, KL_logps, outputs.aux_loss) + else: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, KL_logps) + + def kto_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + policy_KL_logps: torch.FloatTensor, + reference_chosen_logps: torch.FloatTensor, + reference_rejected_logps: torch.FloatTensor, + reference_KL_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute the KTO loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (num(chosen) in batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (num(rejected) in batch_size,) + policy_KL_logps: Log probabilities of the policy model for the KL responses. Shape: (batch_size,) + reference_chosen_logps: + Log probabilities of the reference model for the chosen responses. Shape: (num(chosen) in batch_size,) + reference_rejected_logps: + Log probabilities of the reference model for the rejected responses. Shape: (num(rejected) in + batch_size,) + reference_KL_logps: Log probabilities of the reference model for the KL responses. Shape: (batch_size,) + + Returns: + A tuple of four tensors: (losses, chosen_rewards, rejected_rewards, KL). The losses tensor contains the KTO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. The KL tensor contains the detached KL divergence estimate + between the policy and reference models. + """ + if self.calculate_KL: + kl = (policy_KL_logps - reference_KL_logps).mean().detach() + kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) + else: + kl = torch.zeros(1).to(policy_chosen_logps.device) + + # Chosen losses + if policy_chosen_logps.shape[0] != 0 or reference_chosen_logps.shape[0] != 0: + chosen_logratios = policy_chosen_logps - reference_chosen_logps + + if self.loss_type == "kto": + # Eqn (7) of the KTO paper (https://huggingface.co/papers/2402.01306) + chosen_losses = 1 - F.sigmoid(self.beta * (chosen_logratios - kl)) + elif self.loss_type == "apo_zero_unpaired": + # Unpaired variant of Eqn (7) of the APO paper (https://huggingface.co/papers/2408.06266) + # Use this loss when you believe the chosen outputs are better than your model's default output + chosen_losses = 1 - F.sigmoid(self.beta * chosen_logratios) + + chosen_rewards = self.beta * chosen_logratios.detach() + + else: + # lists can't be empty -- if they are, then accelerate.gather will hang + chosen_losses = torch.Tensor([]).to(self.accelerator.device) + chosen_rewards = torch.Tensor([]).to(self.accelerator.device) + + # Rejected losses + if policy_rejected_logps.shape[0] != 0 or reference_rejected_logps.shape[0] != 0: + rejected_logratios = policy_rejected_logps - reference_rejected_logps + + if self.loss_type == "kto": + rejected_losses = 1 - F.sigmoid(self.beta * (kl - rejected_logratios)) + elif self.loss_type == "apo_zero_unpaired": + rejected_losses = F.sigmoid(self.beta * rejected_logratios) + + rejected_rewards = self.beta * rejected_logratios.detach() + else: + # lists can't be empty -- if they are, then accelerate.gather will hang + rejected_losses = torch.Tensor([]).to(self.accelerator.device) + rejected_rewards = torch.Tensor([]).to(self.accelerator.device) + + losses = torch.cat( + (self.desirable_weight * chosen_losses, self.undesirable_weight * rejected_losses), + 0, + ) + + return losses, chosen_rewards, rejected_rewards, kl + + def _compute_kl_logps(self, model, batch): + """Compute KL log probabilities for a given batch.""" + KL_logps = None + if self.calculate_KL: + if self.is_encoder_decoder: + KL_model_kwargs = { + "input_ids": batch["KL_prompt_input_ids"], + "attention_mask": batch["KL_prompt_attention_mask"], + "labels": batch["KL_completion_labels"], + "decoder_input_ids": batch.get("KL_completion_decoder_input_ids"), + } + else: + KL_model_kwargs = { + "input_ids": batch["KL_completion_input_ids"], + "attention_mask": batch["KL_completion_attention_mask"], + } + + with torch.no_grad(): + KL_logits = model(**KL_model_kwargs).logits + + KL_logps = self.get_batch_logps( + KL_logits, + batch["KL_completion_labels"], + average_log_prob=False, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + return KL_logps + + def _compute_loss_liger(self, model, batch): + """ + Compute the KTO loss using the Liger-Kernel's LigerFusedLinearKTOLoss. + + Args: + model: + The policy model used for generating log probabilities and outputs. It could be an encoder-decoder + model or a regular language model. + batch: A dictionary containing the input data and labels for the batch. + + Returns: + A dictionary containing the following keys: + - "loss": The computed KTO loss for the batch. + - "chosen_logits_sum": Sum of the logits for the chosen responses from the policy model. + - "rejected_logits_sum": Sum of the logits for the rejected responses from the policy model. + - "chosen_logps": Log probabilities of the chosen responses from the policy model. + - "rejected_logps": Log probabilities of the rejected responses from the policy model. + - "chosen_rewards": Rewards for the chosen responses. + - "rejected_rewards": Rewards for the rejected responses. + - "kl": The KL divergence between the policy and reference models (detached). + + If auxiliary loss is enabled, the dictionary will also include: + - "aux_loss": The auxiliary loss from the model outputs. + """ + policy_KL_logps = self._compute_kl_logps(model, batch) + reference_KL_logps = self._compute_kl_logps(self.ref_model, batch) + if self.calculate_KL: + kl = (policy_KL_logps - reference_KL_logps).mean().detach() + kl = self.accelerator.gather_for_metrics(kl).mean().clamp(min=0) + else: + kl = torch.zeros(1).to(self.accelerator.device) + + model_kwargs = ( + { + "labels": batch["completion_labels"], + "decoder_input_ids": batch.get("completion_decoder_input_ids"), + } + if self.is_encoder_decoder + else {} + ) + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + if self.is_encoder_decoder: + # 1. Get encoder outputs + encoder_outputs = model.get_encoder()( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + return_dict=True, + **model_kwargs, + ) + # 2. Get decoder outputs + outputs = model.get_decoder()( + input_ids=model_kwargs["decoder_input_ids"], + encoder_hidden_states=encoder_outputs.last_hidden_state, + use_cache=False, + **model_kwargs, + ) + # 1. Get reference encoder outputs + ref_encoder_outputs = self.ref_model.get_encoder()( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + return_dict=True, + **model_kwargs, + ) + # 2. Get reference decoder outputs + ref_outputs = self.ref_model.get_decoder()( + input_ids=model_kwargs["decoder_input_ids"], + encoder_hidden_states=ref_encoder_outputs.last_hidden_state, + use_cache=False, + **model_kwargs, + ) + else: + # skip the lm head and get the last hidden state + if hasattr(model, "get_decoder"): + base_model = model.get_decoder() + else: + base_model = getattr(model, self.args.base_model_attribute_name) + outputs = base_model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + use_cache=False, + **model_kwargs, + ) + + # reference model + if hasattr(self.ref_model, "get_decoder"): + ref_base_model = self.ref_model.get_decoder() + else: + ref_base_model = getattr(self.ref_model, self.args.base_model_attribute_name) + ref_outputs = ref_base_model( + batch["completion_input_ids"], + attention_mask=batch["completion_attention_mask"], + use_cache=False, + **model_kwargs, + ) + lm_head = model.get_output_embeddings() + ref_lm_head = self.ref_model.get_output_embeddings() + + ( + loss, + ( + chosen_logps_sum, + rejected_logps_sum, + chosen_logits_sum, + rejected_logits_sum, + chosen_rewards_sum, + rejected_rewards_sum, + ), + ) = self.kto_loss_fn( + _input=outputs.last_hidden_state[:, :-1] if not self.is_encoder_decoder else outputs.last_hidden_state, + lin_weight=lm_head.weight, + target=batch["completion_labels"][:, 1:], + bias=lm_head.bias if hasattr(lm_head, "bias") else None, + preference_labels=torch.tensor(batch["label"], dtype=torch.bool).to(self.accelerator.device), + ref_input=ref_outputs.last_hidden_state[:, :-1] + if not self.is_encoder_decoder + else outputs.last_hidden_state, + ref_weight=ref_lm_head.weight, + ref_bias=ref_lm_head.bias if hasattr(lm_head, "bias") else None, + kl=kl, + ) + + output = { + "loss": loss, + "chosen_logits_sum": chosen_logits_sum, + "rejected_logits_sum": rejected_logits_sum, + "chosen_logps_sum": chosen_logps_sum, + "rejected_logps_sum": rejected_logps_sum, + "chosen_rewards_sum": chosen_rewards_sum, + "rejected_rewards_sum": rejected_rewards_sum, + "kl": kl, + } + if self.aux_loss_enabled: + output["aux_loss"] = outputs.aux_loss + + return output + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, Union[list, torch.LongTensor]], + ): + """Compute the KTO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + batch = {k: (v.to(self.accelerator.device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()} + + labels = torch.tensor(batch["label"]) + num_chosen = labels.sum().to(self.accelerator.device) + num_rejected = (len(labels) - num_chosen).to(self.accelerator.device) + + if self.args.use_liger_loss: + model_output = self._compute_loss_liger(model, batch) + losses = model_output["loss"] + policy_chosen_logits = model_output["chosen_logits_sum"] + policy_rejected_logits = model_output["rejected_logits_sum"] + policy_chosen_logps = model_output["chosen_logps_sum"] + policy_rejected_logps = model_output["rejected_logps_sum"] + chosen_rewards = model_output["chosen_rewards_sum"] + rejected_rewards = model_output["rejected_rewards_sum"] + kl = model_output["kl"] + if self.aux_loss_enabled: + aux_loss = model_output["aux_loss"] + else: + forward_output = self.forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_KL_logps, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + # if reference_logps in batch use them, otherwise use the reference model + if "reference_logps" in batch: + chosen_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is True] + rejected_idx = [i for i in range(batch["reference_logps"].shape[0]) if batch["label"][i] is False] + + reference_chosen_logps = batch["reference_logps"][chosen_idx, ...] + reference_rejected_logps = batch["reference_logps"][rejected_idx, ...] + if self.calculate_KL: + reference_KL_logps = batch["reference_KL_logps"] + else: + reference_KL_logps = None + else: + with torch.no_grad(): + if self.ref_model is None: + with self.null_ref_context(): + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + reference_KL_logps, + ) = self.forward(self.model, batch)[:5] + else: + ( + reference_chosen_logps, + reference_rejected_logps, + _, + _, + reference_KL_logps, + ) = self.forward(self.ref_model, batch)[:5] + + losses, chosen_rewards, rejected_rewards, kl = self.kto_loss( + policy_chosen_logps, + policy_rejected_logps, + policy_KL_logps, + reference_chosen_logps, + reference_rejected_logps, + reference_KL_logps, + ) + + metrics["kl"] = kl.item() + + all_num_chosen = self.accelerator.gather_for_metrics(num_chosen).sum().item() + all_num_rejected = self.accelerator.gather_for_metrics(num_rejected).sum().item() + + if all_num_chosen > 0: + metrics["rewards/chosen_sum"] = ( + self.accelerator.gather_for_metrics(chosen_rewards.nansum()).nansum().item() + ) + metrics["logps/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logps.nansum()).nansum().item() + ) + metrics["logits/chosen_sum"] = ( + self.accelerator.gather_for_metrics(policy_chosen_logits.nansum()).nansum().item() + ) + metrics["count/chosen"] = all_num_chosen + + if all_num_rejected > 0: + metrics["rewards/rejected_sum"] = ( + self.accelerator.gather_for_metrics(rejected_rewards.nansum()).nansum().item() + ) + metrics["logps/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logps.nansum()).nansum().item() + ) + metrics["logits/rejected_sum"] = ( + self.accelerator.gather_for_metrics(policy_rejected_logits.nansum()).nansum().item() + ) + metrics["count/rejected"] = all_num_rejected + + loss = losses.nanmean() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, metrics + + def compute_loss( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + return_outputs=False, + num_items_in_batch=None, + ) -> Union[torch.Tensor, tuple[torch.Tensor, dict[str, torch.Tensor]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs) + + # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: + loss = loss.to(self.args.device) + # force log the metrics + if self.accelerator.is_main_process: + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def _get_train_sampler(self, dataset: Optional[Dataset] = None) -> Optional[torch.utils.data.Sampler]: + if dataset is None: + dataset = self.train_dataset + if dataset is None or not has_length(dataset): + return None + return SequentialSampler(dataset) + + def generate_from_model_and_ref(self, model, batch: dict[str, torch.LongTensor]) -> tuple[str, str]: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + # if reference_output in batch use that otherwise use the reference model + if "reference_output" in batch: + reference_output = batch["reference_output"] + else: + if self.ref_model is None: + with self.null_ref_context(): + reference_output = self.model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + else: + reference_output = self.ref_model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + reference_output = pad_to_length(reference_output, self.max_length, self.processing_class.pad_token_id) + reference_output_decoded = self.processing_class.batch_decode(reference_output, skip_special_tokens=True) + + return policy_output_decoded, reference_output_decoded + + def prediction_step( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[list[str]] = None, + ): + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs) + + # force log the metrics + if self.accelerator.is_main_process: + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return (loss.detach(), None, None) + + # logits for the chosen and rejected samples from model + logits_dict = {} + if "logits/chosen_sum" in metrics: + logits_dict["eval_logits/chosen"] = metrics["logits/chosen_sum"] + if "logits/rejected_sum" in metrics: + logits_dict["eval_logits/rejected"] = metrics["logits/rejected_sum"] + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: Optional[bool] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + target_indicies = [i for i in range(len(random_batch["label"])) if random_batch["label"][i] is False] + target_batch = { + "prompt_input_ids": random_batch["prompt_input_ids"][target_indicies], + "prompt_attention_mask": random_batch["prompt_attention_mask"][target_indicies], + "prompt": itemgetter(*target_indicies)(random_batch["prompt"]), + } + policy_output_decoded, ref_output_decoded = self.generate_from_model_and_ref(self.model, target_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy", "Ref Model"], + data=[ + [prompt, pol[len(prompt) :], ref[len(prompt) :]] + for prompt, pol, ref in zip(target_batch["prompt"], policy_output_decoded, ref_output_decoded) + ], + ) + if "wandb" in self.args.report_to: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float` or `None`, *optional*, defaults to `None`): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # train metrics should have no prefix, eval should have 'eval_' + prefix = "eval_" if train_eval == "eval" else "" + # accumulate average metrics from sums and lengths + for split in ["chosen", "rejected"]: + if f"count/{split}" in self._stored_metrics[train_eval]: + count_sum = torch.Tensor(self._stored_metrics[train_eval][f"count/{split}"]).sum().item() + for metric in ["rewards", "logps", "logits"]: + logs[f"{prefix}{metric}/{split}"] = ( + torch.Tensor(self._stored_metrics[train_eval][f"{metric}/{split}_sum"]).sum().item() + / count_sum + ) + # delete obsolete metric + del self._stored_metrics[train_eval][f"{metric}/{split}_sum"] + del self._stored_metrics[train_eval][f"count/{split}"] + # calculate reward margin + if f"{prefix}rewards/chosen" in logs and f"{prefix}rewards/rejected" in logs: + logs[f"{prefix}rewards/margins"] = logs[f"{prefix}rewards/chosen"] - logs[f"{prefix}rewards/rejected"] + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[f"{prefix}{key}"] = torch.Tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{ethayarajh2024kto, + title = {{KTO: Model Alignment as Prospect Theoretic Optimization}}, + author = {Kawin Ethayarajh and Winnie Xu and Niklas Muennighoff and Dan Jurafsky and Douwe Kiela}, + year = 2024, + eprint = {arXiv:2402.01306}, + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="KTO", + trainer_citation=citation, + paper_title="KTO: Model Alignment as Prospect Theoretic Optimization", + paper_id="2402.01306", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/model_config.py b/lib/python3.12/site-packages/trl/trainer/model_config.py new file mode 100644 index 0000000000000000000000000000000000000000..935bd73526dbdd8e91e0f18577903ee6319115ec --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/model_config.py @@ -0,0 +1,179 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Optional + + +@dataclass +class ModelConfig: + """ + Configuration class for the models. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + model_name_or_path (`str` or `None`, *optional*, defaults to `None`): + Model checkpoint for weights initialization. + model_revision (`str`, *optional*, defaults to `"main"`): + Specific model version to use. It can be a branch name, a tag name, or a commit id. + torch_dtype (`Literal["auto", "bfloat16", "float16", "float32"]` or `None`, *optional*, defaults to `None`): + Override the default `torch.dtype` and load the model under this dtype. Possible values are + + - `"bfloat16"`: `torch.bfloat16` + - `"float16"`: `torch.float16` + - `"float32"`: `torch.float32` + - `"auto"`: Automatically derive the dtype from the model's weights. + + trust_remote_code (`bool`, *optional*, defaults to `False`): + Whether to allow for custom models defined on the Hub in their own modeling files. This option should only + be set to `True` for repositories you trust and in which you have read the code, as it will execute code + present on the Hub on your local machine. + attn_implementation (`str` or `None`, *optional*, defaults to `None`): + Which attention implementation to use. You can run `--attn_implementation=flash_attention_2`, in which case + you must install this manually by running `pip install flash-attn --no-build-isolation`. + use_peft (`bool`, *optional*, defaults to `False`): + Whether to use PEFT for training. + lora_r (`int`, *optional*, defaults to `16`): + LoRA R value. + lora_alpha (`int`, *optional*, defaults to `32`): + LoRA alpha. + lora_dropout (`float`, *optional*, defaults to `0.05`): + LoRA dropout. + lora_target_modules (`Union[str, list[str]]` or `None`, *optional*, defaults to `None`): + LoRA target modules. + lora_modules_to_save (`list[str]` or `None`, *optional*, defaults to `None`): + Model layers to unfreeze & train. + lora_task_type (`str`, *optional*, defaults to `"CAUSAL_LM"`): + Task type to pass for LoRA (use `"SEQ_CLS"` for reward modeling). + use_rslora (`bool`, *optional*, defaults to `False`): + Whether to use Rank-Stabilized LoRA, which sets the adapter scaling factor to `lora_alpha/√r`, instead of + the original default value of `lora_alpha/r`. + use_dora (`bool`, *optional*, defaults to `False`): + Enable [Weight-Decomposed Low-Rank Adaptation (DoRA)](https://huggingface.co/papers/2402.09353). This + technique decomposes the updates of the weights into two parts, magnitude and direction. Direction is + handled by normal LoRA, whereas the magnitude is handled by a separate learnable parameter. This can + improve the performance of LoRA, especially at low ranks. Right now, DoRA only supports linear and Conv2D + layers. DoRA introduces a bigger overhead than pure LoRA, so it is recommended to merge weights for + inference. + load_in_8bit (`bool`, *optional*, defaults to `False`): + Whether to use 8 bit precision for the base model. Works only with LoRA. + load_in_4bit (`bool`, *optional*, defaults to `False`): + Whether to use 4 bit precision for the base model. Works only with LoRA. + bnb_4bit_quant_type (`str`, *optional*, defaults to `"nf4"`): + Quantization type (`"fp4"` or `"nf4"`). + use_bnb_nested_quant (`bool`, *optional*, defaults to `False`): + Whether to use nested quantization. + """ + + model_name_or_path: Optional[str] = field( + default=None, + metadata={"help": "Model checkpoint for weights initialization."}, + ) + model_revision: str = field( + default="main", + metadata={"help": "Specific model version to use. It can be a branch name, a tag name, or a commit id."}, + ) + torch_dtype: Optional[str] = field( + default=None, + metadata={ + "help": "Override the default `torch.dtype` and load the model under this dtype.", + "choices": ["auto", "bfloat16", "float16", "float32"], + }, + ) + trust_remote_code: bool = field( + default=False, + metadata={ + "help": "Whether to allow for custom models defined on the Hub in their own modeling files. This option " + "should only be set to `True` for repositories you trust and in which you have read the code, as it will " + "execute code present on the Hub on your local machine." + }, + ) + attn_implementation: Optional[str] = field( + default=None, + metadata={ + "help": "Which attention implementation to use. You can run `--attn_implementation=flash_attention_2`, in " + "which case you must install this manually by running `pip install flash-attn --no-build-isolation`." + }, + ) + use_peft: bool = field( + default=False, + metadata={"help": "Whether to use PEFT for training."}, + ) + lora_r: int = field( + default=16, + metadata={"help": "LoRA R value."}, + ) + lora_alpha: int = field( + default=32, + metadata={"help": "LoRA alpha."}, + ) + lora_dropout: float = field( + default=0.05, + metadata={"help": "LoRA dropout."}, + ) + lora_target_modules: Optional[list[str]] = field( + default=None, + metadata={"help": "LoRA target modules."}, + ) + lora_modules_to_save: Optional[list[str]] = field( + default=None, + metadata={"help": "Model layers to unfreeze & train."}, + ) + lora_task_type: str = field( + default="CAUSAL_LM", + metadata={"help": "Task type to pass for LoRA (use 'SEQ_CLS' for reward modeling)."}, + ) + use_rslora: bool = field( + default=False, + metadata={ + "help": "Whether to use Rank-Stabilized LoRA, which sets the adapter scaling factor to `lora_alpha/√r`, " + "instead of the original default value of `lora_alpha/r`." + }, + ) + use_dora: bool = field( + default=False, + metadata={ + "help": "Enable Weight-Decomposed Low-Rank Adaptation (DoRA). This technique decomposes the updates of " + "the weights into two parts, magnitude and direction. Direction is handled by normal LoRA, whereas the " + "magnitude is handled by a separate learnable parameter. This can improve the performance of LoRA, " + "especially at low ranks. Right now, DoRA only supports linear and Conv2D layers. DoRA introduces a " + "bigger overhead than pure LoRA, so it is recommended to merge weights for inference." + }, + ) + load_in_8bit: bool = field( + default=False, + metadata={"help": "Whether to use 8 bit precision for the base model. Works only with LoRA."}, + ) + load_in_4bit: bool = field( + default=False, + metadata={"help": "Whether to use 4 bit precision for the base model. Works only with LoRA."}, + ) + bnb_4bit_quant_type: str = field( + default="nf4", + metadata={"help": "Quantization type.", "choices": ["fp4", "nf4"]}, + ) + use_bnb_nested_quant: bool = field( + default=False, + metadata={"help": "Whether to use nested quantization."}, + ) + + def __post_init__(self): + if self.load_in_8bit and self.load_in_4bit: + raise ValueError("You can't use 8 bit and 4 bit precision at the same time") + + if hasattr(self.lora_target_modules, "__len__") and len(self.lora_target_modules) == 1: + self.lora_target_modules = self.lora_target_modules[0] diff --git a/lib/python3.12/site-packages/trl/trainer/nash_md_config.py b/lib/python3.12/site-packages/trl/trainer/nash_md_config.py new file mode 100644 index 0000000000000000000000000000000000000000..07d8152f4fae45c9cf16b1815b3a225a9237c695 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/nash_md_config.py @@ -0,0 +1,46 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field + +from trl.trainer.online_dpo_config import OnlineDPOConfig + + +@dataclass +class NashMDConfig(OnlineDPOConfig): + r""" + Configuration class for the [`NashMDTrainer`]. + + Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following: + + Parameters: + mixture_coef (`float` or `list[float]`, *optional*, defaults to `0.5`): + Logit mixture coefficient for the model and reference model. If a list of floats is provided then the + mixture coefficient is selected for each new epoch and the last coefficient is used for the rest of the + epochs. + """ + + mixture_coef: list[float] = field( + default_factory=lambda: [0.5], + metadata={ + "help": "Logit mixture coefficient for the model and reference model. If a list of floats is provided " + "then the mixture coefficient is selected for each new epoch and the last coefficient is used for the " + "rest of the epochs." + }, + ) + + def __post_init__(self): + super().__post_init__() + if hasattr(self.mixture_coef, "__len__") and len(self.mixture_coef) == 1: + self.mixture_coef = self.mixture_coef[0] diff --git a/lib/python3.12/site-packages/trl/trainer/nash_md_trainer.py b/lib/python3.12/site-packages/trl/trainer/nash_md_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..b553ff9dda64fa5cc04d62235ee262f0e769dc45 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/nash_md_trainer.py @@ -0,0 +1,551 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import textwrap +from typing import Any, Callable, Optional, Union + +import jinja2 +import torch +import torch.nn as nn +import torch.nn.functional as F +from datasets import Dataset, IterableDataset +from transformers import ( + BaseImageProcessor, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, + is_wandb_available, +) +from transformers.trainer_utils import EvalPrediction +from transformers.training_args import OptimizerNames +from transformers.utils import is_apex_available, is_peft_available + +from ..data_utils import is_conversational, maybe_apply_chat_template +from ..models.modeling_base import GeometricMixtureWrapper +from ..models.utils import unwrap_model_for_generation +from .judges import BasePairwiseJudge +from .nash_md_config import NashMDConfig +from .online_dpo_trainer import OnlineDPOTrainer +from .utils import ( + SIMPLE_CHAT_TEMPLATE, + empty_cache, + generate_model_card, + get_comet_experiment_url, + get_reward, + selective_log_softmax, + truncate_right, +) + + +if is_apex_available(): + from apex import amp + + +if is_wandb_available(): + import wandb + + +if is_peft_available(): + from peft import PeftModel + + +class NashMDTrainer(OnlineDPOTrainer): + r""" + Initialize NashMDTrainer as a subclass of [`OnlineDPOConfig`]. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForCausalLM`. + ref_model (`PreTrainedModelWrapper`): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + reward_model (`transformers.PreTrainedModel`): + The reward model to score completions with, preferably an `AutoModelForSequenceClassification`. + judge (`BasePairwiseJudge`): + The judge to use for pairwise comparison of model completions. + args (`NashMDConfig`): + The NashMD config arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`DPODataCollatorWithPadding`) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + peft_config (`dict`): + The peft config to use for training. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + """ + + _tag_names = ["trl", "nash-md"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module] = None, + ref_model: Union[PreTrainedModel, nn.Module] = None, + reward_model: Union[PreTrainedModel, nn.Module, None] = None, + judge: Optional[BasePairwiseJudge] = None, + args: Optional[NashMDConfig] = None, + data_collator: Optional[Callable] = None, + train_dataset: Optional[Union[Dataset, IterableDataset]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + peft_config: Optional[dict] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + ) -> None: + super().__init__( + model=model, + ref_model=ref_model, + reward_model=reward_model, + judge=judge, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + reward_processing_class=processing_class, # for now, NashMDTrainer can't use any reward model + peft_config=peft_config, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + self._mixture_coef = self.args.mixture_coef + + # Overwrite the stats dictionary to include NashMD specific statistics + self.stats = { + # Remove "non_score_reward", "rlhf_reward", "scores_margin" + # Add "mixture_coef" + "loss/kl": [], + "objective/entropy": [], + "loss/score": [], + "rewards/probabilities": [], + "rewards/accuracies": [], + "rewards/margins": [], + "logps/chosen": [], + "logps/rejected": [], + "val/model_contain_eos_token": [], + "val/ref_contain_eos_token": [], + "beta": [], + "mixture_coef": [], + } + if self.reward_model is not None: + self.stats["rewards/chosen"] = [] + self.stats["rewards/rejected"] = [] + + @property + def mixture_coef(self): + if isinstance(self._mixture_coef, list): + epoch = self.state.epoch + return self._mixture_coef[epoch] if epoch < len(self._mixture_coef) else self._mixture_coef[-1] + else: + return self._mixture_coef + + def _generate_completions(self, model, prompts): + # Generate completions from the policy model. + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_policy_for_gen_ctx: + model_output = unwrapped_policy_for_gen_ctx.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + # Get the DDP/FSDP unwrapped version of the main model. + # This will be the policy model for GeometricMixtureWrapper (PEFT adapters active if PEFT is used). + policy_model_for_gmw = self.accelerator.unwrap_model(model) + + # Determine the correct reference model for GeometricMixtureWrapper. + # This also needs to be DDP/FSDP unwrapped. + ref_model_for_gmw: torch.nn.Module + if self.ref_model is None: + # No explicit ref_model is provided. + # Use the base of the main `model` if it's a PEFT model. + # policy_model_for_gmw is already DDP-unwrapped. + if is_peft_available() and isinstance(policy_model_for_gmw, PeftModel): + ref_model_for_gmw = policy_model_for_gmw.get_base_model() + else: + # Not a PEFT model (or PEFT not available), or already a base model. + # Use the DDP-unwrapped policy model itself as the reference. + ref_model_for_gmw = policy_model_for_gmw + else: + # An explicit ref_model is provided. Unwrap it for DDP/FSDP. + ref_model_for_gmw = self.accelerator.unwrap_model(self.ref_model) + + # Both models given to GeometricMixtureWrapper (policy_model_for_gmw and ref_model_for_gmw) are DDP-unwrapped. + with torch.no_grad(): # Ensure no_grad context for mixture model generation + mixture_model = GeometricMixtureWrapper( + model=policy_model_for_gmw, + ref_model=ref_model_for_gmw, + generation_config=self.generation_config, + mixture_coef=self.mixture_coef, + device=self.accelerator.device, + ) + + mixture_output = mixture_model.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + return model_output, mixture_output + + def _process_completions(self, model_output, mixture_output, prompts): + context_length = prompts["input_ids"].shape[1] + + # Process model completions + model_completion_ids = model_output[:, context_length:] + model_completion_ids, model_completion_mask = truncate_right( + model_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + model_data = { + "input_ids": torch.cat((prompts["input_ids"], model_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], model_completion_mask), dim=1), + "raw": prompts["raw"], + } + + # Process reference model completions + mixture_completion_ids = mixture_output[:, context_length:] + mixture_completion_ids, mixture_completion_mask = truncate_right( + mixture_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + mixture_data = { + "input_ids": torch.cat((prompts["input_ids"], mixture_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], mixture_completion_mask), dim=1), + "raw": prompts["raw"], + } + + return model_data, mixture_data + + def _compute_rewards(self, model_data, mixture_data, context_length): + with torch.no_grad(): + _, model_scores, _ = get_reward( + self.reward_model, model_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + _, mixture_scores, _ = get_reward( + self.reward_model, mixture_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + + # Apply EOS penalty if needed + if self.args.missing_eos_penalty is not None: + model_contain_eos = torch.any(model_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + mixture_contain_eos = torch.any(mixture_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + model_scores[~model_contain_eos] -= self.args.missing_eos_penalty + mixture_scores[~mixture_contain_eos] -= self.args.missing_eos_penalty + + return model_scores, mixture_scores + + def _compute_judge(self, model_data, mixture_data, context_length): + prompts = model_data["raw"] + model_data_completions = self.processing_class.batch_decode( + model_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + model_data_completions = [completion.strip() for completion in model_data_completions] + + mixture_data_completions = self.processing_class.batch_decode( + mixture_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + mixture_data_completions = [completion.strip() for completion in mixture_data_completions] + if is_conversational({"prompt": prompts[0]}): + model_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in model_data_completions + ] + environment = jinja2.Environment() + template = environment.from_string(SIMPLE_CHAT_TEMPLATE) + prompts = [template.render(messages=message) for message in prompts] + model_data_completions = [template.render(messages=completion) for completion in model_data_completions] + + mixture_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in mixture_data_completions + ] + mixture_data_completions = [ + template.render(messages=completion) for completion in mixture_data_completions + ] + + probability = self.judge.judge( + prompts, + list(zip(model_data_completions, mixture_data_completions)), + return_scores=True, + ) + return torch.tensor(probability, device=model_data["input_ids"].device) + + def _compute_logprobs(self, model, model_data, context_length): + def compute_logprobs_for_data(m, data): + output = m(data["input_ids"], attention_mask=data["attention_mask"]) + logits = output.logits[:, context_length - 1 : -1] + token_logprobs = selective_log_softmax(logits, data["input_ids"][:, context_length:]) + return token_logprobs + + # Compute logprobs for model completions under the model + model_logprobs_model_data = compute_logprobs_for_data(model, model_data) + + # Compute logprobs of model completions under the reference model + with torch.no_grad(): + if self.ref_model is None: + with model.disable_adapter(): + ref_logprobs_model_data = compute_logprobs_for_data(model, model_data) + else: + ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data) + + # Mask padding tokens + model_padding_mask = model_data["attention_mask"][:, context_length:] == 0 + model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + + return (model_logprobs_model_data, ref_logprobs_model_data) + + def _compute_losses( + self, + model_logprobs_model_data, + ref_logprobs_model_data, + probability, + ): + # reinforce score where 0.5 is a control variate + score = (probability - 0.5) * model_logprobs_model_data.sum(1) + + # kl divergence via reinforce + with torch.no_grad(): + log_ratio = model_logprobs_model_data - ref_logprobs_model_data + kl_div_log = log_ratio.sum(1) + kl_div_loss = (log_ratio * model_logprobs_model_data).sum(1) + + # final loss + loss = self.beta * kl_div_loss - score + + return loss.mean(), score, kl_div_log + + def _log_statistics( + self, + model_data, + mixture_data, + model_logprobs_model_data, + ref_logprobs_model_data, + probability, + score, + kl_div, + context_length, + model_scores=None, + mixture_scores=None, + ): + # Helper function to gather and compute mean + def gather_mean(tensor): + return self.accelerator.gather_for_metrics(tensor).mean().item() + + # Log score + self.stats["loss/score"].append(gather_mean(score)) + # Log KL divergence + self.stats["loss/kl"].append(gather_mean(kl_div)) + + # Log logprobs + model_logprobs_model_data_sum = model_logprobs_model_data.sum(1) + ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1) + + self.stats["logps/chosen"].append(gather_mean(model_logprobs_model_data_sum)) + self.stats["logps/rejected"].append(gather_mean(ref_logprobs_model_data_sum)) + + # Log rewards + if self.reward_model is not None: + self.stats["rewards/chosen"].append(gather_mean(model_scores)) + self.stats["rewards/rejected"].append(gather_mean(mixture_scores)) + + # Log probabilities + self.stats["rewards/probabilities"].append(gather_mean(probability)) + + # Calculate entropy for model data + entropy_model_data = -model_logprobs_model_data.sum(1) + self.stats["objective/entropy"].append(gather_mean(entropy_model_data)) + + # Calculate margins + margin = model_logprobs_model_data_sum - ref_logprobs_model_data_sum + self.stats["rewards/margins"].append(gather_mean(margin)) + + # Calculate accuracy + accuracy = (margin > 0).float() + self.stats["rewards/accuracies"].append(gather_mean(accuracy)) + + # Log EOS token statistics + model_eos = (model_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + mixture_eos = (mixture_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + self.stats["val/model_contain_eos_token"].append(gather_mean(model_eos.float())) + self.stats["val/ref_contain_eos_token"].append(gather_mean(mixture_eos.float())) + + # Log beta and mixture coef + self.stats["beta"].append(self.beta) + self.stats["mixture_coef"].append(self.mixture_coef) + + def training_step( + self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None + ) -> torch.Tensor: + model.train() + + # Apply chat template and tokenize the input + batch_size = len(next(iter(inputs.values()))) + prompts = inputs["prompt"] + inputs = [{k: v[i] for k, v in inputs.items()} for i in range(batch_size)] + inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs] + inputs = [self.tokenize_row(x, self.model.config.is_encoder_decoder, self.processing_class) for x in inputs] + inputs = self.data_collator(inputs) + + # need the prompt_ only + inputs = self._prepare_inputs(inputs) + context_length = inputs["prompt_input_ids"].shape[1] + prompts = { + "input_ids": inputs["prompt_input_ids"], + "attention_mask": inputs["prompt_attention_mask"], + "raw": prompts, + } + del inputs + + # Sample completions from both the model and the reference model + model_output, mixture_output = self._generate_completions(model, prompts) + + # Process model completions + model_data, mixture_data = self._process_completions(model_output, mixture_output, prompts) + + # Compute rewards + if self.reward_model is not None: + model_scores, mixture_scores = self._compute_rewards(model_data, mixture_data, context_length) + # probability of the model data vs the mixture data + probability = F.sigmoid(model_scores - mixture_scores) + else: + model_scores, mixture_scores = None, None + probability = self._compute_judge(model_data, mixture_data, context_length) + + # Compute logprobs + model_logprobs_model_data, ref_logprobs_model_data = self._compute_logprobs(model, model_data, context_length) + + # Compute loss + loss, score, kl_div = self._compute_losses(model_logprobs_model_data, ref_logprobs_model_data, probability) + + # Log everything + self._log_statistics( + model_data, + mixture_data, + model_logprobs_model_data.detach(), + ref_logprobs_model_data, + probability, + score.detach(), + kl_div.detach(), + context_length, + model_scores, + mixture_scores, + ) + + if ( + self.args.torch_empty_cache_steps is not None + and self.state.global_step % self.args.torch_empty_cache_steps == 0 + ): + empty_cache() + + kwargs = {} + # For LOMO optimizers you need to explicitly use the learning rate + if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + kwargs["learning_rate"] = self._get_learning_rate() + + if self.args.n_gpu > 1: + loss = loss.mean() # mean() to average on multi-gpu parallel training + + if self.use_apex: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: + scaled_loss.backward() + else: + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{munos2024nash, + title = {{Nash Learning from Human Feedback}}, + author = {R{\'{e}}mi Munos and Michal Valko and Daniele Calandriello and Mohammad Gheshlaghi Azar and Mark Rowland and Zhaohan Daniel Guo and Yunhao Tang and Matthieu Geist and Thomas Mesnard and C{\\^{o}}me Fiegel and Andrea Michi and Marco Selvi and Sertan Girgin and Nikola Momchev and Olivier Bachem and Daniel J. Mankowitz and Doina Precup and Bilal Piot}, + year = 2024, + booktitle = {Forty-first International Conference on Machine Learning, {ICML} 2024, Vienna, Austria, July 21-27, 2024}, + publisher = {OpenReview.net}, + url = {https://openreview.net/forum?id=Y5AmNYiyCQ} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="Nash-MD", + trainer_citation=citation, + paper_title="Nash Learning from Human Feedback", + paper_id="2312.00886", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/online_dpo_config.py b/lib/python3.12/site-packages/trl/trainer/online_dpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..8578a879e7031bbc467268162f9ba0356d7979db --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/online_dpo_config.py @@ -0,0 +1,185 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Optional + +from transformers import TrainingArguments + + +@dataclass +class OnlineDPOConfig(TrainingArguments): + r""" + Configuration class for the [`OnlineDPOTrainer`]. + + This class includes only the parameters that are specific to Online DPO training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + reward_model_path (`str` or `None`, *optional*, defaults to `None`): + Path to the reward model. Either `judge` or `reward_model_path` must be set, but not both. + judge (`str` or `None`, *optional*, defaults to `None`): + Name of the judge to use. Either `judge` or `reward_model_path` must be set, but not both. + max_new_tokens (`int`, *optional*, defaults to `64`): + Maximum number of tokens to generate per completion. + max_length (`int`, *optional*, defaults to `256`): + Maximum total length of the sequence (prompt + completion) used to compute log probabilities. If the + sequence exceeds this limit, the leftmost tokens will be truncated to preserve as much of the completion as + possible. + temperature (`float`, *optional*, defaults to `0.9`): + Temperature for sampling. The higher the temperature, the more random the completions. + missing_eos_penalty (`float` or `None`, *optional*, defaults to `None`): + Penalty applied to the score when the model fails to generate an EOS token. This is useful to encourage to + generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be a positive + value. + beta (`float` or `list[float]`, *optional*, defaults to `0.1`): + Parameter controlling the deviation from the reference model. Higher β means less deviation from the + reference model. For the IPO loss (`loss_type="ipo"`), β is the regularization parameter denoted by τ in + the [paper](https://huggingface.co/papers/2310.12036). If a list of floats is provided then the β is + selected for each new epoch and the last β is used for the rest of the epochs. + loss_type (`str`, *optional*, defaults to `"sigmoid"`): + Type of loss to use. Possible values are: + + - `"sigmoid"`: sigmoid loss from the original [DPO](https://huggingface.co/papers/2305.18290) paper. + - `"ipo"`: IPO loss from the [IPO](https://huggingface.co/papers/2310.12036) paper. + + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model and reference model. + use_vllm (`bool`, *optional*, defaults to `False`): + Whether to use vLLM for generating completions. Requires vLLM to be installed (`pip install vllm`). + gpu_memory_utilization (`float`, *optional*, defaults to `0.55`): + The vLLM memory utilization. The default value is 0.55. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. + """ + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=5e-7, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + reward_model_path: Optional[str] = field( + default=None, + metadata={ + "help": "Path to the reward model. Either `judge` or `reward_model_path` must be set, but not both." + }, + ) + judge: Optional[str] = field( + default=None, + metadata={ + "help": "Name of the judge to use. Either `judge` or `reward_model_path` must be set, but not both." + }, + ) + max_new_tokens: int = field( + default=64, + metadata={"help": "Maximum number of tokens to generate per completion."}, + ) + max_length: int = field( + default=512, + metadata={ + "help": "Maximum total length of the sequence (prompt + completion) used to compute log probabilities. If " + "the sequence exceeds this limit, the leftmost tokens will be truncated to preserve as much of the " + "completion as possible." + }, + ) + temperature: float = field( + default=0.9, + metadata={"help": "Temperature for sampling. The higher the temperature, the more random the completions."}, + ) + missing_eos_penalty: Optional[float] = field( + default=None, + metadata={ + "help": "Penalty applied to the score when the model fails to generate an EOS token. This is useful to " + "encourage to generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be " + "a positive value." + }, + ) + beta: list[float] = field( + default_factory=lambda: [0.1], + metadata={ + "help": "Parameter controlling the deviation from the reference model. Higher β means less deviation from " + "the reference model. For the IPO loss (`loss_type='ipo'`), β is the regularization parameter denoted by " + "τ in the [paper](https://huggingface.co/papers/2310.12036). If a list of floats is provided then the β " + "is selected for each new epoch and the last β is used for the rest of the epochs." + }, + ) + loss_type: str = field( + default="sigmoid", + metadata={ + "help": "Type of loss to use.", + "choices": ["sigmoid", "ipo"], + }, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + use_vllm: bool = field( + default=False, + metadata={ + "help": "Whether to use vLLM for generating completions. Requires vLLM to be installed " + "(`pip install vllm`)." + }, + ) + gpu_memory_utilization: Optional[float] = field( + default=0.55, + metadata={ + "help": "The vLLM memory utilization. The default value is 0.55.", + }, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + if hasattr(self.beta, "__len__") and len(self.beta) == 1: + self.beta = self.beta[0] diff --git a/lib/python3.12/site-packages/trl/trainer/online_dpo_trainer.py b/lib/python3.12/site-packages/trl/trainer/online_dpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..1818f34e28e4095ebf19269c361e7e6493bd8199 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/online_dpo_trainer.py @@ -0,0 +1,808 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import textwrap +import warnings +from functools import wraps +from pathlib import Path +from typing import Any, Callable, Optional, Union + +import datasets +import jinja2 +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data +from datasets import Dataset +from packaging import version +from torch.utils.data import DataLoader, IterableDataset +from transformers import ( + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + GenerationConfig, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + TrainerCallback, + is_apex_available, + is_wandb_available, +) +from transformers.trainer_utils import EvalPrediction, seed_worker +from transformers.training_args import OptimizerNames +from transformers.utils import is_peft_available, is_sagemaker_mp_enabled, logging + +from ..data_utils import apply_chat_template, is_conversational, maybe_apply_chat_template +from ..import_utils import is_vllm_available +from ..models import create_reference_model +from ..models.utils import unwrap_model_for_generation +from .judges import BasePairwiseJudge +from .online_dpo_config import OnlineDPOConfig +from .utils import ( + SIMPLE_CHAT_TEMPLATE, + DPODataCollatorWithPadding, + disable_dropout_in_model, + empty_cache, + generate_model_card, + get_comet_experiment_url, + get_reward, + prepare_deepspeed, + truncate_right, +) + + +if is_peft_available(): + from peft import PeftModel, get_peft_model + +if is_apex_available(): + from apex import amp + + +if is_sagemaker_mp_enabled(): + from smdistributed.modelparallel import __version__ as SMP_VERSION + + IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse("1.10") + +else: + IS_SAGEMAKER_MP_POST_1_10 = False + + +if is_vllm_available(): + from vllm import LLM, SamplingParams + +if is_wandb_available(): + import wandb + +logger = logging.get_logger(__name__) + + +class OnlineDPOTrainer(Trainer): + r""" + Initialize OnlineDPOTrainer. + + Args: + model (`transformers.PreTrainedModel` or `torch.nn.Module`): + The model to train, preferably an `AutoModelForCausalLM`. + ref_model (`transformers.PreTrainedModel` or `torch.nn.Module` or `None`): + The reference model to use for training. If None is specified, the reference model will be created from the + model. + reward_model (`transformers.PreTrainedModel` or `torch.nn.Module` or `None`): + The reward model to score completions with, preferably an `AutoModelForSequenceClassification`. + judge (`BasePairwiseJudge`): + The judge to use for pairwise comparison of model completions. + args (`OnlineDPOConfig`): + The online DPO config arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`DPODataCollatorWithPadding`) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + peft_config (`dict`): + The peft config to use for training. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + """ + + _tag_names = ["trl", "online-dpo"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module], + ref_model: Union[PreTrainedModel, nn.Module, None] = None, + reward_model: Union[PreTrainedModel, nn.Module, None] = None, + judge: Optional[BasePairwiseJudge] = None, + args: Optional[OnlineDPOConfig] = None, + data_collator: Optional[DataCollator] = None, + train_dataset: Optional[Union[Dataset, IterableDataset, "datasets.Dataset"]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset], "datasets.Dataset"]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + reward_processing_class: Optional[PreTrainedTokenizerBase] = None, + peft_config: Optional[dict] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + ) -> None: + if ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, either omit the `ref_model` argument or pass `None`." + ) + + self.ref_model = ref_model + + if reward_model is not None and judge is not None: + warnings.warn( + "Both `reward_model` and `judge` are provided. Please choose provide only one of them. " + "Ignoring `judge` and using `reward_model`.", + UserWarning, + ) + judge = None + elif reward_model is None and judge is None: + raise ValueError("Either `reward_model` or `judge` must be provided.") + + self.reward_model = reward_model + self.reward_processing_class = reward_processing_class + self.judge = judge + self.is_encoder_decoder = model.config.is_encoder_decoder + + if args.missing_eos_penalty is not None and judge is not None: + raise ValueError("`missing_eos_penalty` is not supported when `judge` is provided.") + + if args is None: + raise ValueError("`args` must be provided.") + + # Check that the processing_class is provided + if processing_class is None: + raise ValueError("`processing_class` must be provided.") + + # Convert to PEFT model if peft_config is provided + if peft_config is not None: + # Check if PEFT is available + if not is_peft_available(): + raise ImportError( + "PEFT is not available and passed `peft_config`. Please install PEFT with " + "`pip install peft` to use it." + ) + + # If the model is already a PeftModel, we need to merge and unload it. + # Further information here: https://huggingface.co/docs/trl/dpo_trainer#reference-model-considerations-with-peft + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + # Get peft model with the given config + model = get_peft_model(model, peft_config) + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + if self.ref_model is not None: + disable_dropout_in_model(self.ref_model) + + # Handle the ref_model + # Usually, the user wants the ref model to be the initial version of the model. When using PEFT, it's easy to + # get the ref model, as it's just the model with a disabled adapter. When not using PEFT, we need to create + # the ref model from the model by copying it and disable the gradients and set it in evaluation mode. + if ref_model is None: # No ref model provided, the most common case + if peft_config is None: + self.ref_model = create_reference_model(model) # copy, disable gradients, set eval mode + else: + self.ref_model = None # we don't need a ref model here, we can just disable the adapter. + else: # rare case, the user provided a ref model + self.ref_model = ref_model + self.ref_model.eval() + + # Disable the gradient and set the reward model in eval mode + if self.reward_model is not None: + self.reward_model.eval() + + # Define the collator is not provided + if data_collator is None: + data_collator = DPODataCollatorWithPadding(pad_token_id=processing_class.pad_token_id) + + self.max_length = args.max_length + + self.stats = { + "objective/kl": [], + "objective/entropy": [], + "objective/non_score_reward": [], + "rewards/chosen": [], + "rewards/rejected": [], + "rewards/accuracies": [], + "rewards/margins": [], + "logps/chosen": [], + "logps/rejected": [], + "val/contain_eos_token": [], + "beta": [], + } + if self.reward_model is not None: + self.stats["objective/rlhf_reward"] = [] + self.stats["objective/scores_margin"] = [] + self.stats["objective/scores"] = [] + + if args.use_vllm: + if not is_vllm_available(): + raise ImportError( + "vLLM is not available and `use_vllm` is set to True. Please install vLLM with " + "`pip install vllm` to use it." + ) + self.generation_config = SamplingParams( + n=2, # 2 generations per prompt + max_tokens=args.max_new_tokens, + temperature=args.temperature, + top_k=50, + top_p=1.0, + detokenize=False, # to avoid vllm to decode (we don't need it) + ) + # vLLM dynamically adjusts the size of the key-value cache based on available GPU memory at instantiation. + # A larger cache size improves speed, so we would expect gpu_memory_utilization=1. + # However, at this stage, the optimizer's weights are not yet loaded onto the GPU; they will be loaded + # after the first optimizer step and remain in GPU memory throughout training. So we must reserve enough + # space for them. Setting gpu_memory_utilization to 0.55 seems to work well in practice. + self.llm = LLM( + model=model.name_or_path, + gpu_memory_utilization=args.gpu_memory_utilization, + dtype=torch.float32, + # When release by vLLM, we would be able to distribute the model on multiple GPUs + # See https://github.com/vllm-project/vllm/pull/12071 + # tensor_parallel_size=torch.cuda.device_count(), + # distributed_executor_backend="external_launcher", + ) + else: + self.generation_config = GenerationConfig( + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + top_k=50, + top_p=1.0, + do_sample=True, + use_cache=False if args.gradient_checkpointing else True, + ) + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in Online DPO, the sampled data does not include + # the "input_ids" key. As a result, the trainer issues the warning: "Could not estimate the number of tokens + # of the input, floating-point operations will not be computed." To suppress this warning, we set the + # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate + # that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + self._beta = args.beta + + # Placed after the super().__init__ because we need self.is_deepspeed_enabled and self.accelerator + if self.is_deepspeed_enabled: + if self.reward_model is not None: + self.reward_model = prepare_deepspeed( + self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + if self.ref_model is not None: + self.ref_model = prepare_deepspeed( + self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + else: + if self.ref_model is not None: + self.ref_model = self.ref_model.to(self.accelerator.device) + if self.reward_model is not None: + self.reward_model = self.reward_model.to(self.accelerator.device) + + @property + def beta(self): + if isinstance(self._beta, list): + epoch = self.state.epoch + return self._beta[epoch] if epoch < len(self._beta) else self._beta[-1] + else: + return self._beta + + @staticmethod + def tokenize_row(feature, is_encoder_decoder: bool, tokenizer: PreTrainedTokenizerBase) -> dict[str, Any]: + """Tokenize a single row from a DPO specific dataset.""" + if not is_encoder_decoder: + batch = tokenizer(feature["prompt"], add_special_tokens=False) + # Add BOS token to head of prompt. Avoid adding if it's already there + if tokenizer.bos_token_id is not None: + prompt_len_input_ids = len(batch["input_ids"]) + if prompt_len_input_ids == 0 or tokenizer.bos_token_id != batch["input_ids"][0]: + batch["input_ids"] = [tokenizer.bos_token_id] + batch["input_ids"] + batch["attention_mask"] = [1] + batch["attention_mask"] + else: + batch = tokenizer(feature["prompt"], add_special_tokens=True) + batch = {f"prompt_{key}": value for key, value in batch.items()} + return batch + + # Same as Trainer.get_train_dataloader but skip the "remove_unused_columns". + @wraps(Trainer.get_train_dataloader) + def get_train_dataloader(self) -> DataLoader: + if self.train_dataset is None: + raise ValueError("Trainer: training requires a train_dataset.") + + train_dataset = self.train_dataset + data_collator = self.data_collator + dataloader_params = { + "batch_size": self._train_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(train_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_train_sampler() + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["worker_init_fn"] = seed_worker + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params)) + + # Same as Trainer.get_eval_dataloader but skip the "remove_unused_columns". + @wraps(Trainer.get_eval_dataloader) + def get_eval_dataloader(self, eval_dataset: Optional[Union[str, Dataset]] = None) -> DataLoader: + if eval_dataset is None and self.eval_dataset is None: + raise ValueError("Trainer: evaluation requires an eval_dataset.") + + # If we have persistent workers, don't do a fork bomb especially as eval datasets + # don't change during training + dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval" + if ( + hasattr(self, "_eval_dataloaders") + and dataloader_key in self._eval_dataloaders + and self.args.dataloader_persistent_workers + ): + return self.accelerator.prepare(self._eval_dataloaders[dataloader_key]) + + eval_dataset = ( + self.eval_dataset[eval_dataset] + if isinstance(eval_dataset, str) + else eval_dataset + if eval_dataset is not None + else self.eval_dataset + ) + data_collator = self.data_collator + + dataloader_params = { + "batch_size": self.args.eval_batch_size, + "collate_fn": data_collator, + "num_workers": self.args.dataloader_num_workers, + "pin_memory": self.args.dataloader_pin_memory, + "persistent_workers": self.args.dataloader_persistent_workers, + } + + if not isinstance(eval_dataset, torch.utils.data.IterableDataset): + dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset) + dataloader_params["drop_last"] = self.args.dataloader_drop_last + dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor + + # accelerator.free_memory() will destroy the references, so + # we need to store the non-prepared version + eval_dataloader = DataLoader(eval_dataset, **dataloader_params) + if self.args.dataloader_persistent_workers: + if hasattr(self, "_eval_dataloaders"): + self._eval_dataloaders[dataloader_key] = eval_dataloader + else: + self._eval_dataloaders = {dataloader_key: eval_dataloader} + + return self.accelerator.prepare(eval_dataloader) + + def _generate_vllm(self, model, prompts): + eos_token_id = self.processing_class.eos_token_id + pad_token_id = self.processing_class.pad_token_id + + # Load the latest weights + llm_model = self.llm.llm_engine.model_executor.driver_worker.model_runner.model + llm_model.load_weights(model.state_dict().items()) + + if is_conversational({"prompt": prompts[0]}): + outputs = self.llm.chat(prompts, self.generation_config, use_tqdm=False) + else: + outputs = self.llm.generate(prompts, self.generation_config, use_tqdm=False) + + completion_ids = [list(output.outputs[i].token_ids) for i in range(2) for output in outputs] + prompt_ids = [list(output.prompt_token_ids) for _ in range(2) for output in outputs] + + # Create mask and pad the prompt and completion + max_prompt_length = max(len(ids) for ids in prompt_ids) + prompt_mask = [[0] * (max_prompt_length - len(ids)) + [1] * len(ids) for ids in prompt_ids] + prompt_ids = [[pad_token_id] * (max_prompt_length - len(ids)) + ids for ids in prompt_ids] + max_tokens = self.generation_config.max_tokens + completion_mask = [[1] * len(ids) + [0] * (max_tokens - len(ids)) for ids in completion_ids] + completion_ids = [ + ids + [eos_token_id] if ids[-1] != eos_token_id and len(ids) < max_tokens else ids + for ids in completion_ids + ] + completion_ids = [ids + [pad_token_id] * (max_tokens - len(ids)) for ids in completion_ids] + + # Convert to tensors + prompt_ids = torch.tensor(prompt_ids, device=self.accelerator.device) + prompt_mask = torch.tensor(prompt_mask, device=self.accelerator.device) + completion_ids = torch.tensor(completion_ids, device=self.accelerator.device) + completion_mask = torch.tensor(completion_mask, device=self.accelerator.device) + + return prompt_ids, prompt_mask, completion_ids, completion_mask + + def _generate(self, model, prompts): + eos_token_id = self.processing_class.eos_token_id + pad_token_id = self.processing_class.pad_token_id + + # Apply chat template and tokenize the input. We do this on-the-fly to enable the use of reward models and + # policies with different tokenizers / chat templates. + inputs = [{"prompt": prompt} for prompt in prompts] + inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs] + inputs = [self.tokenize_row(x, self.is_encoder_decoder, self.processing_class) for x in inputs] + inputs = self.data_collator(inputs) + + # Sample 2 completions per prompt of size `max_new_tokens` from the model + inputs = self._prepare_inputs(inputs) + prompt_ids = inputs["prompt_input_ids"].repeat(2, 1) + prompt_mask = inputs["prompt_attention_mask"].repeat(2, 1) + with unwrap_model_for_generation( + model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + output = unwrapped_model.generate( + input_ids=prompt_ids, + attention_mask=prompt_mask, + generation_config=self.generation_config, + ) + + completion_ids = output[:, prompt_ids.size(1) :] + completion_ids, completion_mask = truncate_right(completion_ids, eos_token_id, pad_token_id) + + return prompt_ids, prompt_mask, completion_ids, completion_mask + + def _forward(self, model, prompt_ids, prompt_mask, completion_ids, completion_mask): + # Get the number of tokens to truncate from prompt + num_tokens_to_truncate = max(prompt_ids.size(1) + completion_ids.size(1) - self.max_length, 0) + + # Truncate left to avoid oom + prompt_ids = prompt_ids[:, num_tokens_to_truncate:] + prompt_mask = prompt_mask[:, num_tokens_to_truncate:] + + # Concat the prompt and completion + prompt_completion_ids = torch.cat((prompt_ids, completion_ids), dim=1) + prompt_completion_mask = torch.cat((prompt_mask, completion_mask), dim=1) + + # Get the logprobs of the completions from the model + output = model(prompt_completion_ids, attention_mask=prompt_completion_mask) + + # There is 1 offset, because the model predict the next token + logits = output.logits[:, prompt_ids.size(1) - 1 : -1] + + # Take the completion tokens logprob + logprobs = torch.take_along_dim(logits.log_softmax(dim=-1), completion_ids.unsqueeze(-1), dim=2).squeeze(-1) + return logprobs + + def training_step( + self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None + ) -> torch.Tensor: + model.train() + + prompts = inputs["prompt"] + batch_size = len(prompts) + + if self.args.use_vllm: + prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_vllm(model, prompts) + else: + prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate(model, prompts) + + contain_eos_token = torch.any(completion_ids == self.processing_class.eos_token_id, dim=-1) + + logprobs = self._forward(model, prompt_ids, prompt_mask, completion_ids, completion_mask) + with torch.no_grad(): + if self.ref_model is not None: + ref_logprobs = self._forward(self.ref_model, prompt_ids, prompt_mask, completion_ids, completion_mask) + else: # peft case: we just need to disable the adapter + with self.model.disable_adapter(): + ref_logprobs = self._forward(self.model, prompt_ids, prompt_mask, completion_ids, completion_mask) + + # Decode the completions, and format them if the input is conversational + device = logprobs.device + completions = self.processing_class.batch_decode(completion_ids, skip_special_tokens=True) + if is_conversational({"prompt": prompts[0]}): + completions = [[{"role": "assistant", "content": completion}] for completion in completions] + + # Get the reward from the reward model or judge + if self.judge is not None: + # Once formatted, conversational data may contain special tokens (such as <|im_start|>) that are not + # directly understandable by the judge and could alter its judgment. To avoid this and make the judge + # independent of the model's chat template, we use the raw conversation data, and apply our own chat + # template to it. + if is_conversational({"prompt": prompts[0]}): + environment = jinja2.Environment() + template = environment.from_string(SIMPLE_CHAT_TEMPLATE) + prompts = [template.render(messages=prompt) for prompt in prompts] + completions = [template.render(messages=completion) for completion in completions] + + ranks_of_first_completion = self.judge.judge( + prompts, list(zip(completions[:batch_size], completions[batch_size:])) + ) + + # convert ranks to a True/False mask: + # when rank == 0, it means the first completion is the best + # when rank == 1, it means the second completion is the best + mask = torch.tensor([rank == 0 for rank in ranks_of_first_completion], device=device) + else: + # The reward model may not have the same chat template or tokenizer as the model, so we need to use the + # raw data (string), apply the chat template (if needed), and tokenize it with the reward processing class. + prompts = 2 * prompts # repeat the prompt: [prompt0, prompt1] -> [prompt0, prompt1, prompt0, prompt1] + if is_conversational({"prompt": prompts[0]}): + examples = [{"prompt": p, "completion": c} for p, c in zip(prompts, completions)] + examples = [apply_chat_template(example, self.reward_processing_class) for example in examples] + prompts = [example["prompt"] for example in examples] + completions = [example["completion"] for example in examples] + + # Tokenize the prompts + prompts_ids = self.reward_processing_class( + prompts, padding=True, return_tensors="pt", padding_side="left" + )["input_ids"].to(device) + context_length = prompts_ids.shape[1] + + # Tokenize the completions + completions_ids = self.reward_processing_class( + completions, padding=True, return_tensors="pt", padding_side="right" + )["input_ids"].to(device) + + # Concatenate the prompts and completions and get the reward + prompt_completion_ids = torch.cat((prompts_ids, completions_ids), dim=1) + with torch.inference_mode(): + _, scores, _ = get_reward( + self.reward_model, prompt_completion_ids, self.reward_processing_class.pad_token_id, context_length + ) + + # Filter completion. Ensure that the sample contains stop_token_id + # Completions not passing that filter will receive a lower score. + if self.args.missing_eos_penalty is not None: + scores[~contain_eos_token] -= self.args.missing_eos_penalty + + # Split the scores in 2 (the prompts of the first half are the same as the second half) + first_half, second_half = scores.split(batch_size) + + # Get the indices of the chosen and rejected examples + mask = first_half >= second_half + + batch_range = torch.arange(batch_size, device=device) + chosen_indices = batch_range + (~mask * batch_size) + rejected_indices = batch_range + (mask * batch_size) + + # Build tensor so that the first half is the chosen examples and the second half the rejected examples + cr_indices = torch.cat((chosen_indices, rejected_indices), dim=0) # cr = chosen and rejected + cr_logprobs = logprobs[cr_indices] + cr_ref_logprobs = ref_logprobs[cr_indices] + + # mask out the padding tokens + padding_mask = ~completion_mask.bool() + cr_padding_mask = padding_mask[cr_indices] + + cr_logprobs_sum = (cr_logprobs * ~cr_padding_mask).sum(1) + cr_ref_logprobs_sum = (cr_ref_logprobs * ~cr_padding_mask).sum(1) + + # Split the chosen and rejected examples + chosen_logprobs_sum, rejected_logprobs_sum = torch.split(cr_logprobs_sum, batch_size) + chosen_ref_logprobs_sum, rejected_ref_logprobs_sum = torch.split(cr_ref_logprobs_sum, batch_size) + pi_logratios = chosen_logprobs_sum - rejected_logprobs_sum + ref_logratios = chosen_ref_logprobs_sum - rejected_ref_logprobs_sum + + logits = pi_logratios - ref_logratios + + if self.args.loss_type == "sigmoid": + losses = -F.logsigmoid(self.beta * logits) + elif self.args.loss_type == "ipo": + losses = (logits - 1 / (2 * self.beta)) ** 2 + else: + raise NotImplementedError(f"invalid loss type {self.loss_type}") + + loss = losses.mean() + + # Log everything + if self.reward_model is not None: + scores_margin = scores[chosen_indices] - scores[rejected_indices] + self.stats["objective/scores_margin"].append( + self.accelerator.gather_for_metrics(scores_margin.mean()).mean().item() + ) + self.stats["objective/scores"].append(self.accelerator.gather_for_metrics(scores.mean()).mean().item()) + self.stats["val/contain_eos_token"].append(contain_eos_token.float().mean().item()) + self.stats["logps/chosen"].append(self.accelerator.gather_for_metrics(chosen_logprobs_sum).mean().item()) + self.stats["logps/rejected"].append(self.accelerator.gather_for_metrics(rejected_logprobs_sum).mean().item()) + + kl = logprobs - ref_logprobs + mean_kl = kl.sum(1).mean() + self.stats["objective/kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item()) + non_score_reward = (-self.beta * kl).sum(1) + mean_non_score_reward = non_score_reward.mean() + self.stats["objective/non_score_reward"].append( + self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() + ) + if self.reward_model is not None: + rlhf_reward = scores + non_score_reward + self.stats["objective/rlhf_reward"].append(self.accelerator.gather_for_metrics(rlhf_reward).mean().item()) + mean_entropy = -logprobs.sum(1).mean() + self.stats["objective/entropy"].append(self.accelerator.gather_for_metrics(mean_entropy).mean().item()) + chosen_rewards = self.beta * (chosen_logprobs_sum - chosen_ref_logprobs_sum) + gathered_chosen_rewards = self.accelerator.gather_for_metrics(chosen_rewards) + self.stats["rewards/chosen"].append(gathered_chosen_rewards.mean().item()) + rejected_rewards = self.beta * (rejected_logprobs_sum - rejected_ref_logprobs_sum) + gathered_rejected_rewards = self.accelerator.gather_for_metrics(rejected_rewards) + self.stats["rewards/rejected"].append(gathered_rejected_rewards.mean().item()) + margin = gathered_chosen_rewards - gathered_rejected_rewards + self.stats["rewards/margins"].append(margin.mean().item()) + accuracy = margin > 0 + self.stats["rewards/accuracies"].append(accuracy.float().mean().item()) + self.stats["beta"].append(self.beta) + + if ( + self.args.torch_empty_cache_steps is not None + and self.state.global_step % self.args.torch_empty_cache_steps == 0 + ): + empty_cache() + + kwargs = {} + + # For LOMO optimizers you need to explicitly use the learning rate + if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + kwargs["learning_rate"] = self._get_learning_rate() + + if self.args.n_gpu > 1: + loss = loss.mean() # mean() to average on multi-gpu parallel training + + if self.use_apex: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: + scaled_loss.backward() + else: + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps + + # Same as Trainer._maybe_log_save_evaluate but log our metrics + def _maybe_log_save_evaluate( + self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, start_time, learning_rate=None + ): + if self.control.should_log and self.state.global_step > self._globalstep_last_logged: + logs: dict[str, float] = {} + + # all_gather + mean() to get average loss over all processes + tr_loss_scalar = self._nested_gather(tr_loss).mean().item() + + # reset tr_loss to zero + tr_loss -= tr_loss + + logs["loss"] = round(tr_loss_scalar / (self.state.global_step - self._globalstep_last_logged), 4) + if grad_norm is not None: + logs["grad_norm"] = grad_norm.detach().item() if isinstance(grad_norm, torch.Tensor) else grad_norm + if learning_rate is not None: + logs["learning_rate"] = learning_rate + else: + logs["learning_rate"] = self._get_learning_rate() + + # Add our metrics + for key, val in self.stats.items(): + logs[key] = sum(val) / len(val) + self.stats = {key: [] for key in self.stats} # reset stats + + self._total_loss_scalar += tr_loss_scalar + self._globalstep_last_logged = self.state.global_step + self.store_flos() + self.log(logs, start_time) + + metrics = None + if self.control.should_evaluate: + metrics = self._evaluate(trial, ignore_keys_for_eval) + is_new_best_metric = self._determine_best_metric(metrics=metrics, trial=trial) + + if self.args.save_strategy == "best": + self.control.should_save = is_new_best_metric + + if self.control.should_save: + self._save_checkpoint(model, trial) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{guo2024direct, + title = {{Direct Language Model Alignment from Online AI Feedback}}, + author = {Shangmin Guo and Biao Zhang and Tianlin Liu and Tianqi Liu and Misha Khalman and Felipe Llinares and Alexandre Ram{\'{e}} and Thomas Mesnard and Yao Zhao and Bilal Piot and Johan Ferret and Mathieu Blondel}, + year = 2024, + eprint = {arXiv:2402.04792} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="Online DPO", + trainer_citation=citation, + paper_title="Direct Language Model Alignment from Online AI Feedback", + paper_id="2402.04792", + ) + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/orpo_config.py b/lib/python3.12/site-packages/trl/trainer/orpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..84ee0388372fc66eba8a1d95964637e1f659c03b --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/orpo_config.py @@ -0,0 +1,163 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Any, Optional + +from transformers import TrainingArguments + + +@dataclass +class ORPOConfig(TrainingArguments): + r""" + Configuration class for the [`ORPOTrainer`]. + + This class includes only the parameters that are specific to ORPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch. This argument is required if you want + to use the default data collator. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt. This argument is required if you want to use the default data collator. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion. This argument is required if you want to use the default data collator + and your model is an encoder-decoder. + beta (`float`, *optional*, defaults to `0.1`): + Parameter controlling the relative ratio loss weight in the ORPO loss. In the + [paper](https://huggingface.co/papers/2403.07691), it is denoted by λ. In the + [code](https://github.com/xfactlab/orpo), it is denoted by `alpha`. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + label_pad_token_id (`int`, *optional*, defaults to `-100`): + Label pad token id. This argument is required if you want to use the default data collator. + padding_value (`int` or `None`, *optional*, defaults to `None`): + Padding value to use. If `None`, the padding value of the tokenizer is used. + truncation_mode (`str`, *optional*, defaults to `"keep_end"`): + Truncation mode to use when the prompt is too long. Possible values are `"keep_end"` or `"keep_start"`. + This argument is required if you want to use the default data collator. + generate_during_eval (`bool`, *optional*, defaults to `False`): + If `True`, generates and logs completions from the model to W&B or Comet during evaluation. + is_encoder_decoder (`bool` or `None`, *optional*, defaults to `None`): + When using the `model_init` argument (callable) to instantiate the model instead of the `model` argument, + you need to specify if the model returned by the callable is an encoder-decoder model. + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model from a + string. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-6, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + max_length: Optional[int] = field( + default=1024, + metadata={"help": "Maximum length of the sequences (prompt + completion) in the batch."}, + ) + max_prompt_length: Optional[int] = field( + default=512, + metadata={ + "help": "Maximum length of the prompt. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + max_completion_length: Optional[int] = field( + default=None, + metadata={ + "help": "Maximum length of the completion. This argument is required if you want to use the default data " + "collator and your model is an encoder-decoder." + }, + ) + beta: float = field( + default=0.1, + metadata={ + "help": "Parameter controlling the relative ratio loss weight in the ORPO loss. In the paper, it is " + "denoted by λ." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model."}, + ) + label_pad_token_id: int = field( + default=-100, + metadata={ + "help": "Label pad token id. This argument is required if you want to use the default data collator." + }, + ) + padding_value: Optional[int] = field( + default=None, + metadata={"help": "Padding value to use. If `None`, the padding value of the tokenizer is used."}, + ) + truncation_mode: str = field( + default="keep_end", + metadata={ + "help": "Truncation mode to use when the prompt is too long.", + "choices": ["keep_end", "keep_start"], + }, + ) + generate_during_eval: bool = field( + default=False, + metadata={"help": "If `True`, generates and logs completions from the model to W&B during evaluation."}, + ) + is_encoder_decoder: Optional[bool] = field( + default=None, + metadata={ + "help": "When using the `model_init` argument (callable) to instantiate the model instead of the `model` " + "argument, you need to specify if the model returned by the callable is an encoder-decoder model." + }, + ) + model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments to pass to `AutoModelForCausalLM.from_pretrained` when instantiating the model " + "from a string." + }, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/lib/python3.12/site-packages/trl/trainer/orpo_trainer.py b/lib/python3.12/site-packages/trl/trainer/orpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..784302e8bb1284e7fea3714dbcdef9f785d559ce --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/orpo_trainer.py @@ -0,0 +1,1093 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import random +import textwrap +import warnings +from collections import defaultdict +from contextlib import nullcontext +from pathlib import Path +from typing import Any, Callable, Literal, Optional, Union + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +import torch.nn.functional as F +from accelerate import PartialState +from datasets import Dataset +from torch import autocast +from torch.utils.data import DataLoader +from transformers import ( + AutoModelForCausalLM, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + is_comet_available, + is_torch_xla_available, + is_wandb_available, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalLoopOutput +from transformers.utils import is_peft_available, is_torch_fx_proxy + +from ..data_utils import maybe_apply_chat_template, maybe_extract_prompt +from .orpo_config import ORPOConfig +from .utils import ( + DPODataCollatorWithPadding, + add_bos_token_if_needed, + add_eos_token_if_needed, + disable_dropout_in_model, + generate_model_card, + get_comet_experiment_url, + log_table_to_comet_experiment, + pad_to_length, + peft_module_casting_to_bf16, + selective_log_softmax, +) + + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + + +if is_wandb_available(): + import wandb + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + +class ORPOTrainer(Trainer): + r""" + Initialize ORPOTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForSequenceClassification`. + args (`ORPOConfig`): + The ORPO config arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`DPODataCollatorWithPadding`) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + """ + + _tag_names = ["trl", "orpo"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module, str]] = None, + args: Optional[ORPOConfig] = None, + data_collator: Optional[DataCollator] = None, + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional[dict] = None, + compute_metrics: Optional[Callable[[EvalLoopOutput], dict]] = None, + ): + if args.model_init_kwargs is None: + model_init_kwargs = {} + elif not isinstance(model, str): + raise ValueError("You passed model_kwargs to the ORPOTrainer. But your model is already instantiated.") + else: + model_init_kwargs = args.model_init_kwargs + torch_dtype = model_init_kwargs.get("torch_dtype") + if torch_dtype is not None: + # Convert to `torch.dtype` if an str is passed + if isinstance(torch_dtype, str) and torch_dtype != "auto": + torch_dtype = getattr(torch, torch_dtype) + if torch_dtype != "auto" and not isinstance(torch_dtype, torch.dtype): + raise ValueError( + f"Invalid `torch_dtype` passed to the ORPOConfig. Expected a string with either `torch.dtype` or 'auto', but got {torch_dtype}." + ) + model_init_kwargs["torch_dtype"] = torch_dtype + + if isinstance(model, str): + model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs) + + # Initialize this variable to False. This helps tracking the case when `peft_module_casting_to_bf16` + # has been called in order to properly call autocast if needed. + self._peft_has_been_casted_to_bf16 = False + + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_config, we merge and unload it first + if isinstance(model, PeftModel): + model = model.merge_and_unload() + + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_loaded_in_4bit", False): + _support_gc_kwargs = hasattr( + args, "gradient_checkpointing_kwargs" + ) and "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if _support_gc_kwargs: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + # get peft model with the given config + model = get_peft_model(model, peft_config) + if args.bf16 and getattr(model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(model) + # If args.bf16 we need to explicitly call `generate` with torch amp autocast context manager + self._peft_has_been_casted_to_bf16 = True + + # For models that use gradient_checkpointing, we need to attach a hook that enables input + # to explicitly have `requires_grad=True`, otherwise training will either silently + # fail or completely fail. + elif args.gradient_checkpointing: + # For backward compatibility with older versions of transformers + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + if args.generate_during_eval and not (is_wandb_available() or is_comet_available()): + raise ValueError( + "`generate_during_eval=True` requires Weights and Biases or Comet to be installed." + " Please install `wandb` or `comet-ml` to resolve." + ) + + if model is not None: + self.is_encoder_decoder = model.config.is_encoder_decoder + elif args.is_encoder_decoder is None: + raise ValueError("When no model is provided, you need to pass the parameter is_encoder_decoder.") + else: + self.is_encoder_decoder = args.is_encoder_decoder + + if self.is_encoder_decoder: + self.decoder_start_token_id = model.config.decoder_start_token_id + self.pad_token_id = model.config.pad_token_id + + if processing_class is None: + raise ValueError("processing_class must be specified to tokenize a ORPO dataset.") + if args.max_length is None: + warnings.warn( + "`max_length` is not set in the ORPOConfig's init" + " it will default to `512` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_length = 512 + else: + max_length = args.max_length + if args.max_prompt_length is None: + warnings.warn( + "`max_prompt_length` is not set in the ORPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + max_prompt_length = 128 + else: + max_prompt_length = args.max_prompt_length + + if args.max_completion_length is None and self.is_encoder_decoder: + warnings.warn( + "When using an encoder decoder architecture, you should set `max_completion_length` in the ORPOConfig's init" + " it will default to `128` by default, but you should do it yourself in the future.", + UserWarning, + ) + self.max_completion_length = 128 + else: + self.max_completion_length = args.max_completion_length + + if data_collator is None: + data_collator = DPODataCollatorWithPadding( + pad_token_id=processing_class.pad_token_id, + label_pad_token_id=args.label_pad_token_id, + is_encoder_decoder=self.is_encoder_decoder, + ) + + if args.remove_unused_columns: + args.remove_unused_columns = False + # warn users + warnings.warn( + "When using DPODataCollatorWithPadding, you should set `remove_unused_columns=False` in your TrainingArguments" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_dpo_data_collator = True + else: + self.use_dpo_data_collator = False + + # Disable dropout in the model and reference model + if args.disable_dropout: + disable_dropout_in_model(model) + + self.max_length = max_length + self.generate_during_eval = args.generate_during_eval + self.label_pad_token_id = args.label_pad_token_id + self.padding_value = args.padding_value if args.padding_value is not None else processing_class.pad_token_id + self.max_prompt_length = max_prompt_length + self.truncation_mode = args.truncation_mode + self.processing_class = processing_class + + self.beta = args.beta + self.aux_loss_enabled = getattr(model.config, "output_router_logits", False) + self.aux_loss_coef = getattr(model.config, "router_aux_loss_coef", 0.0) + if self.aux_loss_enabled and self.aux_loss_coef == 0.0: + warnings.warn( + "You set `output_router_logits` to `True` in the model config, but `router_aux_loss_coef` is set to " + "`0.0`, meaning the auxiliary loss will not be used. Either set `router_aux_loss_coef` to a value " + "greater than `0.0`, or set `output_router_logits` to `False` if you don't want to use the auxiliary " + "loss.", + UserWarning, + ) + + self._stored_metrics = defaultdict(lambda: defaultdict(list)) + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in ORPO, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "prompt_input_ids", "chosen_input_ids", and + # "rejected_input_ids". As a result, the trainer issues the warning: "Could not estimate the number of tokens + # of the input, floating-point operations will not be computed." To suppress this warning, we set the + # "estimate_tokens" key in the model's "warnings_issued" dictionary to True. This acts as a flag to indicate + # that the warning has already been issued. + model.warnings_issued["estimate_tokens"] = True + + # Compute that only on the main process for faster data processing. + # see: https://github.com/huggingface/trl/pull/1255 + with PartialState().main_process_first(): + # Extract the prompt if needed, and apply the chat template if needed + train_dataset = train_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + train_dataset = train_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}, num_proc=args.dataset_num_proc + ) + train_dataset = train_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + if eval_dataset is not None: + eval_dataset = eval_dataset.map(maybe_extract_prompt, num_proc=args.dataset_num_proc) + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, + fn_kwargs={"tokenizer": processing_class}, + num_proc=args.dataset_num_proc, + ) + eval_dataset = eval_dataset.map(self.tokenize_row, num_proc=args.dataset_num_proc) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + if not hasattr(self, "accelerator"): + raise AttributeError( + "Your `Trainer` does not have an `accelerator` object. Consider upgrading `transformers`." + ) + + def build_tokenized_answer(self, prompt, answer): + """ + Llama tokenizer does satisfy `enc(a + b) = enc(a) + enc(b)`. It does ensure `enc(a + b) = enc(a) + enc(a + + b)[len(enc(a)):]`. Reference: + https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + """ + + full_tokenized = self.processing_class(prompt + answer, add_special_tokens=False) + prompt_input_ids = self.processing_class(prompt, add_special_tokens=False)["input_ids"] + + answer_input_ids = full_tokenized["input_ids"][len(prompt_input_ids) :] + answer_attention_mask = full_tokenized["attention_mask"][len(prompt_input_ids) :] + + # Concat tokens to form `enc(a) + enc(a + b)[len(enc(a)):]` + full_concat_input_ids = np.concatenate([prompt_input_ids, answer_input_ids]) + + # Prepare input tokens for token by token comparison + full_input_ids = np.array(full_tokenized["input_ids"]) + + if len(full_input_ids) != len(full_concat_input_ids): + raise ValueError("Prompt input ids and answer input ids should have the same length.") + + # On some tokenizers, like Llama-2 tokenizer, there are occasions where tokens + # can be merged together when tokenizing prompt+answer. This could result + # on the last token from the prompt being different when tokenized on its own + # vs when done as prompt+answer. + response_token_ids_start_idx = len(prompt_input_ids) + + # If tokenized prompt is different than both prompt+answer, then it means the + # last token has changed due to merging. + if prompt_input_ids != full_tokenized["input_ids"][:response_token_ids_start_idx]: + response_token_ids_start_idx -= 1 + + prompt_input_ids = full_tokenized["input_ids"][:response_token_ids_start_idx] + prompt_attention_mask = full_tokenized["attention_mask"][:response_token_ids_start_idx] + + if len(prompt_input_ids) != len(prompt_attention_mask): + raise ValueError("Prompt input ids and attention mask should have the same length.") + + answer_input_ids = full_tokenized["input_ids"][response_token_ids_start_idx:] + answer_attention_mask = full_tokenized["attention_mask"][response_token_ids_start_idx:] + + return dict( + prompt_input_ids=prompt_input_ids, + prompt_attention_mask=prompt_attention_mask, + input_ids=answer_input_ids, + attention_mask=answer_attention_mask, + ) + + def tokenize_row(self, feature, model: Optional[Union[PreTrainedModel, nn.Module]] = None) -> dict: + """Tokenize a single row from a ORPO specific dataset. + + At this stage, we don't convert to PyTorch tensors yet; we just handle the truncation in case the prompt + + chosen or prompt + rejected responses is/are too long. First we truncate the prompt; if we're still too long, + we truncate the chosen/rejected. + + We also create the labels for the chosen/rejected responses, which are of length equal to the sum of the length + of the prompt and the chosen/rejected response, with label_pad_token_id for the prompt tokens. + """ + batch = {} + prompt = feature["prompt"] + chosen = feature["chosen"] + rejected = feature["rejected"] + + if not self.is_encoder_decoder: + # Check issues below for more details + # 1. https://github.com/huggingface/trl/issues/907 + # 2. https://github.com/EleutherAI/lm-evaluation-harness/pull/531#issuecomment-1595586257 + # 3. https://github.com/LianjiaTech/BELLE/issues/337 + + if not isinstance(prompt, str): + raise ValueError(f"prompt should be an str but got {type(prompt)}") + prompt_tokens = self.processing_class(prompt, add_special_tokens=False) + prompt_tokens = {f"prompt_{k}": v for k, v in prompt_tokens.items()} + + if not isinstance(chosen, str): + raise ValueError(f"chosen should be an str but got {type(chosen)}") + chosen_tokens = self.build_tokenized_answer(prompt, chosen) + + if not isinstance(rejected, str): + raise ValueError(f"rejected should be an str but got {type(rejected)}") + rejected_tokens = self.build_tokenized_answer(prompt, rejected) + + # Last prompt token might get merged by tokenizer and + # it should not be included for generation if that happens + prompt_len_input_ids = len(prompt_tokens["prompt_input_ids"]) + + chosen_prompt_len_input_ids = len(chosen_tokens["prompt_input_ids"]) + rejected_prompt_len_input_ids = len(rejected_tokens["prompt_input_ids"]) + prompt_len_input_ids = min(chosen_prompt_len_input_ids, rejected_prompt_len_input_ids) + + for k, v in prompt_tokens.items(): + prompt_tokens[k] = v[:prompt_len_input_ids] + + # Make sure prompts only have one different token at most an + # and length only differs by 1 at most + num_diff_tokens = sum( + [a != b for a, b in zip(chosen_tokens["prompt_input_ids"], rejected_tokens["prompt_input_ids"])] + ) + num_diff_len = abs(chosen_prompt_len_input_ids - rejected_prompt_len_input_ids) + if num_diff_tokens > 1 or num_diff_len > 1: + raise ValueError( + "Chosen and rejected prompt_input_ids might only differ on the " + "last token due to tokenizer merge ops." + ) + + # add BOS token to head of prompt. Avoid adding if it's already there + prompt_tokens, chosen_tokens, rejected_tokens = add_bos_token_if_needed( + self.processing_class.bos_token_id, + prompt_len_input_ids, + prompt_tokens, + chosen_prompt_len_input_ids, + chosen_tokens, + rejected_prompt_len_input_ids, + rejected_tokens, + ) + + # add EOS token to end of answer. Avoid adding if it's already there + chosen_tokens, rejected_tokens = add_eos_token_if_needed( + self.processing_class.eos_token_id, chosen_tokens, rejected_tokens + ) + + longer_response_length = max(len(chosen_tokens["input_ids"]), len(rejected_tokens["input_ids"])) + + # if combined sequence is too long, truncate the prompt + for answer_tokens in [chosen_tokens, rejected_tokens, prompt_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + if self.truncation_mode == "keep_start": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_prompt_length] + elif self.truncation_mode == "keep_end": + for k in ["prompt_input_ids", "prompt_attention_mask"]: + answer_tokens[k] = answer_tokens[k][-self.max_prompt_length :] + else: + raise ValueError(f"Unknown truncation mode: {self.truncation_mode}") + + # if that's still too long, truncate the response + for answer_tokens in [chosen_tokens, rejected_tokens]: + if len(answer_tokens["prompt_input_ids"]) + longer_response_length > self.max_length: + for k in ["input_ids", "attention_mask"]: + answer_tokens[k] = answer_tokens[k][: self.max_length - self.max_prompt_length] + + # Create labels + chosen_sequence_tokens = { + k: chosen_tokens[f"prompt_{k}"] + chosen_tokens[k] for k in ["input_ids", "attention_mask"] + } + rejected_sequence_tokens = { + k: rejected_tokens[f"prompt_{k}"] + rejected_tokens[k] for k in ["input_ids", "attention_mask"] + } + chosen_sequence_tokens["labels"] = chosen_sequence_tokens["input_ids"][:] + chosen_sequence_tokens["labels"][: len(chosen_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(chosen_tokens["prompt_input_ids"]) + rejected_sequence_tokens["labels"] = rejected_sequence_tokens["input_ids"][:] + rejected_sequence_tokens["labels"][: len(rejected_tokens["prompt_input_ids"])] = [ + self.label_pad_token_id + ] * len(rejected_tokens["prompt_input_ids"]) + + for k, toks in { + "chosen_": chosen_sequence_tokens, + "rejected_": rejected_sequence_tokens, + "": prompt_tokens, + }.items(): + for type_key, tokens in toks.items(): + if type_key == "token_type_ids": + continue + batch[f"{k}{type_key}"] = tokens + + else: + chosen_tokens = self.processing_class( + chosen, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + rejected_tokens = self.processing_class( + rejected, truncation=True, max_length=self.max_completion_length, add_special_tokens=True + ) + prompt_tokens = self.processing_class( + prompt, truncation=True, max_length=self.max_prompt_length, add_special_tokens=True + ) + + batch["chosen_labels"] = chosen_tokens["input_ids"] + batch["rejected_labels"] = rejected_tokens["input_ids"] + batch["prompt_input_ids"] = prompt_tokens["input_ids"] + batch["prompt_attention_mask"] = prompt_tokens["attention_mask"] + + if model is not None and hasattr(model, "prepare_decoder_input_ids_from_labels"): + batch["rejected_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["rejected_labels"]) + ) + batch["chosen_decoder_input_ids"] = model.prepare_decoder_input_ids_from_labels( + labels=torch.tensor(batch["chosen_labels"]) + ) + + if is_torch_xla_available(): + # Pad the sequences to global max_length to avoid TorchXLA recompilation + for k in batch: + if "labels" in k or self.is_encoder_decoder: + pad_value = self.label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = self.padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + batch[k] = batch[k] + [pad_value] * (self.max_length - len(batch[k])) + return batch + + @staticmethod + def concatenated_inputs( + batch: dict[str, Union[list, torch.LongTensor]], + is_encoder_decoder: bool = False, + label_pad_token_id: int = -100, + padding_value: int = 0, + device: Optional[torch.device] = None, + ) -> dict[str, torch.LongTensor]: + """Concatenate the chosen and rejected inputs into a single tensor. + + Args: + batch: + A batch of data. Must contain the keys 'chosen_input_ids' and 'rejected_input_ids', which are tensors + of shape (batch_size, sequence_length). + is_encoder_decoder: + Whether the model is an encoder-decoder model. + label_pad_token_id: + The label pad token id. + padding_value: + The padding value to use for the concatenated inputs_ids. + device: + The device for the concatenated inputs. + + Returns: + A dictionary containing the concatenated inputs under the key 'concatenated_input_ids'. + """ + concatenated_batch = {} + + if is_encoder_decoder: + max_length = max(batch["chosen_labels"].shape[1], batch["rejected_labels"].shape[1]) + else: + max_length = max(batch["chosen_input_ids"].shape[1], batch["rejected_input_ids"].shape[1]) + + for k in batch: + if k.startswith("chosen") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("chosen", "concatenated") + concatenated_batch[concatenated_key] = pad_to_length(batch[k], max_length, pad_value=pad_value) + for k in batch: + if k.startswith("rejected") and isinstance(batch[k], torch.Tensor): + if "labels" in k or is_encoder_decoder: + pad_value = label_pad_token_id + elif k.endswith("_input_ids"): + pad_value = padding_value + elif k.endswith("_attention_mask"): + pad_value = 0 + concatenated_key = k.replace("rejected", "concatenated") + concatenated_batch[concatenated_key] = torch.cat( + ( + concatenated_batch[concatenated_key], + pad_to_length(batch[k], max_length, pad_value=pad_value), + ), + dim=0, + ).to(device=device) + + if is_encoder_decoder: + concatenated_batch["concatenated_input_ids"] = batch["prompt_input_ids"].repeat(2, 1).to(device=device) + concatenated_batch["concatenated_attention_mask"] = ( + batch["prompt_attention_mask"].repeat(2, 1).to(device=device) + ) + + return concatenated_batch + + def odds_ratio_loss( + self, + policy_chosen_logps: torch.FloatTensor, + policy_rejected_logps: torch.FloatTensor, + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Compute ORPO's odds ratio (OR) loss for a batch of policy and reference model log probabilities. + + Args: + policy_chosen_logps: + Log probabilities of the policy model for the chosen responses. Shape: (batch_size,) + policy_rejected_logps: + Log probabilities of the policy model for the rejected responses. Shape: (batch_size,) + + Returns: + A tuple of three tensors: (losses, chosen_rewards, rejected_rewards). The losses tensor contains the ORPO + loss for each example in the batch. The chosen_rewards and rejected_rewards tensors contain the rewards for + the chosen and rejected responses, respectively. The log odds ratio of the chosen responses over the + rejected responses ratio for logging purposes. The `log(sigmoid(log_odds_chosen))` for logging purposes. + """ + + # Derived from Eqs. (4) and (7) from https://huggingface.co/papers/2403.07691 by using log identities and exp(log(P(y|x)) = P(y|x) + log_odds = (policy_chosen_logps - policy_rejected_logps) - ( + torch.log1p(-torch.exp(policy_chosen_logps)) - torch.log1p(-torch.exp(policy_rejected_logps)) + ) + ratio = F.logsigmoid(log_odds) + losses = self.beta * ratio + + chosen_rewards = self.beta * (policy_chosen_logps.to(self.accelerator.device)).detach() + rejected_rewards = self.beta * (policy_rejected_logps.to(self.accelerator.device)).detach() + + return losses, chosen_rewards, rejected_rewards, torch.mean(ratio), torch.mean(log_odds) + + @staticmethod + def get_batch_logps( + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False, + label_pad_token_id: int = -100, + is_encoder_decoder: bool = False, + ) -> torch.FloatTensor: + """Compute the log probabilities of the given labels under the given logits. + + Args: + logits: Logits of the model (unnormalized). Shape: (batch_size, sequence_length, vocab_size) + labels: + Labels for which to compute the log probabilities. Label tokens with a value of label_pad_token_id are + ignored. Shape: (batch_size, sequence_length) + average_log_prob: + If True, return the average log probability per (non-masked) token. Otherwise, return the sum of the + log probabilities of the (non-masked) tokens. + label_pad_token_id: The label pad token id. + is_encoder_decoder: Whether the model is an encoder-decoder model. + + Returns: + A tensor of shape (batch_size,) containing the average/sum log probabilities of the given labels under the + given logits. + """ + if logits.shape[:-1] != labels.shape: + raise ValueError("Logits (batch and sequence length dim) and labels must have the same shape.") + + if not is_encoder_decoder: + labels = labels[:, 1:].clone() + logits = logits[:, :-1, :] + loss_mask = labels != label_pad_token_id + + # dummy token; we'll ignore the losses on these tokens later + labels = torch.where(labels == label_pad_token_id, 0, labels) + + per_token_logps = selective_log_softmax(logits, labels) + + if average_log_prob: + return (per_token_logps * loss_mask).sum(-1) / loss_mask.sum(-1) + else: + return (per_token_logps * loss_mask).sum(-1) + + def concatenated_forward( + self, model: nn.Module, batch: dict[str, Union[list, torch.LongTensor]] + ) -> tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + """Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together. + + We do this to avoid doing two forward passes, because it's faster for FSDP. + """ + concatenated_batch = self.concatenated_inputs( + batch, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + padding_value=self.padding_value, + device=self.accelerator.device, + ) + len_chosen = batch["chosen_labels"].shape[0] + + model_kwargs = ( + { + "decoder_input_ids": self._shift_right(concatenated_batch["concatenated_labels"]), + } + if self.is_encoder_decoder + else {} + ) + + if self.aux_loss_enabled: + model_kwargs["output_router_logits"] = True + + outputs = model( + concatenated_batch["concatenated_input_ids"], + attention_mask=concatenated_batch["concatenated_attention_mask"], + use_cache=False, + **model_kwargs, + ) + all_logits = outputs.logits + + def cross_entropy_loss(logits, labels): + if not self.is_encoder_decoder: + # Shift so that tokens < n predict n + logits = logits[..., :-1, :].contiguous() + labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = nn.CrossEntropyLoss() + logits = logits.view(-1, logits.shape[-1]) + labels = labels.view(-1) + # Enable model parallelism + labels = labels.to(logits.device) + loss = loss_fct(logits, labels) + return loss + + if self.is_encoder_decoder: + labels = concatenated_batch["concatenated_labels"].clone() + else: + labels = concatenated_batch["concatenated_input_ids"].clone() + attention_mask = concatenated_batch["concatenated_attention_mask"] + labels = torch.where(attention_mask == 1, labels, self.label_pad_token_id) + # orpo chosen nll loss is computed over the full prompt and response + chosen_nll_loss = cross_entropy_loss(all_logits[:len_chosen], labels[:len_chosen]) + + all_logps = self.get_batch_logps( + all_logits, + concatenated_batch["concatenated_labels"], + average_log_prob=True, + is_encoder_decoder=self.is_encoder_decoder, + label_pad_token_id=self.label_pad_token_id, + ) + + chosen_logps = all_logps[:len_chosen] + rejected_logps = all_logps[len_chosen:] + + if not self.is_encoder_decoder: + chosen_logits = all_logits[:len_chosen, :-1, :] + rejected_logits = all_logits[len_chosen:, :-1, :] + else: + chosen_logits = all_logits[:len_chosen] + rejected_logits = all_logits[len_chosen:] + + if self.aux_loss_enabled: + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss, outputs.aux_loss) + + return (chosen_logps, rejected_logps, chosen_logits, rejected_logits, chosen_nll_loss) + + def get_batch_loss_metrics( + self, + model, + batch: dict[str, Union[list, torch.LongTensor]], + train_eval: Literal["train", "eval"] = "train", + ): + """Compute the ORPO loss and other metrics for the given batch of inputs for train or test.""" + metrics = {} + + forward_output = self.concatenated_forward(model, batch) + ( + policy_chosen_logps, + policy_rejected_logps, + policy_chosen_logits, + policy_rejected_logits, + policy_nll_loss, + ) = forward_output[:5] + if self.aux_loss_enabled: + aux_loss = forward_output[5] + + losses, chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen = self.odds_ratio_loss( + policy_chosen_logps, policy_rejected_logps + ) + # full ORPO loss + loss = policy_nll_loss - losses.mean() + + reward_accuracies = (chosen_rewards > rejected_rewards).float() + + prefix = "eval_" if train_eval == "eval" else "" + metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(chosen_rewards).mean() + metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(rejected_rewards).mean() + metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(reward_accuracies).mean() + metrics[f"{prefix}rewards/margins"] = self.accelerator.gather_for_metrics( + chosen_rewards - rejected_rewards + ).mean() + metrics[f"{prefix}logps/rejected"] = self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean() + metrics[f"{prefix}logps/chosen"] = self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean() + metrics[f"{prefix}logits/rejected"] = self.accelerator.gather_for_metrics( + policy_rejected_logits.detach().mean() + ).mean() + metrics[f"{prefix}logits/chosen"] = self.accelerator.gather_for_metrics( + policy_chosen_logits.detach().mean() + ).mean() + metrics[f"{prefix}nll_loss"] = self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean() + metrics[f"{prefix}log_odds_ratio"] = self.accelerator.gather_for_metrics(log_odds_ratio).detach().mean() + metrics[f"{prefix}log_odds_chosen"] = self.accelerator.gather_for_metrics(log_odds_chosen).detach().mean() + if is_torch_xla_available(): + xm.mark_step() # needed because .item() calls + for k, v in metrics.items(): + metrics[k] = v.item() + if self.aux_loss_enabled: + loss += self.aux_loss_coef * aux_loss + + return loss, metrics + + def compute_loss( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + return_outputs=False, + num_items_in_batch=None, + ) -> Union[torch.Tensor, tuple[torch.Tensor, dict[str, torch.Tensor]]]: + compute_loss_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with compute_loss_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="train") + + # Make sure to move the loss to the device the original accumulating loss is at back in the `Trainer` class: + loss = loss.to(self.args.device) + + # force log the metrics + self.store_metrics(metrics, train_eval="train") + + if return_outputs: + return (loss, metrics) + return loss + + def generate_from_model(self, model, batch: dict[str, torch.LongTensor]) -> str: + """Generate samples from the model and reference model for the given batch of inputs.""" + + # If one uses `generate_during_eval` with peft + bf16, we need to explicitly call generate with + # the torch amp context manager as some hidden states are silently casted to full precision. + generate_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with generate_context_manager: + policy_output = model.generate( + input_ids=batch["prompt_input_ids"], + attention_mask=batch["prompt_attention_mask"], + max_length=self.max_length, + do_sample=True, + pad_token_id=self.processing_class.pad_token_id, + ) + + policy_output = pad_to_length(policy_output, self.max_length, self.processing_class.pad_token_id) + policy_output_decoded = self.processing_class.batch_decode(policy_output, skip_special_tokens=True) + + return policy_output_decoded + + def prediction_step( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[list[str]] = None, + ): + if not self.use_dpo_data_collator: + warnings.warn( + "prediction_step is only implemented for DPODataCollatorWithPadding, and you passed a datacollator that is different than " + "DPODataCollatorWithPadding - you might see unexpected behavior. Alternatively, you can implement your own prediction_step method if you are using a custom data collator" + ) + if ignore_keys is None: + if hasattr(model, "config"): + ignore_keys = getattr(model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + prediction_context_manager = ( + autocast(self.accelerator.device.type) if self._peft_has_been_casted_to_bf16 else nullcontext() + ) + + with torch.no_grad(), prediction_context_manager: + loss, metrics = self.get_batch_loss_metrics(model, inputs, train_eval="eval") + + # force log the metrics + self.store_metrics(metrics, train_eval="eval") + + if prediction_loss_only: + return (loss.detach(), None, None) + + # logits for the chosen and rejected samples from model + logits_dict = { + "eval_logits/chosen": metrics["eval_logits/chosen"], + "eval_logits/rejected": metrics["eval_logits/rejected"], + } + logits = [v for k, v in logits_dict.items() if k not in ignore_keys] + logits = torch.tensor(logits, device=self.accelerator.device) + labels = torch.zeros(logits.shape[0], device=self.accelerator.device) + + return (loss.detach(), logits, labels) + + def store_metrics(self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train") -> None: + for key, value in metrics.items(): + self._stored_metrics[train_eval][key].append(value) + + def evaluation_loop( + self, + dataloader: DataLoader, + description: str, + prediction_loss_only: Optional[bool] = None, + ignore_keys: Optional[list[str]] = None, + metric_key_prefix: str = "eval", + ) -> EvalLoopOutput: + """ + Overriding built-in evaluation loop to store metrics for each batch. Prediction/evaluation loop, shared by + `Trainer.evaluate()` and `Trainer.predict()`. + + Works both with or without labels. + """ + + # Sample and save to game log if requested (for one batch to save time) + if self.generate_during_eval: + # Generate random indices within the range of the total number of samples + num_samples = len(dataloader.dataset) + random_indices = random.sample(range(num_samples), k=self.args.eval_batch_size) + + # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader + random_batch_dataset = dataloader.dataset.select(random_indices) + random_batch = self.data_collator(random_batch_dataset) + random_batch = self._prepare_inputs(random_batch) + + policy_output_decoded = self.generate_from_model(self.model, random_batch) + + table = pd.DataFrame( + columns=["Prompt", "Policy"], + data=[ + [prompt, pol[len(prompt) :]] for prompt, pol in zip(random_batch["prompt"], policy_output_decoded) + ], + ) + if "wandb" in self.args.report_to: + wandb.log({"game_log": wandb.Table(data=table)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="game_log.csv", + table=table, + ) + + # Base evaluation + initial_output = super().evaluation_loop( + dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix + ) + + return initial_output + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + """ + Log `logs` on the various objects watching training, including stored metrics. + + Args: + logs (`dict[str, float]`): + The values to log. + start_time (`float` or `None`, *optional*, defaults to `None`): + Start time of the training. + """ + # logs either has 'loss' or 'eval_loss' + train_eval = "train" if "loss" in logs else "eval" + # Add averaged stored metrics to logs + for key, metrics in self._stored_metrics[train_eval].items(): + logs[key] = torch.tensor(metrics).mean().item() + del self._stored_metrics[train_eval] + return super().log(logs, start_time) + + def _shift_right(self, input_ids): + if self.decoder_start_token_id is None: + raise ValueError( + "model.config.decoder_start_token_id has to be defined. It is usually set to the pad_token_id." + ) + + # shift inputs to the right + if is_torch_fx_proxy(input_ids): + # Item assignment is not supported natively for proxies. + shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), self.decoder_start_token_id) + shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1) + else: + shifted_input_ids = input_ids.new_zeros(input_ids.shape) + shifted_input_ids[..., 1:] = input_ids[..., :-1].clone() + shifted_input_ids[..., 0] = self.decoder_start_token_id + + if self.pad_token_id is None: + raise ValueError("model.config.pad_token_id has to be defined.") + # replace possible -100 values in labels by `pad_token_id` + shifted_input_ids.masked_fill_(shifted_input_ids == -100, self.pad_token_id) + + return shifted_input_ids + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{hong2024orpo, + title = {{ORPO: Monolithic Preference Optimization without Reference Model}}, + author = {Jiwoo Hong and Noah Lee and James Thorne}, + year = 2024, + eprint = {arXiv:2403.07691} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="ORPO", + trainer_citation=citation, + paper_title="ORPO: Monolithic Preference Optimization without Reference Model", + paper_id="2403.07691", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/ppo_config.py b/lib/python3.12/site-packages/trl/trainer/ppo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..178063c1261049559d4a1785f58353013bf9a6ae --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/ppo_config.py @@ -0,0 +1,135 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from dataclasses import dataclass, field +from typing import Literal, Optional + +from ..trainer.utils import OnPolicyConfig + + +@dataclass +class PPOConfig(OnPolicyConfig): + r""" + Configuration class for the [`PPOTrainer`]. + + This class includes only the parameters that are specific to PPO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`OnPolicyConfig`] documentation. Note that default + values in this class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[:-3]`): + Name of this experiment. + reward_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`): + Path to the reward model. + model_adapter_name (`str` or `None`, *optional*, defaults to `None`): + Name of the train target PEFT adapter, when using LoRA with multiple adapters. + ref_adapter_name (`str` or `None`, *optional*, defaults to `None`): + Name of the reference PEFT adapter, when using LoRA with multiple adapters. + num_ppo_epochs (`int`, *optional*, defaults to `4`): + Number of epochs to train. + whiten_rewards (`bool`, *optional*, defaults to `False`): + Whether to whiten the rewards. + kl_coef (`float`, *optional*, defaults to `0.05`): + KL coefficient. + kl_estimator (`Literal["k1", "k3"]`, *optional*, defaults to `"k1"`): + Which estimator for KL-Divergence to use from [Approximating KL + Divergence](http://joschu.net/blog/kl-approx.html). Defaults to "k1", a straightforward, unbiased + estimator. Can be set to "k3", an unbiased estimator with lower variance which "appears to be a strictly + better estimator". Cannot be set to "k2", as it is used for logging purposes. + cliprange (`float`, *optional*, defaults to `0.2`): + Clip range. + vf_coef (`float`, *optional*, defaults to `0.1`): + Value function coefficient. + cliprange_value (`float`, *optional*, defaults to `0.2`): + Clip range for the value function. + gamma (`float`, *optional*, defaults to `1.0`): + Discount factor. + lam (`float`, *optional*, defaults to `0.95`): + Lambda value for GAE. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. + """ + + exp_name: str = field( + default=os.path.basename(__file__)[:-3], + metadata={"help": "Name of this experiment."}, + ) + reward_model_path: str = field( + default="EleutherAI/pythia-160m", + metadata={"help": "Path to the reward model."}, + ) + model_adapter_name: Optional[str] = field( + default=None, + metadata={"help": "Name of the train target PEFT adapter, when using LoRA with multiple adapters."}, + ) + ref_adapter_name: Optional[str] = field( + default=None, + metadata={"help": "Name of the reference PEFT adapter, when using LoRA with multiple adapters."}, + ) + num_ppo_epochs: int = field( + default=4, + metadata={"help": "Number of epochs to train."}, + ) + whiten_rewards: bool = field( + default=False, + metadata={"help": "Whether to whiten the rewards."}, + ) + kl_coef: float = field( + default=0.05, + metadata={"help": "KL coefficient."}, + ) + kl_estimator: Literal["k1", "k3"] = field( + default="k1", + metadata={ + "help": "Which estimator for KL-Divergence to use from Approximating KL Divergence " + "(http://joschu.net/blog/kl-approx.html). Defaults to 'k1', a straightforward, unbiased estimator. Can be " + "set to 'k3', an unbiased estimator with lower variance which 'appears to be a strictly better " + "estimator'. Cannot be set to 'k2', as it is used for logging purposes." + }, + ) + cliprange: float = field( + default=0.2, + metadata={"help": "Clip range."}, + ) + vf_coef: float = field( + default=0.1, + metadata={"help": "Value function coefficient."}, + ) + cliprange_value: float = field( + default=0.2, + metadata={"help": "Clip range for the value function."}, + ) + gamma: float = field( + default=1.0, + metadata={"help": "Discount factor."}, + ) + lam: float = field( + default=0.95, + metadata={"help": "Lambda value for GAE."}, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation." + }, + ) diff --git a/lib/python3.12/site-packages/trl/trainer/ppo_trainer.py b/lib/python3.12/site-packages/trl/trainer/ppo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..44156464ec5f800f74349b076336ff65c6428894 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/ppo_trainer.py @@ -0,0 +1,820 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import math +import os +import textwrap +import time +from collections import defaultdict +from contextlib import contextmanager, nullcontext +from pathlib import Path +from typing import Optional, Union + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +from accelerate import Accelerator +from accelerate.utils import broadcast, gather_object +from datasets import Dataset +from torch.utils.data import DataLoader +from transformers import ( + BaseImageProcessor, + DataCollatorWithPadding, + FeatureExtractionMixin, + GenerationConfig, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + TrainerCallback, + TrainerControl, + is_wandb_available, +) +from transformers.integrations import get_reporting_integration_callbacks +from transformers.trainer import DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK +from transformers.trainer_callback import CallbackHandler, ExportableState, PrinterCallback +from transformers.utils import is_peft_available, is_rich_available + +from ..core import masked_mean, masked_whiten +from ..models import create_reference_model +from ..models.utils import unwrap_model_for_generation +from .ppo_config import PPOConfig +from .utils import ( + OnlineTrainerState, + batch_generation, + disable_dropout_in_model, + empty_cache, + exact_div, + first_true_indices, + forward, + generate_model_card, + get_comet_experiment_url, + get_reward, + log_table_to_comet_experiment, + peft_module_casting_to_bf16, + prepare_deepspeed, + print_rich_table, + selective_log_softmax, + truncate_response, +) + + +if is_peft_available(): + from peft import PeftConfig, PeftModel, get_peft_model + +if is_wandb_available(): + import wandb + + +INVALID_LOGPROB = 1.0 + + +# taken from https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/ppo/ppo_trainer.py#L29 +# we did this we can do a single `model = accelerator.prepare(model)` +class PolicyAndValueWrapper(nn.Module): + def __init__(self, policy, value_model) -> None: + super().__init__() + self.policy = policy + self.value_model = value_model + self.critic_backbone = getattr(value_model, value_model.base_model_prefix) + + def forward(self, **kwargs): + output = self.critic_backbone(**kwargs) + logits = self.value_model.score(output.hidden_states[-1]) + return self.policy(**kwargs), logits + + +class PPOTrainer(Trainer): + _tag_names = ["trl", "ppo"] + + def __init__( + self, + args: PPOConfig, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ], + model: nn.Module, + ref_model: Optional[nn.Module], + reward_model: nn.Module, + train_dataset: Dataset, + value_model: nn.Module, + data_collator: Optional[DataCollatorWithPadding] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + # less commonly used + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + callbacks: Optional[list[TrainerCallback]] = None, + peft_config: Optional["PeftConfig"] = None, + ) -> None: + if ref_model is model: + raise ValueError( + "`model` and `ref_model` cannot be the same object. If you want `ref_model` to be the " + "same as `model`, you must make a copy of it, or `None` if you use peft." + ) + + self.args = args + self.processing_class = processing_class + self.policy_model = model + + # Define the collator if not provided + if data_collator is None: + data_collator = DataCollatorWithPadding(self.processing_class) + + # Handle stop token settings: update policy model's generation_config to use provided stop token + if args.stop_token and args.stop_token_id: + raise ValueError("You cannot set both `stop_token` and `stop_token_id`.") + elif args.stop_token: + if args.stop_token == "eos": + self.policy_model.generation_config.eos_token_id = self.stop_token_id = processing_class.eos_token_id + else: + raise ValueError( + f"Unknown `stop_token` {args.stop_token}. Allowed values are: `'eos'` and `None` (no stop token)." + ) + else: + self.policy_model.generation_config.eos_token_id = self.stop_token_id = args.stop_token_id # None or int + + # Check that the kl estimator is valid + if self.args.kl_estimator not in {"k1", "k3"}: + raise ValueError( + "kl_estimator must be either 'k1' (straightforward, unbiased) or 'k3' (lower variance, unbiased, " + "appears to be a strictly better estimator). See " + "[Approximating KL Divergence](http://joschu.net/blog/kl-approx.html) for details." + ) + + # peft support + if not is_peft_available() and peft_config is not None: + raise ImportError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + # if model is a peft model and we have a peft_confg, we merge and unload it first + if isinstance(self.policy_model, PeftModel): + self.policy_model = self.policy_model.merge_and_unload() + + # get peft model with the given config + self.policy_model = get_peft_model(self.policy_model, peft_config) + if args.bf16 and getattr(self.policy_model, "is_loaded_in_4bit", False): + peft_module_casting_to_bf16(self.policy_model) + + self.is_peft_model = is_peft_available() and isinstance(self.policy_model, PeftModel) + self.model_adapter_name = args.model_adapter_name + self.ref_adapter_name = args.ref_adapter_name + + if ref_model: + self.ref_model = ref_model + elif self.is_peft_model: + self.ref_model = None + else: + self.ref_model = create_reference_model(self.policy_model) + + self.reward_model = reward_model + self.train_dataset = train_dataset + self.train_dataset_len = len(train_dataset) + self.value_model = value_model + self.data_collator = data_collator + self.eval_dataset = eval_dataset + self.optimizer, self.lr_scheduler = optimizers + self.optimizer_cls_and_kwargs = None # needed for transformers >= 4.47 + + ######### + # calculate various batch sizes + ######### + if args.total_episodes is None: # allow the users to define episodes in terms of epochs. + args.total_episodes = int(args.num_train_epochs * self.train_dataset_len) + accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps) + self.accelerator = accelerator + args.world_size = accelerator.num_processes + args.local_batch_size = args.per_device_train_batch_size * args.gradient_accumulation_steps + args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size) + args.batch_size = int(args.local_batch_size * args.world_size) + args.mini_batch_size = exact_div( + args.batch_size, args.num_mini_batches, "`batch_size` must be a multiple of `num_mini_batches`" + ) + args.local_mini_batch_size = exact_div( + args.local_batch_size, args.num_mini_batches, "`local_batch_size` must be a multiple of `num_mini_batches`" + ) + if args.whiten_rewards: + assert args.local_mini_batch_size >= 8, ( + f"Per-rank minibatch size {args.local_mini_batch_size} is insufficient for whitening" + ) + # `per_rank_rollout_batch_size` is our `args.local_batch_size` + # `per_rank_minibatch_size` is our `args.local_mini_batch_size` + args.num_total_batches = math.ceil( + args.total_episodes / args.batch_size + ) # we may train for more than `total_episodes` + time_tensor = torch.tensor(int(time.time()), device=accelerator.device) + time_int = broadcast(time_tensor, 0).item() # avoid different timestamps across processes + args.run_name = f"{args.exp_name}__{args.seed}__{time_int}" + self.local_seed = args.seed + accelerator.process_index * 100003 # Prime + if args.num_sample_generations > 0: + self.sample_generations_freq = max(1, args.num_total_batches // args.num_sample_generations) + self.local_dataloader_batch_size = args.local_batch_size + + ######### + # setup model, optimizer, and others + ######### + for module in [self.policy_model, self.ref_model, self.value_model, self.reward_model]: + if module is not None: + disable_dropout_in_model(module) + self.model = PolicyAndValueWrapper(self.policy_model, self.value_model) + self.model.config = self.policy_model.config # needed for pushing to hub + self.create_optimizer_and_scheduler( + num_training_steps=args.num_total_batches + ) # note that we are calling `self.lr_scheduler.step()` manually only at the batch level + + ######### + ### trainer specifics + ######### + default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to) + self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks + self.callback_handler = CallbackHandler( + self.callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler + ) + self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK) + self.control = TrainerControl() + self.state = OnlineTrainerState( + is_local_process_zero=self.is_local_process_zero(), + is_world_process_zero=self.is_world_process_zero(), + stateful_callbacks=[ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ], + ) + self.current_flos = 0 + self.hp_search_backend = None + self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None + self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None + # Create distant repo and output directory if needed + self.hub_model_id = None + if self.args.push_to_hub: + self.init_hf_repo() + if self.args.should_save: + os.makedirs(self.args.output_dir, exist_ok=True) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + ######### + ### setup dataloader + ######### + self.dataloader = DataLoader( + self.train_dataset, + batch_size=self.local_dataloader_batch_size, + shuffle=True, + collate_fn=self.data_collator, + drop_last=True, # needed; otherwise the last batch will be of ragged shape + ) + # sync random states for DataLoader(shuffle=True) before `accelerator.prepare` + # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c + torch.manual_seed(args.seed) + self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader) + torch.manual_seed(self.local_seed) # reset the local seed again + + self.eval_dataloader = DataLoader( + self.eval_dataset, + batch_size=args.per_device_eval_batch_size, + collate_fn=self.data_collator, + drop_last=True, + ) # no need to shuffle eval dataset + self.eval_dataloader = accelerator.prepare(self.eval_dataloader) + + if self.is_deepspeed_enabled: + self.reward_model = prepare_deepspeed( + self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + + if self.ref_model is None: + if not self.is_peft_model: + raise ValueError("No reference model and model is not a Peft model.") + else: + self.ref_model = prepare_deepspeed( + self.ref_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + else: + if self.ref_model is None: + if not self.is_peft_model: + raise ValueError("No reference model and model is not a Peft model.") + else: + self.ref_model = self.ref_model.to(self.accelerator.device) + self.reward_model = self.reward_model.to(self.accelerator.device) + + def get_train_dataloader(self) -> DataLoader: + return self.dataloader + + def get_eval_dataloader(self) -> DataLoader: + return self.eval_dataloader + + @contextmanager + def null_ref_context(self): + """Context manager for handling null reference model (that is, peft adapter manipulation).""" + with ( + self.accelerator.unwrap_model(self.model.policy).disable_adapter() + if self.is_peft_model and not self.ref_adapter_name + else nullcontext() + ): + if self.ref_adapter_name: + self.model.policy.set_adapter(self.ref_adapter_name) + yield + if self.ref_adapter_name: + self.model.policy.set_adapter(self.model_adapter_name or "default") + + def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False): + backup_model = self.model + self.model = self.model.policy # save only the policy + + if self.is_deepspeed_enabled: + backup_deepspeed = self.deepspeed + self.deepspeed = self.model + + super().save_model(output_dir, _internal_call) + + self.model = backup_model + + if self.is_deepspeed_enabled: + self.deepspeed = backup_deepspeed + + def train(self): + args = self.args + accelerator = self.accelerator + optimizer = self.optimizer + model = self.model + ref_policy = self.ref_model + reward_model = self.reward_model + processing_class = self.processing_class + dataloader = self.dataloader + device = accelerator.device + + def repeat_generator(): + while True: + yield from dataloader + + iter_dataloader = iter(repeat_generator()) + generation_config = GenerationConfig( + max_new_tokens=args.response_length, + temperature=(args.temperature + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + accelerator.print("===training policy===") + start_time = time.time() + stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.gradient_accumulation_steps) + approxkl_stats = torch.zeros(stats_shape, device=device) + pg_clipfrac_stats = torch.zeros(stats_shape, device=device) + pg_loss_stats = torch.zeros(stats_shape, device=device) + vf_loss_stats = torch.zeros(stats_shape, device=device) + vf_clipfrac_stats = torch.zeros(stats_shape, device=device) + entropy_stats = torch.zeros(stats_shape, device=device) + ratio_stats = torch.zeros(stats_shape, device=device) + model.train() + + # trainer state initialization + self.state.global_step = 0 + self.state.episode = 0 + self.state.max_steps = args.num_total_batches + self.state.num_train_epochs = args.total_episodes / self.train_dataset_len + # Compute absolute values for logging, eval, and save if given as ratio + if args.logging_steps is not None: + if args.logging_steps < 1: + self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps) + else: + self.state.logging_steps = args.logging_steps + if args.eval_steps is not None: + if args.eval_steps < 1: + self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps) + else: + self.state.eval_steps = args.eval_steps + if args.save_steps is not None: + if args.save_steps < 1: + self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps) + else: + self.state.save_steps = args.save_steps + self.control = self.callback_handler.on_train_begin(args, self.state, self.control) + + # backward compatibility + if self.is_deepspeed_enabled: + self.deepspeed = self.model + self.model_wrapped = self.model + + for update in range(1, args.num_total_batches + 1): + self.state.episode += 1 * args.batch_size + data = next(iter_dataloader) + with torch.no_grad(): + queries = data["input_ids"].to(device) + context_length = queries.shape[1] + responses = [] + postprocessed_responses = [] + logprobs = [] + ref_logprobs = [] + scores = [] + sequence_lengths = [] + values = [] + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + query_responses, logitss = batch_generation( + unwrapped_model.policy, + queries, + args.local_rollout_forward_batch_size, + processing_class.pad_token_id, + generation_config, + ) + + for i in range(0, queries.shape[0], args.local_rollout_forward_batch_size): + query = queries[i : i + args.local_rollout_forward_batch_size] + query_response = query_responses[i : i + args.local_rollout_forward_batch_size] + response = query_response[:, context_length:] + logits = logitss[i : i + args.local_rollout_forward_batch_size] + logprob = selective_log_softmax(logits, response) + del logits + empty_cache() + + if ref_policy is None: + with self.null_ref_context(): + ref_output = forward(model.policy, query_response, processing_class.pad_token_id) + else: + ref_output = forward(ref_policy, query_response, processing_class.pad_token_id) + ref_logits = ref_output.logits[:, context_length - 1 : -1] + ref_logits /= args.temperature + 1e-7 + ref_logprob = selective_log_softmax(ref_logits, response) + del ref_output, ref_logits + empty_cache() + + # Response Processing 1. truncate response after the first occurrence of `stop_token_id` + postprocessed_response = response + if self.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + self.stop_token_id, processing_class.pad_token_id, response + ) + + # Response Processing 2. run reward model on the truncated responses + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + sequence_length = first_true_indices(postprocessed_response == processing_class.pad_token_id) - 1 + unwrapped_value_model = accelerator.unwrap_model(model).value_model + full_value, _, _ = get_reward( + unwrapped_value_model, query_response, processing_class.pad_token_id, context_length + ) + value = full_value[:, context_length - 1 : -1].squeeze(-1) + _, score, _ = get_reward( + reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length + ) + + responses.append(response) + postprocessed_responses.append(postprocessed_response) + logprobs.append(logprob) + ref_logprobs.append(ref_logprob) + sequence_lengths.append(sequence_length) + scores.append(score) + values.append(value) + responses = torch.cat(responses, 0) + postprocessed_responses = torch.cat(postprocessed_responses, 0) + logprobs = torch.cat(logprobs, 0) + ref_logprobs = torch.cat(ref_logprobs, 0) + sequence_lengths = torch.cat(sequence_lengths, 0) + scores = torch.cat(scores, 0) + values = torch.cat(values, 0) + del (logprob, ref_logprob, full_value, value, score, unwrapped_model) + empty_cache() + gc.collect() + + # Response Processing 3. Filter completion. Ensure that the sample contains stop_token_id + # Completions not passing that filter will receive a lower score. + contain_eos_token = torch.any(postprocessed_responses == self.processing_class.eos_token_id, dim=-1) + if self.args.missing_eos_penalty is not None: + scores[~contain_eos_token] -= self.args.missing_eos_penalty + # accelerator.print(f"{scores=}, {(contain_eos_token.sum() / len(contain_eos_token))=}") + + # be very careful with `padding_mask_p1`; see https://excalidraw.com/#json=LWnzG4w2k5DjF_EOL_xPt,e2w3a-hFJ_gX5vOfeyXGTw + response_idxs = torch.arange(responses.shape[1], device=responses.device).repeat(responses.shape[0], 1) + padding_mask = response_idxs > sequence_lengths.unsqueeze(1) + logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB) + ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB) + sequence_lengths_p1 = sequence_lengths + 1 + padding_mask_p1 = response_idxs > (sequence_lengths_p1.unsqueeze(1)) + values = torch.masked_fill(values, padding_mask_p1, 0) + + # 4. compute rewards + # Formula used by http://joschu.net/blog/kl-approx.html for the k1 and k3 estimators + logr = ref_logprobs - logprobs + kl = -logr if args.kl_estimator == "k1" else (logr.exp() - 1) - logr # Else statement is k3 + non_score_reward = -args.kl_coef * kl + rewards = non_score_reward.clone() + actual_start = torch.arange(rewards.size(0), device=rewards.device) + actual_end = torch.where(sequence_lengths_p1 < rewards.size(1), sequence_lengths_p1, sequence_lengths) + rewards[[actual_start, actual_end]] += scores + + # 5. whiten rewards + if args.whiten_rewards: + rewards = masked_whiten(rewards, mask=~padding_mask_p1, shift_mean=False) + rewards = torch.masked_fill(rewards, padding_mask_p1, 0) + + # 6. compute advantages and returns + lastgaelam = 0 + advantages_reversed = [] + gen_length = responses.shape[1] + for t in reversed(range(gen_length)): + nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0 + delta = rewards[:, t] + args.gamma * nextvalues - values[:, t] + lastgaelam = delta + args.gamma * args.lam * lastgaelam + advantages_reversed.append(lastgaelam) + advantages = torch.stack(advantages_reversed[::-1], axis=1) + returns = advantages + values + advantages = masked_whiten(advantages, ~padding_mask) + advantages = torch.masked_fill(advantages, padding_mask, 0) + empty_cache() + + # Do multiple epochs of PPO training, with a fresh random shuffle in each epoch + for ppo_epoch_idx in range(args.num_ppo_epochs): + b_inds = np.random.permutation(args.local_batch_size) + minibatch_idx = 0 + for mini_batch_start in range(0, args.local_batch_size, args.local_mini_batch_size): + mini_batch_end = mini_batch_start + args.local_mini_batch_size + mini_batch_inds = b_inds[mini_batch_start:mini_batch_end] + gradient_accumulation_idx = 0 + for micro_batch_start in range(0, args.local_mini_batch_size, args.per_device_train_batch_size): + with accelerator.accumulate(model): + micro_batch_end = micro_batch_start + args.per_device_train_batch_size + micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end] + mb_advantage = advantages[micro_batch_inds] + mb_responses = responses[micro_batch_inds] + mb_query_responses = query_responses[micro_batch_inds] + mb_logprobs = logprobs[micro_batch_inds] + mb_return = returns[micro_batch_inds] + mb_values = values[micro_batch_inds] + + output, vpred_temp = forward(model, mb_query_responses, processing_class.pad_token_id) + logits = output.logits[:, context_length - 1 : -1] + logits /= args.temperature + 1e-7 + new_logprobs = selective_log_softmax(logits, mb_responses) + new_logprobs = torch.masked_fill( + new_logprobs, padding_mask[micro_batch_inds], INVALID_LOGPROB + ) + vpred = vpred_temp[:, context_length - 1 : -1].squeeze(-1) + vpred = torch.masked_fill(vpred, padding_mask_p1[micro_batch_inds], 0) + vpredclipped = torch.clamp( + vpred, + mb_values - args.cliprange_value, + mb_values + args.cliprange_value, + ) + vf_losses1 = torch.square(vpred - mb_return) + vf_losses2 = torch.square(vpredclipped - mb_return) + vf_loss_max = torch.max(vf_losses1, vf_losses2) + vf_loss = 0.5 * masked_mean(vf_loss_max, ~padding_mask_p1[micro_batch_inds]) + vf_clipfrac = masked_mean( + (vf_losses2 > vf_losses1).float(), ~padding_mask_p1[micro_batch_inds] + ) + logprobs_diff = new_logprobs - mb_logprobs + ratio = torch.exp(logprobs_diff) + pg_losses = -mb_advantage * ratio + pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange) + pg_loss_max = torch.max(pg_losses, pg_losses2) + pg_loss = masked_mean(pg_loss_max, ~padding_mask[micro_batch_inds]) + loss = pg_loss + args.vf_coef * vf_loss + accelerator.backward(loss) + optimizer.step() + optimizer.zero_grad() + with torch.no_grad(): + pg_clipfrac = masked_mean( + (pg_losses2 > pg_losses).float(), ~padding_mask[micro_batch_inds] + ) + prob_dist = torch.nn.functional.softmax(logits, dim=-1) + entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1) + approxkl = 0.5 * (logprobs_diff**2).mean() + approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl + pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( + pg_clipfrac + ) + pg_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = pg_loss + vf_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = vf_loss + vf_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( + vf_clipfrac + ) + entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean() + ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ratio.mean() + gradient_accumulation_idx += 1 + minibatch_idx += 1 + # del everything and empty cache + # fmt: off + del ( + output, vpred_temp, logits, new_logprobs, vpred, vpredclipped, + vf_losses1, vf_losses2, vf_loss, vf_clipfrac, logprobs_diff, ratio, pg_losses, pg_losses2, pg_loss_max, + pg_loss, loss, pg_clipfrac, prob_dist, entropy, approxkl, mb_return, + mb_advantage, mb_values, mb_responses, mb_query_responses, mb_logprobs, + ) + # fmt: on + empty_cache() + with torch.no_grad(): + mean_kl = kl.sum(1).mean() + mean_entropy = (-logprobs).sum(1).mean() + mean_non_score_reward = non_score_reward.sum(1).mean() + rlhf_reward = mean_non_score_reward + scores.mean() + eps = int(self.state.episode / (time.time() - start_time)) + metrics = {} + metrics["eps"] = eps + metrics["objective/kl"] = self.accelerator.gather_for_metrics(mean_kl).mean().item() + metrics["objective/entropy"] = self.accelerator.gather_for_metrics(mean_entropy).mean().item() + metrics["objective/non_score_reward"] = ( + self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() + ) + metrics["objective/rlhf_reward"] = self.accelerator.gather_for_metrics(rlhf_reward).mean().item() + metrics["objective/scores"] = self.accelerator.gather_for_metrics(scores.mean()).mean().item() + metrics["policy/approxkl_avg"] = self.accelerator.gather_for_metrics(approxkl_stats).mean().item() + metrics["policy/clipfrac_avg"] = self.accelerator.gather_for_metrics(pg_clipfrac_stats).mean().item() + metrics["loss/policy_avg"] = self.accelerator.gather_for_metrics(pg_loss_stats).mean().item() + metrics["loss/value_avg"] = self.accelerator.gather_for_metrics(vf_loss_stats).mean().item() + metrics["val/clipfrac_avg"] = self.accelerator.gather_for_metrics(vf_clipfrac_stats).mean().item() + metrics["policy/entropy_avg"] = self.accelerator.gather_for_metrics(entropy_stats).mean().item() + metrics["val/ratio"] = self.accelerator.gather_for_metrics(ratio_stats).mean().item() + metrics["val/ratio_var"] = self.accelerator.gather_for_metrics(ratio_stats).var().item() + metrics["val/num_eos_tokens"] = (responses == processing_class.eos_token_id).sum().item() + metrics["lr"] = self.lr_scheduler.get_last_lr()[0] + metrics["episode"] = self.state.episode + self.state.epoch = self.state.episode / self.train_dataset_len # used by self.log + self.state.global_step += 1 + self.log(metrics) + + self.lr_scheduler.step() + self.control = self.callback_handler.on_step_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + del kl, mean_kl, mean_entropy, mean_non_score_reward, scores, metrics, non_score_reward + empty_cache() + gc.collect() + + if args.num_sample_generations > 0 and (update - 1) % self.sample_generations_freq == 0: + self.generate_completions(sampling=True) + empty_cache() + del ( + query_responses, + responses, + postprocessed_responses, + logprobs, + ref_logprobs, + values, + sequence_lengths, + contain_eos_token, + sequence_lengths_p1, + response_idxs, + padding_mask, + padding_mask_p1, + rewards, + actual_start, + actual_end, + advantages, + returns, + ) + empty_cache() + + # HF trainer specifics + self.control = self.callback_handler.on_train_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None, metrics=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + def generate_completions(self, sampling: bool = False): + args = self.args + processing_class = self.processing_class + generation_config = GenerationConfig( + max_new_tokens=self.args.response_length, + temperature=(0.01 + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + table = defaultdict(list) + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + for batch in self.eval_dataloader: + query = batch["input_ids"] + with torch.no_grad(): + context_length = query.shape[1] + query_response, _ = batch_generation( + unwrapped_model.policy, + query, + query.shape[0], + processing_class.pad_token_id, + generation_config, + ) + response = query_response[:, context_length:] + postprocessed_response = response + if self.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + self.stop_token_id, processing_class.pad_token_id, response + ) + table["query"].extend( + gather_object(processing_class.batch_decode(query, skip_special_tokens=True)) + ) + table["model response"].extend( + gather_object(processing_class.batch_decode(postprocessed_response)) + ) + + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + _, score, _ = get_reward( + self.reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length + ) + table["score"].extend(self.accelerator.gather_for_metrics(score).float().cpu().numpy()) + + if sampling: + break + df = pd.DataFrame(table) + + if self.accelerator.is_main_process: + if is_rich_available(): + print_rich_table(df.iloc[0 : 0 + 5]) + if "wandb" in args.report_to: + import wandb + + if wandb.run is not None: + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=df, + ) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{mziegler2019fine-tuning, + title = {{Fine-Tuning Language Models from Human Preferences}}, + author = {Daniel M. Ziegler and Nisan Stiennon and Jeffrey Wu and Tom B. Brown and Alec Radford and Dario Amodei and Paul F. Christiano and Geoffrey Irving}, + year = 2019, + eprint = {arXiv:1909.08593} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="PPO", + trainer_citation=citation, + paper_title="Fine-Tuning Language Models from Human Preferences", + paper_id="1909.08593", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/prm_config.py b/lib/python3.12/site-packages/trl/trainer/prm_config.py new file mode 100644 index 0000000000000000000000000000000000000000..af152ddf949f9d1e2c53e023051eac28347dd521 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/prm_config.py @@ -0,0 +1,114 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Optional + +from transformers import TrainingArguments + + +@dataclass +class PRMConfig(TrainingArguments): + r""" + Configuration class for the [`PRMTrainer`]. + + This class includes only the parameters that are specific to PRM training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) used for truncation. + max_prompt_length (`int` or `None`, *optional*, defaults to `512`): + Maximum length of the prompt used for truncation. + max_completion_length (`int` or `None`, *optional*, defaults to `None`): + Maximum length of the completion used for truncation. The completion is the concatenation of the steps. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + step_separator (`str`, *optional*, defaults to `"\n"`): + Separator used to separate each step of the reasoning process. + train_on_last_step_only (`bool`, *optional*, defaults to `False`): + Whether to train only on the last step. + dataset_num_proc (`int`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + """ + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=1e-5, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + average_tokens_across_devices: bool = field( + default=True, + metadata={ + "help": "Whether or not to average tokens across devices. If enabled, will use all_reduce to synchronize " + "num_tokens_in_batch for precise loss calculation. Reference: https://github.com/huggingface/transformers/issues/34242 " + }, + ) + + max_length: Optional[int] = field( + default=1024, + metadata={"help": "Maximum length of the sequences (prompt + completion) used for truncation."}, + ) + max_prompt_length: Optional[int] = field( + default=512, + metadata={"help": "Maximum length of the prompt used for truncation."}, + ) + max_completion_length: Optional[int] = field( + default=None, + metadata={ + "help": "Maximum length of the completion used for truncation. The completion is the concatenation of the " + "steps." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model and reference model."}, + ) + step_separator: str = field( + default="\n", + metadata={"help": "Separator used to separate each step of the reasoning process."}, + ) + train_on_last_step_only: bool = field( + default=False, + metadata={"help": "Whether to train only on the last step."}, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/lib/python3.12/site-packages/trl/trainer/prm_trainer.py b/lib/python3.12/site-packages/trl/trainer/prm_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..6cec4792dc37f2eced7d35ea5eddccaa13ae544d --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/prm_trainer.py @@ -0,0 +1,373 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import textwrap +import warnings +from itertools import chain +from pathlib import Path +from typing import Callable, Optional, Union + +import torch +import torch.nn as nn +from accelerate import PartialState +from datasets import Dataset, features +from transformers import ( + BaseImageProcessor, + DataCollator, + DataCollatorForTokenClassification, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + is_wandb_available, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_peft_available + +from .prm_config import PRMConfig +from .utils import compute_accuracy, disable_dropout_in_model, generate_model_card + + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + +if is_wandb_available(): + import wandb + + +class PRMTrainer(Trainer): + """ + Initialize PRMTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForTokenClassification`. + args (`PRMConfig`): + The arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`DataCollatorForTokenClassification`) will be used which will pad the sequences to the maximum length of + the sequences in the batch, given a dataset of paired sequences. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + compute_metrics (`Callable[[transformers.EvalPrediction], dict]`, *optional* defaults to `compute_accuracy`): + The metrics to use for evaluation. If no metrics are specified, the default metric (`compute_accuracy`) + will be used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped in + a PEFT model. + """ + + _tag_names = ["trl", "prm"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module]] = None, + args: Optional[PRMConfig] = None, + data_collator: Optional[DataCollator] = None, + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = ( + None, + None, + ), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional[dict] = None, + ): + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + if not isinstance(model, PeftModel): + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_quantized", False): + _supports_gc_kwargs = "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if not _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None: + warnings.warn( + "You passed `gradient_checkpointing_kwargs` in the trainer's kwargs, but your peft version does not support it. " + "please update to the latest version of peft to use `gradient_checkpointing_kwargs`." + ) + elif _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + + model = get_peft_model(model, peft_config) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + if compute_metrics is None: + compute_metrics = compute_accuracy + + if data_collator is None: + if processing_class is None: + raise ValueError( + "A processing_class must be specified when using the default DataCollatorForTokenClassification" + ) + data_collator = DataCollatorForTokenClassification(processing_class, max_length=args.max_length) + + if "input_ids" not in train_dataset.column_names: + with PartialState().main_process_first(): + fn_kwargs = { + "tokenizer": processing_class, + "step_separator": args.step_separator, + "max_length": args.max_length, + "max_prompt_length": args.max_prompt_length, + "max_completion_length": args.max_completion_length, + "train_on_last_step_only": args.train_on_last_step_only, + } + train_fn_kwargs = {**fn_kwargs, "is_eval": False} + train_dataset = train_dataset.map( + self.tokenize_row, + fn_kwargs=train_fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=train_dataset.features, + desc="Tokenizing train dataset", + features=features.Features( # needed to avoid map to cast labels to bool + { + "labels": features.Sequence(features.Value("int64")), + "input_ids": features.Sequence(features.Value("int64")), + } + ), + ) + + eval_fn_kwargs = {**fn_kwargs, "is_eval": True} + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + self.tokenize_row, + fn_kwargs=eval_fn_kwargs, + num_proc=args.dataset_num_proc, + remove_columns=eval_dataset.features, + desc="Tokenizing eval dataset", + features=features.Features( # needed to avoid map to cast labels to bool + { + "labels": features.Sequence(features.Value("int64")), + "input_ids": features.Sequence(features.Value("int64")), + } + ), + ) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + @staticmethod + def tokenize_row( + features, + tokenizer, + step_separator, + max_length, + max_prompt_length, + max_completion_length, + train_on_last_step_only, + is_eval, + ): + r""" + Tokenize a row of the dataset. + + Args: + features (`dict[str, str]`): + Row of the dataset, should contain the keys `"prompt"`, `"completions"`, and `"labels"`. + tokenizer (`PreTrainedTokenizerBase`): + Tokenizer used to process the data. + step_separator (`str`): + Separator between steps in the completion. + max_length (`int` or `None`): + Maximum length of the sequences (prompt + completion). If `None`, the sequences are not truncated. + max_prompt_length (`int` or `None`): + Maximum length of the prompt. If `None`, the prompt is not truncated. + max_completion_length (`int` or `None`): + Maximum length of the completion sequences. If `None`, the completion sequences are not truncated. + train_on_last_step_only (`bool`): + Whether to train only on the last step. If `True`, the labels are `-100` for all tokens except the last + token of the completion. + is_eval (`bool`): + Whether the function is used to tokenize samples from a training or an evaluation dataset. Used only if + `train_on_last_step_only` is set to `True`. + + Returns: + `dict[str, list[int]]`: + Tokenized sequences with the keys `"input_ids"`, and `"labels". + + Example: + ```python + >>> from transformers import AutoTokenizer + + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B") + >>> features = { + ... "prompt": "Which number is larger, 9.8 or 9.11?", + ... "completions": ["11 is greater than 8.", "Hence, 9.11 > 9.8."], + ... "labels": [True, False], + ... } + >>> PRMTrainer.tokenize_row( + ... features, tokenizer, "\n", max_completion_length=None, train_on_last_step_only=False, is_eval=False + ... ) + {'input_ids': [23085, 1372, 374, 8131, 11, 220, 24, 13, 23, 476, 220, 24, 13, 16, 16, 30, 16, 16, 374, 7046, 1091, 220, 23, 13, 198, 39, 763, 11, 220, 24, 13, 16, 16, 861, 220, 24, 13, 23, 13, 198], + 'labels': [-100, -100, -100, -100, -100, -100, -100, -100, 1, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, 0]} + ``` + """ + # Tokenize the prompt and completions + prompt_ids = tokenizer(features["prompt"], add_special_tokens=False)["input_ids"] + completions_ids = [ + tokenizer(completion, add_special_tokens=False)["input_ids"] for completion in features["completions"] + ] + if train_on_last_step_only and not is_eval: + labels = [-100] * (len(features["labels"]) - 1) + [int(features["labels"][-1])] + else: + labels = [int(label) for label in features["labels"]] + + # Get the ID of the separator token and add it to the completions + separator_ids = tokenizer.encode(step_separator, add_special_tokens=False) + completions_ids = [completion + separator_ids for completion in completions_ids] + + # Create the label + labels = [[-100] * (len(completion) - 1) + [label] for completion, label in zip(completions_ids, labels)] + + # Join the completions and labels steps + completion_ids = list(chain(*completions_ids)) + labels = list(chain(*labels)) + + if tokenizer.bos_token_id is not None: + prompt_ids = [tokenizer.bos_token_id] + prompt_ids + + # Truncate prompt and completion sequences + if max_prompt_length is not None: + prompt_ids = prompt_ids[-max_prompt_length:] + if max_completion_length is not None: + completion_ids = completion_ids[:max_completion_length] + labels = labels[:max_completion_length] + + input_ids = prompt_ids + completion_ids + labels = [-100] * len(prompt_ids) + labels + + if max_length is not None: + input_ids = input_ids[:max_length] + labels = labels[:max_length] + + return {"input_ids": input_ids, "labels": labels} + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{uesato2022solving, + title = {{Solving Math Word Problems With Process- and Outcome-Based Feedback}}, + author = {Uesato, Jonathan and Kushman, Nate and Kumar, Ramana and Song, Francis and Siegel, Noah and Wang, Lisa and Creswell, Antonia and Irving, Geoffrey and Higgins, Irina}, + year = 2022, + journal = {arXiv preprint arXiv:2211.14275} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + trainer_name="PRM", + trainer_citation=citation, + paper_title="Solving math word problems with process-and outcome-based feedback", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/reward_config.py b/lib/python3.12/site-packages/trl/trainer/reward_config.py new file mode 100644 index 0000000000000000000000000000000000000000..ab25b23df47b8aa8095a7022241b661789e2e5a0 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/reward_config.py @@ -0,0 +1,107 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field +from typing import Optional + +from transformers import TrainingArguments + + +@dataclass +class RewardConfig(TrainingArguments): + r""" + Configuration class for the [`RewardTrainer`]. + + This class includes only the parameters that are specific to Reward training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the sequences (prompt + completion) in the batch, filters out entries that exceed the + limit. This argument is required if you want to use the default data collator. + disable_dropout (`bool`, *optional*, defaults to `True`): + Whether to disable dropout in the model. + dataset_num_proc (`int`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + center_rewards_coefficient (`float`, *optional*, defaults to `None`): + Coefficient to incentivize the reward model to output mean-zero rewards (proposed by + https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`. + remove_unused_columns (`bool`, *optional*, defaults to `False`): + Whether to remove the columns that are not used by the model's forward pass. Can be `True` only if the + dataset is pretokenized. + """ + + # Parameters whose default values are overridden from TrainingArguments + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + average_tokens_across_devices: bool = field( + default=True, + metadata={ + "help": "Whether or not to average tokens across devices. If enabled, will use all_reduce to synchronize " + "num_tokens_in_batch for precise loss calculation. Reference: https://github.com/huggingface/transformers/issues/34242 " + }, + ) + + max_length: Optional[int] = field( + default=1024, + metadata={ + "help": "Maximum length of the sequences (prompt + completion) in the batch, filters out entries that " + "exceed the limit. This argument is required if you want to use the default data collator." + }, + ) + disable_dropout: bool = field( + default=True, + metadata={"help": "Whether to disable dropout in the model and reference model."}, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + center_rewards_coefficient: Optional[float] = field( + default=None, + metadata={ + "help": "Coefficient to incentivize the reward model to output mean-zero rewards (proposed by " + "https://huggingface.co/papers/2312.09244, Eq. 2). Recommended value: `0.01`." + }, + ) + remove_unused_columns: bool = field( + default=False, + metadata={ + "help": "Whether to remove the columns that are not used by the model's forward pass. Can be `True` only " + "if the dataset is pretokenized." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() diff --git a/lib/python3.12/site-packages/trl/trainer/reward_trainer.py b/lib/python3.12/site-packages/trl/trainer/reward_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..8fa0f082db986720e8289418c723835723f73cf2 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/reward_trainer.py @@ -0,0 +1,430 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import os +import warnings +from collections import defaultdict +from dataclasses import FrozenInstanceError, replace +from pathlib import Path +from typing import Any, Callable, Optional, Union + +import pandas as pd +import torch +import torch.nn as nn +from accelerate import PartialState +from accelerate.utils import gather_object +from datasets import Dataset +from transformers import ( + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + is_wandb_available, +) +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_pt_utils import nested_detach +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_peft_available, is_rich_available + +from ..data_utils import maybe_apply_chat_template +from .reward_config import RewardConfig +from .utils import ( + RewardDataCollatorWithPadding, + compute_accuracy, + decode_and_strip_padding, + disable_dropout_in_model, + generate_model_card, + get_comet_experiment_url, + log_table_to_comet_experiment, + print_rich_table, +) + + +if is_peft_available(): + from peft import PeftModel, get_peft_model, prepare_model_for_kbit_training + +if is_wandb_available(): + import wandb + + +def _tokenize(batch: dict[str, list[Any]], tokenizer: "PreTrainedTokenizerBase") -> dict[str, list[Any]]: + """Tokenize a batch from a reward modelling dataset.""" + new_examples = { + "input_ids_chosen": [], + "attention_mask_chosen": [], + "input_ids_rejected": [], + "attention_mask_rejected": [], + } + for chosen, rejected in zip(batch["chosen"], batch["rejected"]): + tokenized_chosen = tokenizer(chosen) + tokenized_rejected = tokenizer(rejected) + new_examples["input_ids_chosen"].append(tokenized_chosen["input_ids"]) + new_examples["attention_mask_chosen"].append(tokenized_chosen["attention_mask"]) + new_examples["input_ids_rejected"].append(tokenized_rejected["input_ids"]) + new_examples["attention_mask_rejected"].append(tokenized_rejected["attention_mask"]) + + return new_examples + + +class RewardTrainer(Trainer): + _tag_names = ["trl", "reward-trainer"] + + def __init__( + self, + model: Optional[Union[PreTrainedModel, nn.Module]] = None, + args: Optional[RewardConfig] = None, + data_collator: Optional[DataCollator] = None, + train_dataset: Optional[Dataset] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + model_init: Optional[Callable[[], PreTrainedModel]] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = ( + None, + None, + ), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional[dict] = None, + ): + """ + Initialize RewardTrainer. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForSequenceClassification`. + args (`RewardConfig`): + The arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`RewardDataCollatorWithPadding`) will be used which will pad the sequences to the maximum length of + the sequences in the batch, given a dataset of paired sequences. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the + inputs for the model, and it will be saved along the model to make it easier to rerun an interrupted + training or reuse the fine-tuned model. + model_init (`Callable[[], transformers.PreTrainedModel]`): + The model initializer to use for training. If None is specified, the default model initializer will be + used. + compute_metrics (`Callable[[transformers.EvalPrediction], dict]`, *optional* defaults to `compute_accuracy`): + The metrics to use for evaluation. If no metrics are specified, the default metric (`compute_accuracy`) + will be used. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + peft_config (`dict`, defaults to `None`): + The PEFT configuration to use for training. If you pass a PEFT configuration, the model will be wrapped + in a PEFT model. + """ + if not is_peft_available() and peft_config is not None: + raise ValueError( + "PEFT is not installed and you passed a `peft_config` in the trainer's kwargs, please install it to use the PEFT models" + ) + elif is_peft_available() and peft_config is not None: + if not isinstance(model, PeftModel): + if getattr(model, "is_loaded_in_8bit", False) or getattr(model, "is_quantized", False): + _supports_gc_kwargs = "gradient_checkpointing_kwargs" in list( + inspect.signature(prepare_model_for_kbit_training).parameters + ) + + prepare_model_kwargs = {"use_gradient_checkpointing": args.gradient_checkpointing} + + if not _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None: + warnings.warn( + "You passed `gradient_checkpointing_kwargs` in the trainer's kwargs, but your peft version does not support it. " + "please update to the latest version of peft to use `gradient_checkpointing_kwargs`.", + UserWarning, + ) + elif _supports_gc_kwargs and args.gradient_checkpointing_kwargs is not None: + prepare_model_kwargs["gradient_checkpointing_kwargs"] = args.gradient_checkpointing_kwargs + + model = prepare_model_for_kbit_training(model, **prepare_model_kwargs) + + model = get_peft_model(model, peft_config) + + # Disable dropout in the model + if args.disable_dropout: + disable_dropout_in_model(model) + + if compute_metrics is None: + compute_metrics = compute_accuracy + + if data_collator is None: + if processing_class is None: + raise ValueError( + "A processing_class must be specified when using the default RewardDataCollatorWithPadding" + ) + + max_length = args.max_length + + data_collator = RewardDataCollatorWithPadding(processing_class) + + if args.remove_unused_columns: + try: # for bc before https://github.com/huggingface/transformers/pull/25435 + args.remove_unused_columns = False + except FrozenInstanceError: + args = replace(args, remove_unused_columns=False) + # warn users + warnings.warn( + "When using RewardDataCollatorWithPadding, you should set `remove_unused_columns=False` in your RewardConfig" + " we have set it for you, but you should do it yourself in the future.", + UserWarning, + ) + + self.use_reward_data_collator = True + else: + self.use_reward_data_collator = False + + # The trainer estimates the number of FLOPs (floating-point operations) using the number of elements in the + # input tensor associated with the key "input_ids". However, in Reward, the sampled data does not include the + # "input_ids" key. Instead, the available keys are "input_ids_chosen" and "input_ids_rejected". As a result, + # the trainer issues the warning: "Could not estimate the number of tokens of the input, floating-point + # operations will not be computed." To suppress this warning, we set the "estimate_tokens" key in the model's + # "warnings_issued" dictionary to True. This acts as a flag to indicate that the warning has already been + # issued. + model.warnings_issued["estimate_tokens"] = True + + if "input_ids_chosen" not in train_dataset.column_names: + with PartialState().main_process_first(): + fn_kwargs = {"tokenizer": processing_class} + train_dataset = train_dataset.map(maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class}) + train_dataset = train_dataset.map( + _tokenize, + batched=True, + fn_kwargs=fn_kwargs, + num_proc=args.dataset_num_proc, + ) + # This filter is important because otherwise you get samples that exceed the model's context length and + # get truncated => noisy signal the chosen/rejected label gets lost. The downside is that the + # user might get surprised if N samples are missing from training. + train_dataset = train_dataset.filter( + lambda x: len(x["input_ids_chosen"]) <= max_length and len(x["input_ids_rejected"]) <= max_length, + num_proc=args.dataset_num_proc, + ) + if eval_dataset is not None: + eval_dataset = eval_dataset.map( + maybe_apply_chat_template, fn_kwargs={"tokenizer": processing_class} + ) + eval_dataset = eval_dataset.map( + _tokenize, + fn_kwargs=fn_kwargs, + batched=True, + num_proc=args.dataset_num_proc, + ) + # This filter is important because otherwise you get samples that exceed the model's context length and + # get truncated => noisy signal the chosen/rejected label gets lost. The downside is that the + # user might get surprised if N samples are missing from training. + eval_dataset = eval_dataset.filter( + lambda x: len(x["input_ids_chosen"]) <= max_length + and len(x["input_ids_rejected"]) <= max_length, + num_proc=args.dataset_num_proc, + ) + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + model_init=model_init, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + def compute_loss( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + return_outputs=False, + num_items_in_batch=None, + ) -> Union[torch.Tensor, tuple[torch.Tensor, dict[str, torch.Tensor]]]: + rewards_chosen = model( + input_ids=inputs["input_ids_chosen"], + attention_mask=inputs["attention_mask_chosen"], + return_dict=True, + )["logits"] + rewards_rejected = model( + input_ids=inputs["input_ids_rejected"], + attention_mask=inputs["attention_mask_rejected"], + return_dict=True, + )["logits"] + # calculate loss, optionally modulate with margin + if "margin" in inputs: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - inputs["margin"]).mean() + else: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean() + + if self.args.center_rewards_coefficient is not None: + loss += self.args.center_rewards_coefficient * torch.mean((rewards_chosen + rewards_rejected) ** 2) + + if return_outputs: + return loss, { + "rewards_chosen": rewards_chosen, + "rewards_rejected": rewards_rejected, + } + return loss + + def prediction_step( + self, + model: Union[PreTrainedModel, nn.Module], + inputs: dict[str, Union[torch.Tensor, Any]], + prediction_loss_only: bool, + ignore_keys: Optional[list[str]] = None, + ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + inputs = self._prepare_inputs(inputs) + if ignore_keys is None: + if hasattr(self.model, "config"): + ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", []) + else: + ignore_keys = [] + + with torch.no_grad(): + loss, logits_dict = self.compute_loss(model, inputs, return_outputs=True) + + if prediction_loss_only: + return (loss, None, None) + + loss = loss.detach() + logits = tuple(v for k, v in logits_dict.items() if k not in ignore_keys) + logits = nested_detach(logits) + # Stack accepted against rejected, mean over logits + # and softmax to get preferences between accepted and rejected to sum to 1 + logits = torch.stack(logits).mean(dim=2).softmax(dim=0).T + + labels = torch.zeros(logits.shape[0]) + labels = self._prepare_inputs(labels) + + return loss, logits, labels + + def evaluate(self, *args, **kwargs): + num_print_samples = kwargs.pop("num_print_samples", 4) + self.visualize_samples(num_print_samples) + return super().evaluate(*args, **kwargs) + + def visualize_samples(self, num_print_samples: int): + """ + Visualize the reward model logits prediction + + Args: + num_print_samples (`int`, defaults to `4`): + The number of samples to print. Set to `-1` to print all samples. + """ + eval_dataloader = self.get_eval_dataloader() + table = defaultdict(list) + for _, inputs in enumerate(eval_dataloader): + _, logits, _ = self.prediction_step(self.model, inputs, prediction_loss_only=False) + chosen_text = decode_and_strip_padding(inputs["input_ids_chosen"], self.processing_class) + rejected_text = decode_and_strip_padding(inputs["input_ids_rejected"], self.processing_class) + table["chosen_text"].extend(gather_object(chosen_text)) + table["rejected_text"].extend(gather_object(rejected_text)) + table["logits"].extend( + gather_object([[round(inner_item, 4) for inner_item in item] for item in logits.tolist()]) + ) + if num_print_samples >= 0 and len(table["chosen_text"]) >= num_print_samples: + break + df = pd.DataFrame(table) + if self.accelerator.process_index == 0: + if is_rich_available(): + print_rich_table(df[:num_print_samples]) + if "wandb" in self.args.report_to: + import wandb + + if wandb.run is not None: + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in self.args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=df, + ) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="Reward", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/rloo_config.py b/lib/python3.12/site-packages/trl/trainer/rloo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..e13680a1cc8231a25ea829b7b5972f48a0f0019f --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/rloo_config.py @@ -0,0 +1,114 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from dataclasses import dataclass, field + +from ..trainer.utils import OnPolicyConfig + + +@dataclass +class RLOOConfig(OnPolicyConfig): + r""" + Configuration class for the [`RLOOTrainer`]. + + This class includes only the parameters that are specific to RLOO training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] and [`OnPolicyConfig`] documentation. Note that default + values in this class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + exp_name (`str`, *optional*, defaults to `os.path.basename(__file__)[: -len(".py")]`): + Name of this experiment. + reward_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`): + Path to the reward model. + num_ppo_epochs (`int`, *optional*, defaults to `4`): + Number of epochs to train. + whiten_rewards (`bool`, *optional*, defaults to `False`): + Whether to whiten the rewards. + kl_coef (`float`, *optional*, defaults to `0.05`): + KL coefficient. + cliprange (`float`, *optional*, defaults to `0.2`): + Clip range. + rloo_k (`int`, *optional*, defaults to `2`): + REINFORCE Leave-One-Out (RLOO) number of online samples per prompt. + normalize_reward (`bool`, *optional*, defaults to `False`): + Whether to normalize rewards. + reward_clip_range (`float`, *optional*, defaults to `10.0`): + Clip range for rewards. + normalize_advantage (`bool`, *optional*, defaults to `False`): + Whether to normalize advantages. + token_level_kl (`bool`, *optional*, defaults to `True`): + Whether to use token-level KL penalty or sequence-level KL penalty. + ds3_gather_for_generation (`bool`, *optional*, defaults to `True`): + This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for generation, + improving generation speed. However, disabling this option allows training models that exceed the VRAM + capacity of a single GPU, albeit at the cost of slower generation. + """ + + exp_name: str = field( + default=os.path.basename(__file__)[:-3], + metadata={"help": "Name of this experiment."}, + ) + reward_model_path: str = field( + default="EleutherAI/pythia-160m", + metadata={"help": "Path to the reward model."}, + ) + num_ppo_epochs: int = field( + default=4, + metadata={"help": "Number of epochs to train."}, + ) + whiten_rewards: bool = field( + default=False, + metadata={"help": "Whether to whiten the rewards."}, + ) + kl_coef: float = field( + default=0.05, + metadata={"help": "KL coefficient."}, + ) + cliprange: float = field( + default=0.2, + metadata={"help": "Clip range."}, + ) + rloo_k: int = field( + default=2, + metadata={"help": "REINFORCE Leave-One-Out (RLOO) number of online samples per prompt."}, + ) + normalize_reward: bool = field( + default=False, + metadata={"help": "Whether to normalize rewards"}, + ) + reward_clip_range: float = field( + default=10.0, + metadata={"help": "Clip range for rewards"}, + ) + normalize_advantage: bool = field( + default=False, + metadata={"help": "Whether to normalize advantages"}, + ) + token_level_kl: bool = field( + default=False, + metadata={"help": "Whether to use token-level KL penalty or sequence-level KL penalty"}, + ) + ds3_gather_for_generation: bool = field( + default=True, + metadata={ + "help": "This setting applies to DeepSpeed ZeRO-3. If enabled, the policy model weights are gathered for " + "generation, improving generation speed. However, disabling this option allows training models that " + "exceed the VRAM capacity of a single GPU, albeit at the cost of slower generation." + }, + ) diff --git a/lib/python3.12/site-packages/trl/trainer/rloo_trainer.py b/lib/python3.12/site-packages/trl/trainer/rloo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..2c5d68ecce40b0f1065d9da33859a1ff2ebb2124 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/rloo_trainer.py @@ -0,0 +1,716 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import gc +import math +import os +import textwrap +import time +from collections import defaultdict +from pathlib import Path +from typing import Callable, Optional, Union + +import numpy as np +import pandas as pd +import torch +import torch.nn as nn +from accelerate import Accelerator +from accelerate.utils import broadcast, gather_object +from datasets import Dataset +from torch.utils.data import DataLoader +from transformers import ( + BaseImageProcessor, + DataCollatorWithPadding, + FeatureExtractionMixin, + GenerationConfig, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + TrainerCallback, + TrainerControl, + is_wandb_available, +) +from transformers.integrations import get_reporting_integration_callbacks +from transformers.trainer import DEFAULT_CALLBACKS, DEFAULT_PROGRESS_CALLBACK +from transformers.trainer_callback import CallbackHandler, ExportableState, PrinterCallback +from transformers.utils import is_rich_available + +from ..models.utils import unwrap_model_for_generation +from ..trainer.utils import ( + OnlineTrainerState, + batch_generation, + disable_dropout_in_model, + exact_div, + first_true_indices, + forward, + get_reward, + prepare_deepspeed, + print_rich_table, + selective_log_softmax, + truncate_response, +) +from .rloo_config import RLOOConfig +from .utils import empty_cache, generate_model_card, get_comet_experiment_url, log_table_to_comet_experiment + + +if is_wandb_available(): + import wandb + +INVALID_LOGPROB = 1.0 + + +class RLOOTrainer(Trainer): + _tag_names = ["trl", "rloo"] + + def __init__( + self, + config: RLOOConfig, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ], + policy: nn.Module, + ref_policy: nn.Module, + reward_model: Union[nn.Module, Callable[[list[str]], list[float]]], + train_dataset: Dataset, + data_collator: Optional[DataCollatorWithPadding] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + # less commonly used + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + callbacks: Optional[list[TrainerCallback]] = None, + ) -> None: + if ref_policy is policy: + raise ValueError( + "`policy` and `ref_policy` cannot be the same object. If you want `ref_policy` to be the " + "same as `policy`, you must mass a copy of it, or `None` if you use peft." + ) + + self.args = config + args = config + self.processing_class = processing_class + self.policy = policy + + # Define the collator if not provided + if data_collator is None: + data_collator = DataCollatorWithPadding(self.processing_class) + + self.policy.generation_config.eos_token_id = ( + None # disable `pad_token_id` and `eos_token_id` because we just want to + ) + self.policy.generation_config.pad_token_id = None # generate tokens without truncation / padding + + self.ref_policy = ref_policy + self.reward_model = reward_model + self.train_dataset = train_dataset + self.train_dataset_len = len(train_dataset) + self.data_collator = data_collator + self.eval_dataset = eval_dataset + self.optimizer, self.lr_scheduler = optimizers + self.optimizer_cls_and_kwargs = None # needed for transformers >= 4.47 + + ######### + # calculate various batch sizes + ######### + if args.total_episodes is None: # allow the users to define episodes in terms of epochs. + args.total_episodes = int(args.num_train_epochs * self.train_dataset_len) + accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps) + self.accelerator = accelerator + args.world_size = accelerator.num_processes + args.local_batch_size = ( + args.per_device_train_batch_size * args.gradient_accumulation_steps * args.num_mini_batches + ) + args.micro_batch_size = int(args.per_device_train_batch_size * args.world_size) + args.batch_size = int(args.local_batch_size * args.world_size) + args.mini_batch_size = exact_div( + args.batch_size, args.num_mini_batches, "`batch_size` must be a multiple of `num_mini_batches`" + ) + args.local_mini_batch_size = exact_div( + args.local_batch_size, args.num_mini_batches, "`local_batch_size` must be a multiple of `num_mini_batches`" + ) + args.num_total_batches = math.ceil( + args.total_episodes / args.batch_size + ) # we may train for more than `total_episodes` + time_tensor = torch.tensor(int(time.time()), device=accelerator.device) + time_int = broadcast(time_tensor, 0).item() # avoid different timestamps across processes + args.run_name = f"{args.exp_name}__{args.seed}__{time_int}" + self.local_seed = args.seed + accelerator.process_index * 100003 # Prime + if args.num_sample_generations > 0: + self.sample_generations_freq = max(1, args.num_total_batches // args.num_sample_generations) + self.local_dataloader_batch_size = exact_div( + args.local_batch_size, args.rloo_k, "`local_batch_size` must be a multiple of rloo_k" + ) # RLOO logic: needed because RLOO repeats the same prompt args.rloo_k times + + ######### + # setup model, optimizer, and others + ######### + for module in [policy, ref_policy, reward_model]: + if isinstance(module, nn.Module): + disable_dropout_in_model(module) + if args.stop_token and args.stop_token == "eos": + args.stop_token_id = self.processing_class.eos_token_id + self.model = policy + self.create_optimizer_and_scheduler( + num_training_steps=args.num_total_batches + ) # note that we are calling `self.lr_scheduler.step()` manually only at the batch level + + ######### + ### trainer specifics + ######### + default_callbacks = DEFAULT_CALLBACKS + get_reporting_integration_callbacks(self.args.report_to) + self.callbacks = default_callbacks if callbacks is None else default_callbacks + callbacks + self.callback_handler = CallbackHandler( + self.callbacks, self.model, self.processing_class, self.optimizer, self.lr_scheduler + ) + self.add_callback(PrinterCallback if self.args.disable_tqdm else DEFAULT_PROGRESS_CALLBACK) + self.control = TrainerControl() + self.state = OnlineTrainerState( + is_local_process_zero=self.is_local_process_zero(), + is_world_process_zero=self.is_world_process_zero(), + stateful_callbacks=[ + cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) + ], + ) + + self.current_flos = 0 + self.hp_search_backend = None + self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None + self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None + # Create distant repo and output directory if needed + self.hub_model_id = None + if self.args.push_to_hub: + self.init_hf_repo() + if self.args.should_save: + os.makedirs(self.args.output_dir, exist_ok=True) + self.backup_model = None + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + ######### + ### setup dataloader + ######### + self.dataloader = DataLoader( + self.train_dataset, + batch_size=self.local_dataloader_batch_size, + shuffle=True, + collate_fn=self.data_collator, + drop_last=True, # needed; otherwise the last batch will be of ragged shape + ) + # sync random states for DataLoader(shuffle=True) before `accelerator.prepare` + # see https://gist.github.com/vwxyzjn/2581bff1e48e185e0b85b6dfe1def79c + torch.manual_seed(args.seed) + self.model, self.optimizer, self.dataloader = accelerator.prepare(self.model, self.optimizer, self.dataloader) + torch.manual_seed(self.local_seed) # reset the local seed again + + self.eval_dataloader = DataLoader( + self.eval_dataset, + batch_size=args.per_device_eval_batch_size, + collate_fn=self.data_collator, + drop_last=True, + ) # no need to shuffle eval dataset + self.eval_dataloader = accelerator.prepare(self.eval_dataloader) + + if self.is_deepspeed_enabled: + if isinstance(self.reward_model, nn.Module): + self.reward_model = prepare_deepspeed( + self.reward_model, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + self.ref_policy = prepare_deepspeed( + self.ref_policy, args.per_device_train_batch_size, args.fp16, args.bf16 + ) + self.deepspeed = self.model + else: + self.ref_policy = self.ref_policy.to(self.accelerator.device) + if isinstance(self.reward_model, nn.Module): + self.reward_model = self.reward_model.to(self.accelerator.device) + + def get_train_dataloader(self) -> DataLoader: + return self.dataloader + + def get_eval_dataloader(self) -> DataLoader: + return self.eval_dataloader + + def train(self): + args = self.args + accelerator = self.accelerator + optimizer = self.optimizer + model = self.model + self.model_wrapped = self.model + ref_policy = self.ref_policy + reward_model = self.reward_model + processing_class = self.processing_class + dataloader = self.dataloader + device = accelerator.device + + def repeat_generator(): + while True: + yield from dataloader + + iter_dataloader = iter(repeat_generator()) + generation_config = GenerationConfig( + max_new_tokens=args.response_length, + temperature=(args.temperature + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + accelerator.print("===training policy===") + start_time = time.time() + stats_shape = (args.num_ppo_epochs, args.num_mini_batches, args.gradient_accumulation_steps) + approxkl_stats = torch.zeros(stats_shape, device=device) + pg_clipfrac_stats = torch.zeros(stats_shape, device=device) + pg_loss_stats = torch.zeros(stats_shape, device=device) + vf_clipfrac_stats = torch.zeros(stats_shape, device=device) + entropy_stats = torch.zeros(stats_shape, device=device) + ratio_stats = torch.zeros(stats_shape, device=device) + model.train() + + # trainer state initialization + self.state.global_step = 0 + self.state.episode = 0 + self.state.max_steps = (args.num_total_batches * args.num_mini_batches) // 2 + self.state.num_train_epochs = args.total_episodes / self.train_dataset_len + # Compute absolute values for logging, eval, and save if given as ratio + if args.logging_steps is not None: + if args.logging_steps < 1: + self.state.logging_steps = math.ceil(self.state.max_steps * args.logging_steps) + else: + self.state.logging_steps = args.logging_steps + if args.eval_steps is not None: + if args.eval_steps < 1: + self.state.eval_steps = math.ceil(self.state.max_steps * args.eval_steps) + else: + self.state.eval_steps = args.eval_steps + if args.save_steps is not None: + if args.save_steps < 1: + self.state.save_steps = math.ceil(self.state.max_steps * args.save_steps) + else: + self.state.save_steps = args.save_steps + self.control = self.callback_handler.on_train_begin(args, self.state, self.control) + + for update in range(1, args.num_total_batches + 1): + self.state.episode += 1 * args.batch_size + data = next(iter_dataloader) + with torch.no_grad(): + queries = data["input_ids"].to(device) + queries = queries.repeat(args.rloo_k, 1) + context_length = queries.shape[1] + responses = [] + postprocessed_responses = [] + logprobs = [] + ref_logprobs = [] + scores = [] + sequence_lengths = [] + + # Generate responses and compute logprobs + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + query_responses, logitss = batch_generation( + unwrapped_model, + queries, + args.local_rollout_forward_batch_size, + processing_class.pad_token_id, + generation_config, + ) + + # Process responses in batches + for i in range(0, queries.shape[0], args.local_rollout_forward_batch_size): + query = queries[i : i + args.local_rollout_forward_batch_size] + query_response = query_responses[i : i + args.local_rollout_forward_batch_size] + response = query_response[:, context_length:] + logits = logitss[i : i + args.local_rollout_forward_batch_size] + logprob = selective_log_softmax(logits, response) + del logits + empty_cache() + + ref_output = forward(ref_policy, query_response, processing_class.pad_token_id) + ref_logits = ref_output.logits[:, context_length - 1 : -1] + ref_logits /= args.temperature + 1e-7 + ref_logprob = selective_log_softmax(ref_logits, response) + del ref_output, ref_logits + empty_cache() + + # Response Processing 1. truncate response after the first occurrence of `stop_token_id` + postprocessed_response = response + if args.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + args.stop_token_id, processing_class.pad_token_id, response + ) + + # Response Processing 2. run reward model on the truncated responses + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + sequence_length = first_true_indices(postprocessed_response == processing_class.pad_token_id) - 1 + + if isinstance(reward_model, nn.Module): + _, score, _ = get_reward( + reward_model, postprocessed_query_response, processing_class.pad_token_id, context_length + ) + else: + score = torch.tensor( + reward_model( + processing_class.batch_decode(postprocessed_query_response, skip_special_tokens=True) + ), + dtype=torch.float, + ).to(device) + + # Store batch results + responses.append(response) + postprocessed_responses.append(postprocessed_response) + logprobs.append(logprob) + ref_logprobs.append(ref_logprob) + sequence_lengths.append(sequence_length) + scores.append(score) + + # Concatenate all batched results + responses = torch.cat(responses, 0) + postprocessed_responses = torch.cat(postprocessed_responses, 0) + logprobs = torch.cat(logprobs, 0) + ref_logprobs = torch.cat(ref_logprobs, 0) + sequence_lengths = torch.cat(sequence_lengths, 0) + scores = torch.cat(scores, 0) + del (logprob, ref_logprob, score) + empty_cache() + gc.collect() + + # Response Processing 3. filter response. Ensure that the sample contains stop_token_id + # responses not passing that filter will receive a low (fixed) score + # only query humans on responses that pass that filter + contain_eos_token = torch.any(postprocessed_responses == processing_class.eos_token_id, dim=-1) + if args.missing_eos_penalty is not None: + scores[~contain_eos_token] -= self.args.missing_eos_penalty + # accelerator.print(f"{scores=}, {(contain_eos_token.sum() / len(contain_eos_token))=}") + + # be very careful with `padding_mask_p1`; see https://excalidraw.com/#json=LWnzG4w2k5DjF_EOL_xPt,e2w3a-hFJ_gX5vOfeyXGTw + response_idxs = torch.arange(responses.shape[1], device=responses.device).repeat(responses.shape[0], 1) + padding_mask = response_idxs > sequence_lengths.unsqueeze(1) + logprobs = torch.masked_fill(logprobs, padding_mask, INVALID_LOGPROB) + ref_logprobs = torch.masked_fill(ref_logprobs, padding_mask, INVALID_LOGPROB) + + # 4. compute rewards + # Compute KL divergence + kl = logprobs - ref_logprobs + + # Normalize rewards + if args.normalize_reward: + scores = (scores - scores.mean()) / (scores.std() + 1e-8) + scores = torch.clamp(scores, -args.reward_clip_range, args.reward_clip_range) + + # Compute total reward with KL penalty + if args.token_level_kl: + # Token-level KL penalty: apply KL penalty per token + kl_reward = -args.kl_coef * kl + + # Get the index of the last non-padded token for each sequence + eos_indices = padding_mask.size(1) - 1 - padding_mask.long().fliplr().argmax(dim=1, keepdim=True) + last_reward = torch.zeros_like(kl) + # Ensure scores has correct shape and type + scores_shaped = scores.reshape(-1, 1).to(kl.dtype) + last_reward.scatter_(dim=1, index=eos_indices, src=scores_shaped) + + # Combine KL reward and last reward + non_score_reward = kl_reward.sum(1) # Keep this for logging + reward = last_reward + kl_reward + rlhf_reward = reward.sum(1) # Sum across sequence length + else: + # Sequence-level KL penalty: sum KL across tokens first + sequence_kl = kl.sum(1) + non_score_reward = -args.kl_coef * sequence_kl + rlhf_reward = non_score_reward + scores + + # vectorized RLOO advantages implementation + rlhf_reward = rlhf_reward.reshape(args.rloo_k, -1) + baseline = (rlhf_reward.sum(0) - rlhf_reward) / (args.rloo_k - 1) + advantages = rlhf_reward - baseline + advantages = advantages.flatten() + + # Normalize advantages + if args.normalize_advantage: + advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) + + empty_cache() + + # Do multiple epochs of PPO training, with a fresh random shuffle in each epoch + for ppo_epoch_idx in range(args.num_ppo_epochs): + b_inds = np.random.permutation(args.local_batch_size) + minibatch_idx = 0 + for mini_batch_start in range(0, args.local_batch_size, args.local_mini_batch_size): + mini_batch_end = mini_batch_start + args.local_mini_batch_size + mini_batch_inds = b_inds[mini_batch_start:mini_batch_end] + gradient_accumulation_idx = 0 + for micro_batch_start in range(0, args.local_mini_batch_size, args.per_device_train_batch_size): + with accelerator.accumulate(model): + micro_batch_end = micro_batch_start + args.per_device_train_batch_size + micro_batch_inds = mini_batch_inds[micro_batch_start:micro_batch_end] + + # Get batch data + mb_advantage = advantages[micro_batch_inds] + mb_responses = responses[micro_batch_inds] + mb_query_responses = query_responses[micro_batch_inds] + mb_logprobs = logprobs[micro_batch_inds] + + # Forward pass + output = forward(model, mb_query_responses, processing_class.pad_token_id) + logits = output.logits[:, context_length - 1 : -1] + logits /= args.temperature + 1e-7 + + # Compute new logprobs + new_logprobs = selective_log_softmax(logits, mb_responses) + new_logprobs = torch.masked_fill( + new_logprobs, padding_mask[micro_batch_inds], INVALID_LOGPROB + ) + + # Compute probability ratios + new_ratio = (new_logprobs - mb_logprobs).exp() + new_logprobs = new_logprobs.sum(1) + mb_logprobs = mb_logprobs.sum(1) + logprobs_diff = new_logprobs - mb_logprobs + ratio = torch.exp(logprobs_diff) + + # PPO clipped loss + pg_losses = -mb_advantage * ratio + pg_losses2 = -mb_advantage * torch.clamp(ratio, 1.0 - args.cliprange, 1.0 + args.cliprange) + pg_loss_max = torch.max(pg_losses, pg_losses2) + pg_loss = pg_loss_max.mean() + + # Final loss + loss = pg_loss + + # Optimization step + accelerator.backward(loss) + optimizer.step() + optimizer.zero_grad() + + with torch.no_grad(): + pg_clipfrac = (pg_losses2 > pg_losses).float().mean() + prob_dist = torch.nn.functional.softmax(logits, dim=-1) + entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1) + approxkl = 0.5 * (logprobs_diff**2).mean() + approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl + pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ( + pg_clipfrac + ) + pg_loss_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = pg_loss + entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean() + ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = new_ratio.mean() + gradient_accumulation_idx += 1 + minibatch_idx += 1 + + # del everything and empty cache + # fmt: off + del ( + output, logits, new_logprobs, logprobs_diff, ratio, pg_losses, + pg_losses2, pg_loss, loss, pg_clipfrac, prob_dist, entropy, approxkl, + mb_advantage, mb_responses, mb_query_responses, mb_logprobs, + ) + # fmt: on + empty_cache() + + # Compute metrics + with torch.no_grad(): + mean_kl = kl.sum(1).mean() + mean_entropy = (-logprobs).sum(1).mean() + mean_non_score_reward = non_score_reward.mean() + eps = int(self.state.episode / (time.time() - start_time)) + metrics = {} + metrics["eps"] = eps + metrics["objective/kl"] = self.accelerator.gather_for_metrics(mean_kl).mean().item() + metrics["objective/entropy"] = self.accelerator.gather_for_metrics(mean_entropy).mean().item() + metrics["objective/non_score_reward"] = ( + self.accelerator.gather_for_metrics(mean_non_score_reward).mean().item() + ) + metrics["objective/rlhf_reward"] = self.accelerator.gather_for_metrics(rlhf_reward).mean().item() + metrics["objective/scores"] = self.accelerator.gather_for_metrics(scores.mean()).mean().item() + metrics["policy/approxkl_avg"] = self.accelerator.gather_for_metrics(approxkl_stats).mean().item() + metrics["policy/clipfrac_avg"] = self.accelerator.gather_for_metrics(pg_clipfrac_stats).mean().item() + metrics["loss/policy_avg"] = self.accelerator.gather_for_metrics(pg_loss_stats).mean().item() + metrics["val/clipfrac_avg"] = self.accelerator.gather_for_metrics(vf_clipfrac_stats).mean().item() + metrics["policy/entropy_avg"] = self.accelerator.gather_for_metrics(entropy_stats).mean().item() + metrics["val/ratio"] = self.accelerator.gather_for_metrics(ratio_stats).mean().item() + metrics["val/ratio_var"] = self.accelerator.gather_for_metrics(ratio_stats).var().item() + metrics["val/num_eos_tokens"] = (responses == processing_class.eos_token_id).sum().item() + metrics["lr"] = self.lr_scheduler.get_last_lr()[0] + metrics["episode"] = self.state.episode + self.state.epoch = self.state.episode / (args.rloo_k * self.train_dataset_len) # used by self.log + self.log(metrics) + del kl, mean_kl, mean_entropy, scores + + self.lr_scheduler.step() + self.state.global_step += 1 + self.control = self.callback_handler.on_step_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + empty_cache() + gc.collect() + + if args.num_sample_generations > 0 and (update - 1) % self.sample_generations_freq == 0: + self.generate_completions(sampling=True) + + # HF trainer specifics + self.control = self.callback_handler.on_train_end(args, self.state, self.control) + if self.control.should_save: + self._save_checkpoint(model, trial=None, metrics=None) + self.control = self.callback_handler.on_save(self.args, self.state, self.control) + + def generate_completions(self, sampling: bool = False): + args = self.args + processing_class = self.processing_class + generation_config = GenerationConfig( + max_new_tokens=self.args.response_length, + temperature=(0.01 + 1e-7), + top_k=0.0, + top_p=1.0, + do_sample=True, + ) + + table = defaultdict(list) + with unwrap_model_for_generation( + self.model, self.accelerator, gather_deepspeed3_params=self.args.ds3_gather_for_generation + ) as unwrapped_model: + for batch in self.eval_dataloader: + query = batch["input_ids"] + with torch.no_grad(): + context_length = query.shape[1] + query_response, _ = batch_generation( + unwrapped_model, + query, + query.shape[0], + processing_class.pad_token_id, + generation_config, + ) + response = query_response[:, context_length:] + postprocessed_response = response + if args.stop_token_id is not None: # handle the edge case when stop_token_id exists but is 0 + postprocessed_response = truncate_response( + args.stop_token_id, processing_class.pad_token_id, response + ) + table["query"].extend( + gather_object(processing_class.batch_decode(query, skip_special_tokens=True)) + ) + table["model response"].extend( + gather_object(processing_class.batch_decode(postprocessed_response)) + ) + + postprocessed_query_response = torch.cat((query, postprocessed_response), 1) + + if isinstance(self.reward_model, nn.Module): + _, score, _ = get_reward( + self.reward_model, + postprocessed_query_response, + processing_class.pad_token_id, + context_length, + ) + else: + score = torch.tensor( + self.reward_model( + processing_class.batch_decode(postprocessed_query_response, skip_special_tokens=True) + ), + dtype=torch.float, + ).to(postprocessed_query_response.device) + table["score"].extend(self.accelerator.gather_for_metrics(score).float().cpu().numpy()) + + if sampling: + break + df = pd.DataFrame(table) + + if self.accelerator.is_main_process: + if is_rich_available(): + print_rich_table(df.iloc[0 : 0 + 5]) + if "wandb" in args.report_to: + import wandb + + if wandb.run is not None: + wandb.log({"completions": wandb.Table(dataframe=df)}) + + if "comet_ml" in args.report_to: + log_table_to_comet_experiment( + name="completions.csv", + table=df, + ) + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @inproceedings{ahmadian2024back, + title = {{Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs}}, + author = {Arash Ahmadian and Chris Cremer and Matthias Gall{\'{e}} and Marzieh Fadaee and Julia Kreutzer and Olivier Pietquin and Ahmet {\"{U}}st{\"{u}}n and Sara Hooker}, + year = 2024, + booktitle = {Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), {ACL} 2024, Bangkok, Thailand, August 11-16, 2024}, + publisher = {Association for Computational Linguistics}, + pages = {12248--12267}, + editor = {Lun{-}Wei Ku and Andre Martins and Vivek Srikumar}, + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="RLOO", + trainer_citation=citation, + paper_title="Back to Basics: Revisiting REINFORCE-Style Optimization for Learning from Human Feedback in LLMs", + paper_id="2402.14740", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/sft_config.py b/lib/python3.12/site-packages/trl/trainer/sft_config.py new file mode 100644 index 0000000000000000000000000000000000000000..983943835e54a5545a3dad2afe772e9bb0f02a66 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/sft_config.py @@ -0,0 +1,261 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import warnings +from dataclasses import dataclass, field +from typing import Any, Optional + +from transformers import TrainingArguments + + +@dataclass +class SFTConfig(TrainingArguments): + r""" + Configuration class for the [`SFTTrainer`]. + + This class includes only the parameters that are specific to SFT training. For a full list of training arguments, + please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this class may + differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + > Parameters that control the model + + model_init_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Keyword arguments for [`~transformers.AutoModelForCausalLM.from_pretrained`], used when the `model` + argument of the [`SFTTrainer`] is provided as a string. + chat_template_path (`str` or `None`, *optional*, defaults to `None`): + If specified, sets the model's chat template. This can either be the path to a tokenizer (local directory + or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, you must + ensure that any special tokens referenced in the template are added to the tokenizer and that the model's + embedding layer is resized accordingly. + + > Parameters that control the data preprocessing + + dataset_text_field (`str`, *optional*, defaults to `"text"`): + Name of the column that contains text data in the dataset. + dataset_kwargs (`dict[str, Any]` or `None`, *optional*, defaults to `None`): + Dictionary of optional keyword arguments for the dataset preparation. The only supported key is + `skip_prepare_dataset`. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + eos_token (`str` or `None`, *optional*, defaults to `None`): + Token used to indicate the end of a turn or sequence. If `None`, it defaults to + `processing_class.eos_token`. + pad_token (`int` or `None`, *optional*, defaults to `None`): + Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that is also `None`, + it falls back to `processing_class.eos_token`. + max_length (`int` or `None`, *optional*, defaults to `1024`): + Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from the right. + If `None`, no truncation is applied. When packing is enabled, this value sets the sequence length. + packing (`bool`, *optional*, defaults to `False`): + Whether to group multiple sequences into fixed-length blocks to improve computational efficiency and reduce + padding. Uses `max_length` to define sequence length. + packing_strategy (`str`, *optional*, defaults to `"ffd"`): + Strategy for packing sequences. Can be either `"ffd"` (first-fit decreasing, default), or `"wrapped"`. + padding_free (`bool`, *optional*, defaults to `False`): + Whether to perform forward passes without padding by flattening all sequences in the batch into a single + continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, this is only + supported with the `flash_attention_2` attention implementation, which can efficiently handle the flattened + batch structure. When packing is enabled with strategy `"ffd"`, padding-free is enabled, regardless of the + value of this parameter. + pad_to_multiple_of (`int` or `None`, *optional*, defaults to `None`): + If set, the sequences will be padded to a multiple of this value. + eval_packing (`bool` or `None`, *optional*, defaults to `None`): + Whether to pack the eval dataset. If `None`, uses the same value as `packing`. + + > Parameters that control the training + + completion_only_loss (`bool` or `None`, *optional*, defaults to `None`): + Whether to compute loss only on the completion part of the sequence. If set to `True`, loss is computed + only on the completion, which is supported only for [prompt-completion](#prompt-completion) datasets. If + `False`, loss is computed on the entire sequence. If `None` (default), the behavior depends on the dataset: + loss is computed on the completion for [prompt-completion](#prompt-completion) datasets, and on the full + sequence for [language modeling](#language-modeling) datasets. + assistant_only_loss (`bool`, *optional*, defaults to `False`): + Whether to compute loss only on the assistant part of the sequence. If set to `True`, loss is computed + only on the assistant responses, which is supported only for [conversational](#conversational) datasets. If `False`, + loss is computed on the entire sequence. + activation_offloading (`bool`, *optional*, defaults to `False`): + Whether to offload the activations to the CPU. + """ + + _VALID_DICT_FIELDS = TrainingArguments._VALID_DICT_FIELDS + ["model_init_kwargs"] + + # Parameters whose default values are overridden from TrainingArguments + learning_rate: float = field( + default=2e-5, + metadata={"help": "The initial learning rate for AdamW."}, + ) + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + average_tokens_across_devices: bool = field( + default=True, + metadata={ + "help": "Whether or not to average tokens across devices. If enabled, will use all_reduce to synchronize " + "num_tokens_in_batch for precise loss calculation. Reference: https://github.com/huggingface/transformers/issues/34242 " + }, + ) + + # Parameters that control the model + model_init_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of " + "the `SFTTrainer` is provided as a string." + }, + ) + chat_template_path: Optional[str] = field( + default=None, + metadata={ + "help": "If specified, sets the model's chat template. This can either be the path to a tokenizer (local " + "directory or Hugging Face Hub model) or a direct path to a Jinja template file. When using a Jinja file, " + "you must ensure that any special tokens referenced in the template are added to the tokenizer and " + "that the model's embedding layer is resized accordingly." + }, + ) + + # Parameters that control the data preprocessing + dataset_text_field: str = field( + default="text", + metadata={"help": "Name of the column that contains text data in the dataset."}, + ) + dataset_kwargs: Optional[dict[str, Any]] = field( + default=None, + metadata={ + "help": "Dictionary of optional keyword arguments for the dataset preparation. The only supported key is " + "`skip_prepare_dataset`." + }, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + eos_token: Optional[str] = field( + default=None, + metadata={ + "help": "Token used to indicate the end of a turn or sequence. If `None`, it defaults to `processing_class.eos_token`." + }, + ) + pad_token: Optional[str] = field( + default=None, + metadata={ + "help": "Token used for padding. If `None`, it defaults to `processing_class.pad_token`, or if that " + "is also `None`, it falls back to `processing_class.eos_token`." + }, + ) + max_length: Optional[int] = field( + default=1024, + metadata={ + "help": "Maximum length of the tokenized sequence. Sequences longer than `max_length` are truncated from" + "the right. If `None`, no truncation is applied. When packing is enabled, this value sets the " + "sequence length." + }, + ) + packing: bool = field( + default=False, + metadata={ + "help": "Whether to group multiple sequences into fixed-length blocks to improve computational efficiency " + "and reduce padding. Uses `max_length` to define sequence length." + }, + ) + packing_strategy: str = field( + default="ffd", + metadata={ + "help": "Strategy for packing sequences. Can be either `'ffd'` (first-fit decreasing, default), or " + "`'wrapped'`." + }, + ) + padding_free: bool = field( + default=False, + metadata={ + "help": "Whether to perform forward passes without padding by flattening all sequences in the batch into " + "a single continuous sequence. This reduces memory usage by eliminating padding overhead. Currently, " + "this is only supported with the `flash_attention_2` attention implementation, which can efficiently " + "handle the flattened batch structure. When packing is enabled with strategy `'ffd'`, padding-free is " + "enabled, regardless of the value of this parameter." + }, + ) + pad_to_multiple_of: Optional[int] = field( + default=None, + metadata={"help": "If set, the sequences will be padded to a multiple of this value."}, + ) + eval_packing: Optional[bool] = field( + default=None, + metadata={"help": "Whether to pack the eval dataset. If `None`, uses the same value as `packing`."}, + ) + + # Parameters that control the training + completion_only_loss: Optional[bool] = field( + default=None, + metadata={ + "help": ( + "Whether to compute loss only on the completion part of the sequence. If set to `True`, loss is " + "computed only on the completion, which is supported only for prompt-completion datasets. If `False`, " + "loss is computed on the entire sequence. If `None` (default), the behavior depends on the dataset: " + "loss is computed on the completion for prompt-completion datasets, and on the full sequence for " + "language modeling datasets." + ) + }, + ) + assistant_only_loss: bool = field( + default=False, + metadata={ + "help": ( + "Whether to compute loss only on the assistant part of the sequence. If set to `True`, loss is " + "computed only on the assistant responses, which is supported only for conversational datasets. If `False`, " + "loss is computed on the entire sequence." + ) + }, + ) + activation_offloading: bool = field( + default=False, + metadata={"help": "Whether to offload the activations to the CPU."}, + ) + + # Deprecated parameters + max_seq_length: Optional[int] = field( + default=None, + metadata={ + "help": "This parameter is deprecated and will be removed in version 0.20.0. Use `max_length` instead." + }, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + if self.max_seq_length is not None: + warnings.warn( + "`max_seq_length` is deprecated and will be removed in version 0.20.0. Use `max_length` instead.", + DeprecationWarning, + ) + self.max_length = self.max_seq_length diff --git a/lib/python3.12/site-packages/trl/trainer/sft_trainer.py b/lib/python3.12/site-packages/trl/trainer/sft_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..3250cc56aea7e8717f986c1fbb8f3ae1680ae0f3 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/sft_trainer.py @@ -0,0 +1,937 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextlib +import dataclasses +import os +import warnings +from collections import defaultdict +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Optional, TypeVar, Union + +import torch +import torch.nn as nn +from accelerate import PartialState +from datasets import Dataset, IterableDataset +from packaging import version +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BaseImageProcessor, + DataCollator, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + Trainer, + TrainingArguments, + is_wandb_available, +) +from transformers.data.data_collator import DataCollatorMixin +from transformers.trainer_callback import TrainerCallback +from transformers.trainer_utils import EvalPrediction +from transformers.utils import is_peft_available + +from ..data_utils import ( + is_conversational, + is_conversational_from_value, + maybe_convert_to_chatml, + pack_dataset, + truncate_dataset, +) +from ..models import clone_chat_template, get_act_offloading_ctx_manager +from .sft_config import SFTConfig +from .utils import ( + ConstantLengthDataset, + generate_model_card, + get_comet_experiment_url, + pad, + peft_module_casting_to_bf16, +) + + +if is_peft_available(): + import peft + from peft import PeftConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training + +if is_wandb_available(): + import wandb + + +TListOrMapping = TypeVar("TListOrMapping", list, Mapping) + + +def remove_none_values(example: TListOrMapping) -> TListOrMapping: + """ + Recursively removes entries with `None` values from a nested structure (list or dictionary). + + Args: + example (`list` or `Mapping`): + Input nested structure (list or dictionary) from which to remove `None`. + + Example: + ```python + >>> [{ + ... "a": {"aa": None, + ... "ab": 1}, + ... "b": "my_string", + ... }] + >>> remove_none_values(example) + [{'a': {'ab': 1}, 'b': 'my_string'}] + ``` + """ + if isinstance(example, list): + return [remove_none_values(value) if isinstance(value, (dict, list)) else value for value in example] + elif isinstance(example, Mapping): + return { + key: remove_none_values(value) if isinstance(value, (dict, list)) else value + for key, value in example.items() + if value is not None + } + else: + raise TypeError("Input must be a list or a dictionary.") + + +@dataclass +class DataCollatorForLanguageModeling(DataCollatorMixin): + """ + Data collator used for language modeling data. Inputs are dynamically padded to the maximum length of a batch. + + This collator expects each example in the input list to be a dictionary containing at least the `"input_ids"` key. + If the input contains a `"completion_mask"`, it is used to set the labels to `-100` for tokens that are not in the + completion. If `"assistant_masks"` are present, they are used to set the labels to `-100` for tokens that are not + in the assistant part of the sequence. The collator returns a dictionary containing the following keys: + - `"input_ids"`: Tensor of input IDs, padded to the maximum length of the batch. + - `"attention_mask"`: Tensor of attention mask, padded to the maximum length of the batch. + - `"position_ids"`: Tensor of position IDs, padded to the maximum length of the batch. + - `"labels"`: Tensor of labels, padded to the maximum length of the batch. If `completion_only_loss` is set to + `True`, tokens that are not in the completion are set to -100. If `assistant_masks` are present, tokens that are + not in the assistant part of the sequence are set to -100. + + Args: + pad_token_id (`int`): + Token ID to use for padding. + completion_only_loss (`bool`, *optional*, defaults to `True`): + When the input contains a completion mask (`completion_mask`), the labels are set to -100 for the tokens + that are no in the completion. + padding_free (`bool`, *optional*, defaults to `False`): + If set to `True`, the sequences will be flattened into a single sequence, and the position IDs will be + generated accordingly. The attention mask will be set to 1 for all tokens. + pad_to_multiple_of (`int` or `None`, *optional*, defaults to `None`): + If set, the sequences will be padded to a multiple of this value. + return_tensors (`str`, *optional*, defaults to `"pt"`): + Type of Tensor to return. Only `"pt"` is currently supported. + + Examples: + ```python + >>> from trl import DataCollatorForLanguageModeling + + >>> collator = DataCollatorForLanguageModeling(pad_token_id=0) + >>> examples = [{"input_ids": [1, 2, 3]}, {"input_ids": [4, 5]}] + >>> collator(examples) + {'input_ids': tensor([[ 1, 2, 3], + [ 4, 5, 0]]), + 'attention_mask': tensor([[ 1, 1, 1], + [ 1, 1, 0]]), + 'position_ids': tensor([[0, 1, 2], + [0, 1, 0]]), + 'labels': tensor([[ 1, 2, 3], + [ 4, 5, -100]])} + + >>> # With completion mask + >>> examples = [ + ... {"input_ids": [1, 2, 3], "completion_mask": [0, 1, 1]}, + ... {"input_ids": [4, 5], "completion_mask": [0, 1]}, + ... ] + >>> collator(examples) + {'input_ids': tensor([[ 1, 2, 3], + [ 4, 5, 0]]), + 'attention_mask': tensor([[ 1, 1, 1], + [ 1, 1, 0]]), + 'position_ids': tensor([[0, 1, 2], + [0, 1, 0]]), + 'labels': tensor([[-100, 2, 3], + [-100, 5, -100]])} + + >>> # With padding_free + >>> collator = DataCollatorForLanguageModeling(pad_token_id=0, padding_free=True) + >>> collator(examples) + {'input_ids': tensor([[ 1, 2, 3, 4, 5]]), + 'attention_mask': tensor([[1, 1, 1, 1, 1]]), + 'position_ids': tensor([[0, 1, 2, 0, 1]]), + 'labels': tensor([[1, 2, 3, 4, 5]])} + ``` + """ + + pad_token_id: int + completion_only_loss: bool = True + padding_free: bool = False + return_position_ids: bool = True + pad_to_multiple_of: Optional[int] = None + return_tensors: str = "pt" + + def torch_call(self, examples: list[Union[list[int], Any, dict[str, Any]]]) -> dict[str, Any]: + # Convert to tensor + input_ids = [torch.tensor(example["input_ids"]) for example in examples] + attention_mask = [torch.ones_like(input_ids) for input_ids in input_ids] + if self.return_position_ids: + if "position_ids" in examples[0]: + position_ids = [torch.tensor(example["position_ids"]) for example in examples] + else: + position_ids = [torch.arange(len(ids)) for ids in input_ids] + if "labels" in examples[0]: + labels = [torch.tensor(example["labels"]) for example in examples] + else: + labels = [torch.tensor(example["input_ids"]) for example in examples] + if self.completion_only_loss and "completion_mask" in examples[0]: + completion_mask = [torch.tensor(example["completion_mask"]) for example in examples] + if "assistant_masks" in examples[0]: + assistant_masks = [torch.tensor(example["assistant_masks"]) for example in examples] + + # Pad + output = {} + if self.padding_free: + output["input_ids"] = torch.cat(input_ids, dim=0).unsqueeze(0) + output["attention_mask"] = torch.cat(attention_mask, dim=0).unsqueeze(0) + if self.return_position_ids: + output["position_ids"] = torch.cat(position_ids, dim=0).unsqueeze(0) + output["labels"] = torch.cat(labels, dim=0).unsqueeze(0) + if self.completion_only_loss and "completion_mask" in examples[0]: + completion_mask = torch.cat(completion_mask, dim=0).unsqueeze(0) + output["labels"][completion_mask == 0] = -100 + if "assistant_masks" in examples[0]: + assistant_masks = torch.cat(assistant_masks, dim=0).unsqueeze(0) + output["labels"][assistant_masks == 0] = -100 + + else: + output["input_ids"] = pad( + input_ids, + padding_value=self.pad_token_id, + padding_side="right", + pad_to_multiple_of=self.pad_to_multiple_of, + ) + output["attention_mask"] = pad( + attention_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + if self.return_position_ids: + output["position_ids"] = pad( + position_ids, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + output["labels"] = pad( + labels, padding_value=-100, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + if self.completion_only_loss and "completion_mask" in examples[0]: + completion_mask = pad( + completion_mask, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + output["labels"][completion_mask == 0] = -100 # mask everything that is not in the completion + if "assistant_masks" in examples[0]: + assistant_masks = pad( + assistant_masks, padding_value=0, padding_side="right", pad_to_multiple_of=self.pad_to_multiple_of + ) + output["labels"][assistant_masks == 0] = -100 + return output + + +class SFTTrainer(Trainer): + """ + Trainer for Supervised Fine-Tuning (SFT) method. + + This class is a wrapper around the [`transformers.Trainer`] class and inherits all of its attributes and methods. + + Example: + + ```python + from datasets import load_dataset + from trl import SFTTrainer + + dataset = load_dataset("roneneldan/TinyStories", split="train[:1%]") + + trainer = SFTTrainer(model="Qwen/Qwen2-0.5B-Instruct", train_dataset=dataset) + trainer.train() + ``` + + Args: + model (`Union[str, PreTrainedModel]`): + Model to be trained. Can be either: + + - A string, being the *model id* of a pretrained model hosted inside a model repo on huggingface.co, or a + path to a *directory* containing model weights saved using + [`~transformers.PreTrainedModel.save_pretrained`], e.g., `'./my_model_directory/'`. The model is loaded + using [`~transformers.AutoModelForCausalLM.from_pretrained`] with the keyword arguments in + `args.model_init_kwargs`. + - A [`~transformers.PreTrainedModel`] object. Only causal language models are supported. + args ([`SFTConfig`], *optional*, defaults to `None`): + Configuration for this trainer. If `None`, a default configuration is used. + data_collator (`DataCollator`, *optional*): + Function to use to form a batch from a list of elements of the processed `train_dataset` or `eval_dataset`. + Will default to a custom [`DataCollatorForLanguageModeling`]. + train_dataset ([`~datasets.Dataset`] or [`~datasets.IterableDataset`]): + Dataset to use for training. SFT supports both [language modeling](#language-modeling) type and + [prompt-completion](#prompt-completion) type. The format of the samples can be either: + + - [Standard](dataset_formats#standard): Each sample contains plain text. + - [Conversational](dataset_formats#conversational): Each sample contains structured messages (e.g., role + and content). + + The trainer also supports processed datasets (tokenized) as long as they contain an `input_ids` field. + eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`): + Dataset to use for evaluation. It must meet the same requirements as `train_dataset`. + processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`): + Processing class used to process the data. If `None`, the processing class is loaded from the model's name + with [`~transformers.AutoTokenizer.from_pretrained`]. + callbacks (list of [`~transformers.TrainerCallback`], *optional*, defaults to `None`): + List of callbacks to customize the training loop. Will add those to the list of default callbacks detailed + in [here](https://huggingface.co/docs/transformers/main_classes/callback). + + If you want to remove one of the default callbacks used, use the [`~transformers.Trainer.remove_callback`] + method. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`, *optional*, defaults to `(None, None)`): + A tuple containing the optimizer and the scheduler to use. Will default to an instance of [`AdamW`] on your + model and a scheduler given by [`get_linear_schedule_with_warmup`] controlled by `args`. + optimizer_cls_and_kwargs (`Tuple[Type[torch.optim.Optimizer], Dict[str, Any]]`, *optional*, defaults to `None`): + A tuple containing the optimizer class and keyword arguments to use. Overrides `optim` and `optim_args` in + `args`. Incompatible with the `optimizers` argument. + + Unlike `optimizers`, this argument avoids the need to place model parameters on the correct devices before + initializing the Trainer. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`, *optional*, defaults to `None`): + A function that preprocess the logits right before caching them at each evaluation step. Must take two + tensors, the logits and the labels, and return the logits once processed as desired. The modifications made + by this function will be reflected in the predictions received by `compute_metrics`. + + Note that the labels (second parameter) will be `None` if the dataset does not have them. + peft_config ([`~peft.PeftConfig`], *optional*, defaults to `None`): + PEFT configuration used to wrap the model. If `None`, the model is not wrapped. + formatting_func (`Optional[Callable]`): + Formatting function applied to the dataset before tokenization. Applying the formatting function explicitly + converts the dataset into a [language modeling](#language-modeling) type. + """ + + _tag_names = ["trl", "sft"] + + def __init__( + self, + model: Union[str, nn.Module, PreTrainedModel], + args: Optional[Union[SFTConfig, TrainingArguments]] = None, + data_collator: Optional[DataCollator] = None, # type: ignore + train_dataset: Optional[Union[Dataset, IterableDataset]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + compute_loss_func: Optional[Callable] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None), + optimizer_cls_and_kwargs: Optional[tuple[type[torch.optim.Optimizer], dict[str, Any]]] = None, + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + peft_config: Optional["PeftConfig"] = None, + formatting_func: Optional[Callable[[dict], str]] = None, + ): + # Args + model_id = model if isinstance(model, str) else model.config._name_or_path + if args is None: + model_name = model_id.split("/")[-1] + args = SFTConfig(f"{model_name}-SFT") + elif isinstance(args, TrainingArguments) and not isinstance(args, SFTConfig): + dict_args = args.to_dict() + dict_args["hub_token"] = args.hub_token # to_dict hides the hub_token + dict_args.pop("push_to_hub_token") + args = SFTConfig(**dict_args) + + # Handle the tokenizer + if processing_class is None: + processing_class = AutoTokenizer.from_pretrained(model_id) + + if args.eos_token is not None: + eos_token = args.eos_token + eos_token_id = processing_class.convert_tokens_to_ids(eos_token) + if eos_token_id is None: + raise ValueError( + f"The specified `eos_token` ('{eos_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `eos_token` exists " + "in the vocabulary before using it as an EOS token." + ) + processing_class.eos_token_id = eos_token_id + + # Model + if args.model_init_kwargs is not None and not isinstance(model, str): + warnings.warn( + "You passed model_init_kwargs to the `SFTConfig`, but your model is already instantiated. " + "The `model_init_kwargs` will be ignored." + ) + if isinstance(model, str): + model = self._create_model_from_path(model, args) + + if args.chat_template_path is not None: + if os.path.isfile(args.chat_template_path) and args.chat_template_path.endswith((".jinja", ".j2")): + with open(args.chat_template_path, encoding="utf-8") as chat_template_file: + processing_class.chat_template = chat_template_file.read() + else: + model, processing_class = clone_chat_template(model, processing_class, args.chat_template_path) + + # PEFT configuration and model wrapping + if peft_config is not None: + model = self._prepare_peft_model(model, peft_config, args) + + # Data collator + # FFD packing requires padding-free mode; otherwise, the collator outputs padded attention masks, causing + # FlashAttention to ignore position_ids and recompute them incorrectly from the padded attention mask. + self.padding_free = args.padding_free or (args.packing and args.packing_strategy == "ffd") + if self.padding_free: + if data_collator is not None: + raise ValueError("Passing a custom data collator is not supported when using padding-free.") + if args.packing and args.packing_strategy == "wrapped": + warnings.warn( + "You are passing `padding_free=True` with the 'wrapped' packing strategy, which is not " + "recommended. Please refer to the documentation to understand why this is not recommended." + ) + if model.config._attn_implementation != "flash_attention_2": + warnings.warn( + "Padding-free training is enabled, but the attention implementation is not set to " + "'flash_attention_2'. Padding-free training flattens batches into a single sequence, and " + "'flash_attention_2' is the only known attention mechanism that reliably supports this. Using " + "other implementations may lead to unexpected behavior. To ensure compatibility, set " + "`attn_implementation='flash_attention_2'` in the model configuration, or verify that your " + "attention mechanism can handle flattened sequences." + ) + if args.per_device_train_batch_size == 1 and not args.packing: + warnings.warn( + "You are using a per_device_train_batch_size of 1 with padding-free training. Using a batch size " + "of 1 anihilate the benefits of padding-free training. Please consider increasing the batch size " + "to at least 2." + ) + + if args.completion_only_loss is None: + first_example = next(iter(train_dataset)) + self.completion_only_loss = "prompt" in first_example + else: + self.completion_only_loss = args.completion_only_loss + + if data_collator is None: + # Get the pad token: if not provided, use the one from the processing class or the eos token + # if the processing class does not have a pad token. + pad_token = args.pad_token or processing_class.pad_token or processing_class.eos_token + pad_token_id = processing_class.convert_tokens_to_ids(pad_token) + if pad_token_id is None: + raise ValueError( + f"The specified `pad_token` ('{pad_token}') is not found in the vocabulary of the given " + f"`processing_class` ({processing_class.__class__.__name__}). Ensure that the `pad_token` exists " + "in the vocabulary before using it as a padding token." + ) + data_collator = DataCollatorForLanguageModeling( + pad_token_id=pad_token_id, + completion_only_loss=self.completion_only_loss, + padding_free=self.padding_free, + # Using position_ids without flash_attn hurts the training + return_position_ids=model.config._attn_implementation == "flash_attention_2", + pad_to_multiple_of=args.pad_to_multiple_of, + ) + + if ( + args.packing + and args.packing_strategy == "ffd" + and model.config._attn_implementation != "flash_attention_2" + ): + warnings.warn( + "You are using packing, but the attention implementation is not set to 'flash_attention_2'. Packing " + "flattens batches into a single sequence, and 'flash_attention_2' is the only known attention " + "mechanism that reliably supports this. Using other implementations may lead to cross-contamination " + "between batches. To avoid this, either disable packing by setting `packing=False`, or set " + "`attn_implementation='flash_attention_2'` in the model configuration." + ) + if args.assistant_only_loss and not is_conversational(train_dataset[0]): + raise ValueError( + "You set `assistant_only_loss=True`, but the dataset is not conversational. This option is only " + "supported for conversational datasets." + ) + + # Dataset + preprocess_dataset = args.dataset_kwargs is None or not args.dataset_kwargs.get("skip_prepare_dataset", False) + if preprocess_dataset: + if self.completion_only_loss and formatting_func: + raise ValueError( + "A formatting function was provided while `completion_only_loss=True`, which is incompatible. " + "Using a formatter converts the dataset to a language modeling type, conflicting with " + "completion-only loss. To resolve this, apply your formatting function before passing the " + "dataset, or disable `completion_only_loss` in `SFTConfig`." + ) + train_dataset = self._prepare_dataset( + train_dataset, processing_class, args, args.packing, formatting_func, "train" + ) + if eval_dataset is not None: + packing = args.packing if args.eval_packing is None else args.eval_packing + if isinstance(eval_dataset, dict): + eval_dataset = { + key: self._prepare_dataset(dataset, processing_class, args, packing, formatting_func, key) + for key, dataset in eval_dataset.items() + } + else: + eval_dataset = self._prepare_dataset( + eval_dataset, processing_class, args, packing, formatting_func, "eval" + ) + + # Initialize the metrics + self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} + self._total_train_tokens = 0 + + # Initialize the Trainer. Parent class will handle: + # - DeepSpeed configuration (through create_accelerator_and_postprocess) + # - FSDP setup + # - Distributed training setup + # - Optimizer and scheduler creation + + super().__init__( + model=model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + compute_loss_func=compute_loss_func, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + optimizer_cls_and_kwargs=optimizer_cls_and_kwargs, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + # Initialize activation offloading context + if self.args.activation_offloading: + self.maybe_activation_offload_context = get_act_offloading_ctx_manager(model=self.model) + else: + self.maybe_activation_offload_context = contextlib.nullcontext() + + # Add tags for models that have been loaded with the correct transformers version + if hasattr(self.model, "add_model_tags"): + self.model.add_model_tags(self._tag_names) + + def _create_model_from_path(self, model_path: str, args: SFTConfig) -> PreTrainedModel: + """Creates a model from a path or model identifier.""" + model_init_kwargs = args.model_init_kwargs or {} + # Handle torch dtype + torch_dtype = model_init_kwargs.get("torch_dtype") + if isinstance(torch_dtype, torch.dtype) or torch_dtype == "auto" or torch_dtype is None: + pass # torch_dtype is already a torch.dtype or "auto" or None + elif isinstance(torch_dtype, str): # it's a str, but not "auto" + torch_dtype = getattr(torch, torch_dtype) + model_init_kwargs["torch_dtype"] = torch_dtype + else: + raise ValueError( + "Invalid `torch_dtype` passed to `SFTConfig`. Expected either 'auto' or a string representing " + f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}." + ) + # Disable caching if gradient checkpointing is enabled (not supported) + # if args.gradient_checkpointing: + # model_init_kwargs["use_cache"] = False + + # Create model + model = AutoModelForCausalLM.from_pretrained(model_path, **model_init_kwargs) + return model + + def _prepare_peft_model(self, model: PreTrainedModel, peft_config: Any, args: SFTConfig) -> PreTrainedModel: + """Prepares a model for PEFT training.""" + if not is_peft_available(): + raise ImportError("To use PeftModel, you need to install the `peft` library.") + + if not isinstance(peft_config, PeftConfig): + raise ValueError( + f"Expected PeftConfig object but got {type(peft_config)}. If you want to use the PeftModel, you need " + "to pass a PeftConfig object to the SFTTrainer." + ) + + if isinstance(model, PeftModel): + return model + + # Handle quantized models (QLoRA) + is_qlora = getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False) + + is_sharded_qlora = False + if getattr(model, "is_loaded_in_4bit", False): + # Check if model is sharded (FSDP/DS-Zero3) + for _, param in model.named_parameters(): + if param.__class__.__name__ == "Params4bit": + is_sharded_qlora = param.data.device.type in {"cpu", "meta"} + break + + # Prepare model for kbit training if needed + if is_qlora and not is_sharded_qlora: + model = self._prepare_model_for_kbit_training(model, args) + # Disable gradient checkpointing as it's handled by prepare_model_for_kbit_training + args = dataclasses.replace(args, gradient_checkpointing=False) + elif args.gradient_checkpointing: + model = self._enable_gradient_checkpointing(model, args) + + # Create PEFT model + if ( + version.parse(peft.__version__) >= version.parse("0.12") # autocast_adapter_dtype introduced in 0.12 + and getattr(model, "is_loaded_in_4bit", False) + and is_sharded_qlora + ): + model = get_peft_model(model, peft_config, autocast_adapter_dtype=False) + else: + model = get_peft_model(model, peft_config) + + # Handle bf16 casting for 4-bit models + if args.bf16 and getattr(model, "is_loaded_in_4bit", False) and not is_sharded_qlora: + peft_module_casting_to_bf16(model) + + return model + + def _prepare_model_for_kbit_training(self, model: PreTrainedModel, args: SFTConfig) -> PreTrainedModel: + """Prepares a quantized model for kbit training.""" + prepare_model_kwargs = { + "use_gradient_checkpointing": args.gradient_checkpointing, + "gradient_checkpointing_kwargs": args.gradient_checkpointing_kwargs or {}, + } + + return prepare_model_for_kbit_training(model, **prepare_model_kwargs) + + def _enable_gradient_checkpointing(self, model: PreTrainedModel, args: SFTConfig) -> PreTrainedModel: + """Enables gradient checkpointing for the model.""" + gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs or {} + use_reentrant = ( + "use_reentrant" not in gradient_checkpointing_kwargs or gradient_checkpointing_kwargs["use_reentrant"] + ) + + if use_reentrant: + if hasattr(model, "enable_input_require_grads"): + model.enable_input_require_grads() + else: + + def make_inputs_require_grad(module, input, output): + output.requires_grad_(True) + + model.get_input_embeddings().register_forward_hook(make_inputs_require_grad) + + return model + + def _prepare_dataset( + self, + dataset: Union[Dataset, IterableDataset], + processing_class: Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin], + args: SFTConfig, + packing: bool, + formatting_func: Optional[Callable[[dict], str]], + dataset_name: str, + ) -> Union[Dataset, IterableDataset]: + # Convert the dataset to an IterableDataset if it is a ConstantLengthDataset + if isinstance(dataset, ConstantLengthDataset): + return dataset + + # Tabular backends like Arrow/Parquet insert `None` for mismatched keys in nested structures. Clean them from + # sampled data. + if isinstance(dataset, Dataset): # IterableDataset does not support `with_transform` + dataset = dataset.with_transform(remove_none_values) + + # If the dataset is already preprocessed (tokenized), skip the processing steps. + column_names = list(next(iter(dataset)).keys()) + is_processed = "input_ids" in column_names + + # Build the kwargs for the `map` function + map_kwargs = {} + if isinstance(dataset, Dataset): # IterableDataset does not support num_proc + map_kwargs["num_proc"] = args.dataset_num_proc + + with PartialState().main_process_first(): + # Apply the formatting function if any + if formatting_func is not None and is_processed: + warnings.warn( + "You passed a dataset that is already processed (contains an `input_ids` field) together with a " + "formatting function. Therefore `formatting_func` will be ignored. Either remove the " + "`formatting_func` or pass a dataset that is not already processed.", + UserWarning, + ) + + if formatting_func is not None and not is_processed: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Applying formatting function to {dataset_name} dataset" + + def _func(example): + return {"text": formatting_func(example)} + + try: + dataset = dataset.map(_func, batched=False, **map_kwargs) + except Exception as e: + warnings.warn( + f"Failed to apply the formatting function due to the following error: {e}. This may be " + "because the function is designed for batched input. Please update it to process one example " + "at a time (i.e., accept and return a single example). For now, we will attempt to apply the " + "function in batched mode, but note that batched formatting is deprecated and will be removed " + "in version 0.21.", + DeprecationWarning, + ) + dataset = dataset.map(_func, batched=True, **map_kwargs) + + if not is_processed: + # Convert the dataset to ChatML if needed + first_example = next(iter(dataset)) + if is_conversational_from_value(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Converting {dataset_name} dataset to ChatML" + column_names = next(iter(dataset)).keys() + dataset = dataset.map( + maybe_convert_to_chatml, + remove_columns="conversations" if "conversations" in column_names else None, + **map_kwargs, + ) + + # Apply the chat template if needed + first_example = next(iter(dataset)) + if not is_conversational(first_example): + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Adding EOS to {dataset_name} dataset" + + def add_eos(example, eos_token): + if "text" in example and not example["text"].endswith(eos_token): # language modeling case + example["text"] = example["text"] + eos_token + elif "completion" in example and not example["completion"].endswith(eos_token): + example["completion"] = example["completion"] + eos_token + return example + + dataset = dataset.map( + add_eos, + fn_kwargs={"eos_token": processing_class.eos_token}, + remove_columns="messages" if "messages" in column_names else None, # renamed to "text" + **map_kwargs, + ) + + # Tokenize the dataset + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset" + + def tokenize(example, processing_class, dataset_text_field, assistant_only_loss): + if "prompt" in example: # prompt-completion case + if is_conversational(example): + prompt_ids = processing_class.apply_chat_template( + example["prompt"], + tools=example.get("tools"), + **example.get("chat_template_kwargs", {}), + ) + prompt_completion_ids = processing_class.apply_chat_template( + example["prompt"] + example["completion"], + tools=example.get("tools"), + **example.get("chat_template_kwargs", {}), + ) + else: + prompt_ids = processing_class(text=example["prompt"]).input_ids + prompt_completion_ids = processing_class( + text=example["prompt"] + example["completion"] + ).input_ids + + # Check if the tokenized prompt starts with the tokenized prompt+completion + if not prompt_completion_ids[: len(prompt_ids)] == prompt_ids: + warnings.warn( + "Mismatch between tokenized prompt and the start of tokenized prompt+completion. " + "This may be due to unexpected tokenizer behavior, whitespace issues, or special " + "token handling. Verify that the tokenizer is processing text consistently." + ) + + # Create a completion mask + completion_mask = [0] * len(prompt_ids) + [1] * (len(prompt_completion_ids) - len(prompt_ids)) + processed = {"input_ids": prompt_completion_ids, "completion_mask": completion_mask} + + else: # language modeling case + if is_conversational(example): + processed = processing_class.apply_chat_template( + example["messages"], + return_dict=True, + return_assistant_tokens_mask=assistant_only_loss, + tools=example.get("tools"), + **example.get("chat_template_kwargs", {}), + ) + if "assistant_masks" in processed and 1 not in processed["assistant_masks"]: + raise RuntimeError( + "You're using `assistant_only_loss=True`, but at least one example has no " + "assistant tokens. This usually means the tokenizer's chat template doesn't " + "generate assistant masks — it may be missing the `{% generation %}` keyword. Please " + "check the template and ensure it's correctly configured to support assistant " + "masking." + ) + processed = {k: processed[k] for k in ("input_ids", "assistant_masks") if k in processed} + else: + processed = {"input_ids": processing_class(text=example[dataset_text_field]).input_ids} + return processed + + dataset = dataset.map( + tokenize, + fn_kwargs={ + "processing_class": processing_class, + "dataset_text_field": args.dataset_text_field, + "assistant_only_loss": args.assistant_only_loss, + }, + **map_kwargs, + ) + + # Pack or truncate + if packing: + if args.max_length is None: + raise ValueError("When packing is enabled, `max_length` can't be `None`.") + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Packing {dataset_name} dataset" + dataset = dataset.select_columns("input_ids") + # Packing adds new column "position_ids" needed for document aware flash attention + dataset = pack_dataset(dataset, args.max_length, args.packing_strategy, map_kwargs) + elif args.max_length is not None: + if isinstance(dataset, Dataset): # `IterableDataset.map` does not support `desc` + map_kwargs["desc"] = f"Truncating {dataset_name} dataset" + dataset = truncate_dataset(dataset, args.max_length, map_kwargs) + # For Liger kernel, ensure only input_ids is present + if args.use_liger_kernel: + dataset = dataset.select_columns({"input_ids", "position_ids"}.intersection(dataset.column_names)) + + return dataset + + def _set_signature_columns_if_needed(self): + # If `self.args.remove_unused_columns` is True, non-signature columns are removed. + # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids" + # and "attention_mask"). When using `train_on_completion_only` we add a "completion_mask" column to the + # dataset. So we need to override the default signature columns to include "completion_mask" as well. + if self._signature_columns is None: + self._signature_columns = [ + "input_ids", + "labels", + "position_ids", + "completion_mask", + "assistant_masks", + ] + + def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): + """ + Compute training loss and additionally compute token accuracies + """ + mode = "train" if self.model.training else "eval" + (loss, outputs) = super().compute_loss( + model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch + ) + if mode == "train": + # When using padding-free, the attention_mask is not present in the inputs, instead we have cu_seq_lens_q, + # cu_seq_lens_k, and max_length_k, max_length_q and position_ids. + if "attention_mask" in inputs: + num_tokens_in_batch = self.accelerator.gather_for_metrics(inputs["attention_mask"].sum()).sum().item() + elif "position_ids" in inputs: + local_num_tokens = torch.tensor(inputs["position_ids"].size(1), device=inputs["position_ids"].device) + num_tokens_in_batch = self.accelerator.gather_for_metrics(local_num_tokens).sum().item() + else: + raise ValueError("Expected 'attention_mask' or 'position_ids' in inputs.") + self._total_train_tokens += num_tokens_in_batch + self._metrics[mode]["num_tokens"] = [self._total_train_tokens] + + # Compute token accuracy if we have labels and if the model is not using Liger (no logits) + if "labels" in inputs and not self.args.use_liger_kernel: + shift_logits = outputs.logits[..., :-1, :].contiguous() + shift_labels = inputs["labels"][..., 1:].contiguous() + + # Get predictions + predictions = shift_logits.argmax(dim=-1) + + # Create mask for non-padding tokens (assuming ignore_index is -100) + mask = shift_labels != -100 + + # Calculate accuracy only on non-padding tokens + correct_predictions = (predictions == shift_labels) & mask + total_tokens = mask.sum() + correct_tokens = correct_predictions.sum() + + # Gather the correct_tokens and total_tokens across all processes + correct_tokens = self.accelerator.gather_for_metrics(correct_tokens) + total_tokens = self.accelerator.gather_for_metrics(total_tokens) + + # Compute the mean token accuracy and log it + total_sum = total_tokens.sum() + accuracy = (correct_tokens.sum() / total_sum).item() if total_sum > 0 else 0.0 + self._metrics[mode]["mean_token_accuracy"].append(accuracy) + + return (loss, outputs) if return_outputs else loss + + # Override training step to add activation offloading context. + def training_step(self, *args, **kwargs): + with self.maybe_activation_offload_context: + return super().training_step(*args, **kwargs) + + def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> None: + mode = "train" if self.model.training else "eval" + metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()} # average the metrics + + # This method can be called both in training and evaluation. When called in evaluation, the keys in `logs` + # start with "eval_". We need to add the prefix "eval_" to the keys in `metrics` to match the format. + if mode == "eval": + metrics = {f"eval_{key}": val for key, val in metrics.items()} + + logs = {**logs, **metrics} + super().log(logs, start_time) + self._metrics[mode].clear() + + # Ensure the model card is saved along with the checkpoint + def _save_checkpoint(self, model, trial): + if self.args.hub_model_id is None: + model_name = Path(self.args.output_dir).name + else: + model_name = self.args.hub_model_id.split("/")[-1] + self.create_model_card(model_name=model_name) + super()._save_checkpoint(model, trial) + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=list(tags), + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="SFT", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/trl/trainer/utils.py b/lib/python3.12/site-packages/trl/trainer/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9b1eca384138a68c40b594bbc55a4187bed74535 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/utils.py @@ -0,0 +1,1868 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import dataclasses +import importlib.resources as pkg_resources +import json +import random +import warnings +from collections import deque +from dataclasses import dataclass, field +from importlib.metadata import version +from typing import Any, Literal, Optional, Union + +import datasets +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as F +import torch.utils.data +from accelerate import Accelerator, PartialState +from accelerate.state import AcceleratorState +from huggingface_hub import ModelCard, ModelCardData +from torch.nn.utils.rnn import pad_sequence +from torch.utils.data import IterableDataset +from transformers import ( + BitsAndBytesConfig, + DataCollatorForLanguageModeling, + EvalPrediction, + GenerationConfig, + PreTrainedTokenizerBase, + TrainerState, + TrainingArguments, + is_comet_available, +) +from transformers.utils import ( + ModelOutput, + is_peft_available, + is_rich_available, + is_torch_mlu_available, + is_torch_npu_available, + is_torch_xpu_available, +) + +from ..trainer.model_config import ModelConfig + + +if is_rich_available(): + from rich.console import Console + from rich.panel import Panel + from rich.table import Table + from rich.text import Text + +if is_comet_available(): + import comet_ml + +if is_peft_available(): + from peft import LoraConfig, PeftConfig + + +class DataCollatorForCompletionOnlyLM(DataCollatorForLanguageModeling): + """ + Data collator used for completion tasks. It ensures that all the tokens of the labels are set to an 'ignore_index' + when they do not come from the assistant. This ensure that the loss is only calculated on the completion made by + the assistant. + + Args: + response_template (`Union[str, list[int]]`): + the template form that indicates the start of the response, typically something like '### Response:\n'. It + can also be passed as tokenized ids, which can be useful when using a tokenizer that encodes the response + differently if it does not have proper context. + instruction_template (`Union[str, list[int]]`): + the template form that indicates the start of the human instruction, typically something like '### + Human:\n'. Useful for assistant-style conversation datasets. It can also be passed as tokenized ids. + mlm (`bool`, *optional*, defaults to `False`): Whether to use masked language modeling in the underlying + `DataCollatorForLanguageModeling` class. Note that this option currently has no effect but is present + for flexibility and backwards-compatibility. + ignore_index (`int`, *optional*, defaults to `-100`): + The index to use to ignore the initial tokens with + """ + + def __init__( + self, + response_template: Union[str, list[int]], + instruction_template: Optional[Union[str, list[int]]] = None, + *args, + mlm: bool = False, + ignore_index: int = -100, + padding_free: bool = False, + **kwargs, + ): + super().__init__(*args, mlm=mlm, **kwargs) + warnings.warn( + "This class is deprecated and will be removed in version 0.20.0. To train on completion only, please use " + "the parameter `completion_only_loss` of `SFTConfig` instead.", + DeprecationWarning, + ) + + self.instruction_template = instruction_template + if isinstance(instruction_template, str): + # The user provides a string, must tokenize + self.instruction_token_ids = self.tokenizer.encode(self.instruction_template, add_special_tokens=False) + else: + # The user already provides the token ids + self.instruction_token_ids = instruction_template + + self.response_template = response_template + if isinstance(response_template, str): + # The user provides a string, must tokenize + self.response_token_ids = self.tokenizer.encode(self.response_template, add_special_tokens=False) + else: + # The user already provides the token ids + self.response_token_ids = response_template + + if not self.mlm and self.instruction_template and self.tokenizer.pad_token_id == self.tokenizer.eos_token_id: + warnings.warn( + "The pad_token_id and eos_token_id values of this tokenizer are identical. " + "If you are planning for multi-turn training, " + "it can result in the model continuously generating questions and answers without eos token. " + "To avoid this, set the pad_token_id to a different value.", + UserWarning, + ) + + self.ignore_index = ignore_index + self.padding_free = padding_free + + def torch_call(self, examples: list[Union[list[int], Any, dict[str, Any]]]) -> dict[str, Any]: + batch = super().torch_call(examples) + + if self.instruction_template is None: + for i in range(len(examples)): + response_token_ids_start_idx = None + + for idx in np.where(batch["labels"][i] == self.response_token_ids[0])[0]: + # `response_token_ids` is `'### Response:\n'`, here we are just making sure that the token IDs match + if ( + self.response_token_ids + == batch["labels"][i][idx : idx + len(self.response_token_ids)].tolist() + ): + response_token_ids_start_idx = idx + + if response_token_ids_start_idx is None: + warnings.warn( + f"Could not find response key `{self.response_template}` in the following instance: " + f"{self.tokenizer.decode(batch['input_ids'][i])}. This instance will be ignored in loss " + "calculation. Note, if this happens often, consider increasing the `max_length`.", + UserWarning, + ) + batch["labels"][i, :] = self.ignore_index + else: + response_token_ids_end_idx = response_token_ids_start_idx + len(self.response_token_ids) + + # Make pytorch loss function ignore all tokens up through the end of the response key + batch["labels"][i, :response_token_ids_end_idx] = self.ignore_index + + else: + for i in range(len(examples)): + response_token_ids_idxs = [] + human_token_ids_idxs = [] + + for assistant_idx in np.where(batch["labels"][i] == self.response_token_ids[0])[0]: + # find the indexes of the start of a response. + if ( + self.response_token_ids + == batch["labels"][i][assistant_idx : assistant_idx + len(self.response_token_ids)].tolist() + ): + response_token_ids_idxs.append(assistant_idx + len(self.response_token_ids)) + + if len(response_token_ids_idxs) == 0: + warnings.warn( + f"Could not find response key `{self.response_template}` in the following instance: " + f"{self.tokenizer.decode(batch['input_ids'][i])}. This instance will be ignored in loss " + "calculation. Note, if this happens often, consider increasing the `max_length`.", + UserWarning, + ) + batch["labels"][i, :] = self.ignore_index + + human_token_ids = self.instruction_token_ids + for human_idx in np.where(batch["labels"][i] == human_token_ids[0])[0]: + # find the indexes of the start of a human answer. + if human_token_ids == batch["labels"][i][human_idx : human_idx + len(human_token_ids)].tolist(): + human_token_ids_idxs.append(human_idx) + + if len(human_token_ids_idxs) == 0: + warnings.warn( + f"Could not find instruction key `{self.instruction_template}` in the following instance: " + f"{self.tokenizer.decode(batch['input_ids'][i])}. This instance will be ignored in loss " + "calculation. Note, if this happens often, consider increasing the `max_length`.", + UserWarning, + ) + batch["labels"][i, :] = self.ignore_index + + if ( + len(human_token_ids_idxs) > 0 + and len(response_token_ids_idxs) > 0 + and human_token_ids_idxs[0] > response_token_ids_idxs[0] + ): + human_token_ids_idxs = [0] + human_token_ids_idxs + + for idx, (start, end) in enumerate(zip(human_token_ids_idxs, response_token_ids_idxs)): + # Make pytorch loss function ignore all non response tokens + if idx != 0: + batch["labels"][i, start:end] = self.ignore_index + else: + batch["labels"][i, :end] = self.ignore_index + + if len(response_token_ids_idxs) < len(human_token_ids_idxs): + batch["labels"][i, human_token_ids_idxs[-1] :] = self.ignore_index + + if self.padding_free: + # remove padding, `attention_mask` and add `position_ids` + attn_mask = batch.pop("attention_mask") + batch["input_ids"] = batch["input_ids"][attn_mask.bool()].unsqueeze(0) + batch["position_ids"] = attn_mask.cumsum(1)[attn_mask.bool()].unsqueeze(0) - 1 + batch["labels"] = batch["labels"][attn_mask.bool()].unsqueeze(0) + batch["labels"][batch["position_ids"] == 0] = self.ignore_index + + # Calculate cumulative sequence lengths for queries and keys to prevent graph breaks during further computations. + flattened_position_ids = batch["position_ids"].flatten() + indices_q = torch.arange( + flattened_position_ids.size(0), device=flattened_position_ids.device, dtype=torch.int32 + ) + batch["cu_seq_lens_q"] = torch.cat( + ( + indices_q[flattened_position_ids == 0], + torch.tensor( + flattened_position_ids.size(), device=flattened_position_ids.device, dtype=torch.int32 + ), + ) + ).unsqueeze(0) + batch["cu_seq_lens_k"] = batch["cu_seq_lens_q"] + + # Determine maximum sequence lengths to prevent graph breaks during further computations. + batch["max_length_k"] = torch.tensor([flattened_position_ids.max().item() + 1]) + batch["max_length_q"] = batch["max_length_k"] + + return batch + + +@dataclass +class DataCollatorForChatML: + """ + Data collator for ChatML format datasets. + """ + + tokenizer: PreTrainedTokenizerBase + ignore_index: int = -100 + max_length: int = None + prompt_key: str = "prompt" + messages_key: str = "messages" + + def __post_init__(self): + if self.tokenizer.pad_token_id is None: + raise ValueError("The tokenizer does not have a pad token. Please set `pad_token_id` in the tokenizer.") + if self.max_length is None: + # set a sensible default + self.max_length = min(self.tokenizer.model_max_length, 1024) + + def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + input_ids = [] + attention_mask = [] + prompts_input_ids = [] + prompt_attention_mask = [] + labels = [] + + for example in examples: + formatted_prompt = example.get(self.prompt_key, None) + if formatted_prompt is None: + prompt = example[self.messages_key][:-1] + formatted_prompt = self.tokenizer.apply_chat_template( + prompt, tokenize=False, add_generation_prompt=True + ) + + if "input_ids" not in example: + message = example[self.messages_key] + formatted_message = self.tokenizer.apply_chat_template( + message, tokenize=False, add_generation_prompt=False + ) + tokenized_message = self.tokenizer( + formatted_message, + truncation=True, + max_length=self.max_length, + padding=False, + return_tensors=None, + add_special_tokens=False, + ) + input_ids.append(tokenized_message["input_ids"]) + if "attention_mask" in example: + attention_mask.append(tokenized_message["attention_mask"]) + else: + attention_mask.append([1] * len(tokenized_message["input_ids"])) + else: + input_ids.append(example["input_ids"]) + if "attention_mask" in example: + attention_mask.append(example["attention_mask"]) + else: + attention_mask.append([1] * len(example["input_ids"])) + + tokenized_prompt = self.tokenizer( + formatted_prompt, + truncation=True, + max_length=len(input_ids[-1]), + padding=False, + return_tensors=None, + add_special_tokens=False, + ) + + prompts_input_ids.append(tokenized_prompt["input_ids"]) + prompt_attention_mask.append(tokenized_prompt["attention_mask"]) + + # Create the labels that will have all but the completion tokens of the example["input_ids"] set to ignore_index + label = [self.ignore_index] * len(input_ids[-1]) + completion_start_idx = len(tokenized_prompt["input_ids"]) + label[completion_start_idx:] = input_ids[-1][completion_start_idx:] + labels.append(label) + + # convert to list of tensors and pad + input_ids = [torch.tensor(ids, dtype=torch.long) for ids in input_ids] + attention_mask = [torch.tensor(mask, dtype=torch.long) for mask in attention_mask] + labels = [torch.tensor(label, dtype=torch.long) for label in labels] + input_ids = pad(input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id) + attention_mask = pad(attention_mask, padding_side="left", padding_value=0) + labels = pad(labels, padding_side="left", padding_value=self.ignore_index) + + prompts_input_ids = [torch.tensor(ids, dtype=torch.long) for ids in prompts_input_ids] + prompt_attention_mask = [torch.tensor(mask, dtype=torch.long) for mask in prompt_attention_mask] + prompts_input_ids = pad(prompts_input_ids, padding_side="left", padding_value=self.tokenizer.pad_token_id) + prompt_attention_mask = pad(prompt_attention_mask, padding_side="left", padding_value=0) + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + "prompts": prompts_input_ids, + "prompt_attention_mask": prompt_attention_mask, + } + + +@dataclass +class RewardDataCollatorWithPadding: + r""" + Reward DataCollator class that pads the inputs to the maximum length of the batch. + + Args: + tokenizer (`PreTrainedTokenizerBase`): + The tokenizer used for encoding the data. + padding (`Union[bool, str, `PaddingStrategy`]`, `optional`, defaults to `True`): + padding_strategy to pass to the tokenizer. + pad_to_multiple_of (`int` or `None`, `optional`, defaults to `None`): + If set will pad the sequence to a multiple of the provided value. + return_tensors (`str`, `optional`, defaults to `"pt"`): + The tensor type to use. + """ + + tokenizer: PreTrainedTokenizerBase + padding: Union[bool, str] = True + pad_to_multiple_of: Optional[int] = None + return_tensors: str = "pt" + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: + features_chosen = [] + features_rejected = [] + margin = [] + # check if we have a margin. If we do, we need to batch it as well + has_margin = "margin" in features[0] + for feature in features: + # check if the keys are named as expected + if ( + "input_ids_chosen" not in feature + or "input_ids_rejected" not in feature + or "attention_mask_chosen" not in feature + or "attention_mask_rejected" not in feature + ): + raise ValueError( + "The features should include `input_ids_chosen`, `attention_mask_chosen`, `input_ids_rejected` and `attention_mask_rejected`" + ) + + features_chosen.append( + { + "input_ids": feature["input_ids_chosen"], + "attention_mask": feature["attention_mask_chosen"], + } + ) + features_rejected.append( + { + "input_ids": feature["input_ids_rejected"], + "attention_mask": feature["attention_mask_rejected"], + } + ) + if has_margin: + margin.append(feature["margin"]) + batch_chosen = self.tokenizer.pad( + features_chosen, + padding=self.padding, + pad_to_multiple_of=self.pad_to_multiple_of, + return_tensors=self.return_tensors, + ) + batch_rejected = self.tokenizer.pad( + features_rejected, + padding=self.padding, + pad_to_multiple_of=self.pad_to_multiple_of, + return_tensors=self.return_tensors, + ) + batch = { + "input_ids_chosen": batch_chosen["input_ids"], + "attention_mask_chosen": batch_chosen["attention_mask"], + "input_ids_rejected": batch_rejected["input_ids"], + "attention_mask_rejected": batch_rejected["attention_mask"], + "return_loss": True, + } + if has_margin: + margin = torch.tensor(margin, dtype=torch.float) + batch["margin"] = margin + return batch + + +def pad( + tensors: list[torch.Tensor], + padding_value: int = 0, + padding_side: str = "right", + pad_to_multiple_of: Optional[int] = None, +) -> torch.Tensor: + """ + Pads a list of tensors to the same shape along the first dimension. + + Args: + tensors (`list[torch.Tensor]`): + List of input tensors to pad. + padding_value (`int`): + Value to use for padding. Default is 0. + padding_side (`str`): + Side on which to add padding. Must be 'left' or 'right'. Default is 'right'. + pad_to_multiple_of (`int`, *optional*, defaults to `None`): + If set will pad the sequence to a multiple of the provided value. + + Returns: + `torch.Tensor`: + A single tensor containing the padded tensors. + + Examples: + ```python + >>> import torch + + >>> pad([torch.tensor([1, 2, 3]), torch.tensor([4, 5])]) + tensor([[1, 2, 3], + [4, 5, 0]]) + + >>> pad([torch.tensor([[1, 2], [3, 4]]), torch.tensor([[5, 6]])]) + tensor([[[1, 2], + [3, 4]], + [[5, 6], + [0, 0]]]) + ``` + """ + # Determine the maximum shape for each dimension + output_shape = np.max([t.shape for t in tensors], 0).tolist() + + # Apply pad_to_multiple_of to the first (sequence) dimension + if pad_to_multiple_of is not None: + remainder = output_shape[0] % pad_to_multiple_of + if remainder != 0: + output_shape[0] += pad_to_multiple_of - remainder + + # Create an output tensor filled with the padding value + output = torch.full((len(tensors), *output_shape), padding_value, dtype=tensors[0].dtype, device=tensors[0].device) + + for i, t in enumerate(tensors): + if padding_side == "left": + seq_start = output_shape[0] - t.shape[0] + elif padding_side == "right": + seq_start = 0 + else: + raise ValueError("padding_side must be 'left' or 'right'") + + # Define the slices + seq_slice = slice(seq_start, seq_start + t.shape[0]) + slices = (seq_slice,) + tuple(slice(0, s) for s in t.shape[1:]) + output[i][slices] = t + + return output + + +@dataclass +class DPODataCollatorWithPadding: + r""" + DPO DataCollator class that pads the tokenized inputs to the maximum length of the batch. + + Args: + pad_token_id (`int` defaults to 0): + The tokenizer's pad_token_id. + label_pad_token_id (`int`, defaults to -100): + The label used for masking. + is_encoder_decoder (`bool` or `None`, `optional`, defaults to `None`): + Whether you model has an encoder_decoder architecture. + """ + + pad_token_id: int = 0 + label_pad_token_id: int = -100 + is_encoder_decoder: Optional[bool] = False + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]: + # first, pad everything to the same length + padded_batch = {} + for k in features[0].keys(): + if k.endswith(("_input_ids", "_attention_mask", "_labels", "_pixel_values")): + if self.is_encoder_decoder: + to_pad = [torch.LongTensor(ex[k]) for ex in features] + + if (k.startswith("prompt")) and (k.endswith("input_ids")): + if self.pad_token_id is None: + raise ValueError( + "Padding is enabled, but the tokenizer is not configured with a padding token." + " Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)" + " before calling the trainer." + ) + padding_value = self.pad_token_id + elif k.endswith("_attention_mask"): + padding_value = 0 + elif k.startswith(("chosen", "rejected", "completion")) or ("decoder" in k): + padding_value = self.label_pad_token_id + else: + raise ValueError(f"Unexpected key in batch '{k}'") + padded_batch[k] = pad_sequence(to_pad, batch_first=True, padding_value=padding_value) + else: + # Set padding value based on the key + if k.endswith("_input_ids"): + if self.pad_token_id is None: + raise ValueError( + "Padding is enabled, but the tokenizer is not configured with a padding token." + " Explicitly set `tokenizer.pad_token` (e.g. `tokenizer.pad_token = tokenizer.eos_token`)" + " before calling the trainer." + ) + padding_value = self.pad_token_id + elif k.endswith("_labels"): + padding_value = self.label_pad_token_id + elif k.endswith("_attention_mask"): + padding_value = 0 + elif k.endswith("_pixel_values"): + padding_value = 0 # TODO: check if this is correct + else: + raise ValueError(f"Unexpected key in batch '{k}'") + + # Set padding side based on the key + if k in ["prompt_input_ids", "prompt_attention_mask"]: + padding_side = "left" + else: + padding_side = "right" + + # Set the dtype + if k.endswith("_pixel_values"): + dtype = torch.float32 # will be downcasted if necessary by the Trainer + else: + dtype = torch.int64 + + # Convert to tensor and pad + to_pad = [torch.tensor(ex[k], dtype=dtype) for ex in features] + padded_batch[k] = pad(to_pad, padding_value=padding_value, padding_side=padding_side) + elif k.endswith("_logps"): + # the cached reference model logprobs + padded_batch[k] = torch.tensor([ex[k] for ex in features]) + else: + padded_batch[k] = [ex[k] for ex in features] + + return padded_batch + + +class ConstantLengthDataset(IterableDataset): + """ + Iterable dataset that returns constant length chunks of tokens from stream of text files. The dataset also formats + the text before tokenization with a specific format that is provided by the user. + + Args: + tokenizer (`transformers.PreTrainedTokenizer`): + The processor used for processing the data. + dataset (`dataset.Dataset`): + Dataset with text files. + dataset_text_field (`str` or `None`, *optional*, defaults to `None`): + Name of the field in the dataset that contains the text. Only one of `dataset_text_field` and + `formatting_func` should be provided. + formatting_func (`Callable`, *optional*): + Function that formats the text before tokenization. Usually it is recommended to follow a certain pattern + such as `"### Question: {question} ### Answer: {answer}"`. Only one of `dataset_text_field` and + `formatting_func` should be provided. + infinite (`bool`, *optional*, defaults to `False`): + If True the iterator is reset after dataset reaches end else stops. + seq_length (`int`, *optional*, defaults to `1024`): + Length of token sequences to return. + num_of_sequences (`int`, *optional*, defaults to `1024`): + Number of token sequences to keep in buffer. + chars_per_token (`int`, *optional*, defaults to `3.6`): + Number of characters per token used to estimate number of tokens in text buffer. + eos_token_id (`int`, *optional*, defaults to `0`): + Id of the end of sequence token if the passed tokenizer does not have an EOS token. + shuffle (`bool`, *optional*, defaults to `True`) + Shuffle the examples before they are returned + append_concat_token (`bool`, *optional*, defaults to `True`) + If true, appends `eos_token_id` at the end of each sample being packed. + add_special_tokens (`bool`, *optional*, defaults to `True`) + If true, tokenizers adds special tokens to each sample being packed. + """ + + def __init__( + self, + tokenizer, + dataset, + dataset_text_field=None, + formatting_func=None, + infinite=False, + seq_length=1024, + num_of_sequences=1024, + chars_per_token=3.6, + eos_token_id=0, + shuffle=True, + append_concat_token=True, + add_special_tokens=True, + ): + warnings.warn( + "This class is deprecated and will be removed in version 0.20.0. To use packing, use the argument " + "`packing` of `SFTConfig` instead.", + DeprecationWarning, + ) + self.tokenizer = tokenizer + self.concat_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id else eos_token_id + self.dataset = dataset + self.seq_length = seq_length + self.infinite = infinite + self.current_size = 0 + self.max_buffer_size = seq_length * chars_per_token * num_of_sequences + self.shuffle = shuffle + self.append_concat_token = append_concat_token + self.add_special_tokens = add_special_tokens + + if dataset_text_field is not None and formatting_func is not None: + warnings.warn( + "Only one of `dataset_text_field` and `formatting_func` should be provided. " + "Ignoring `dataset_text_field` and using `formatting_func`.", + UserWarning, + ) + + if formatting_func is not None: + self.formatting_func = formatting_func + elif dataset_text_field is not None: + self.formatting_func = lambda x: x[dataset_text_field] + else: # neither is provided + raise ValueError("Either `dataset_text_field` or `formatting_func` should be provided.") + + self.pretokenized = False + column_names = ( + dataset.column_names if isinstance(dataset, (datasets.Dataset, datasets.IterableDataset)) else None + ) + if column_names is not None and "input_ids" in column_names: + self.pretokenized = True + # since the dataset is tokenized, the unit of buffer size should be tokens + self.max_buffer_size = seq_length * num_of_sequences + + def __len__(self): + return len(self.dataset) + + def __iter__(self): + iterator = iter(self.dataset) + more_examples = True + while more_examples: + buffer, buffer_len = [], 0 + while True: + if buffer_len >= self.max_buffer_size: + break + try: + buffer.append(self.formatting_func(next(iterator))) + buffer_len += len(buffer[-1]) + except StopIteration: + if self.infinite: + iterator = iter(self.dataset) + else: + more_examples = False + break + if self.shuffle: + random.shuffle(buffer) + if self.pretokenized: + tokenized_inputs = buffer + else: + tokenized_inputs = self.tokenizer( + buffer, add_special_tokens=self.add_special_tokens, truncation=False + )["input_ids"] + all_token_ids = [] + for tokenized_input in tokenized_inputs: + if self.append_concat_token: + tokenized_input = tokenized_input + [self.concat_token_id] + all_token_ids.extend(tokenized_input) + examples = [] + for i in range(0, len(all_token_ids), self.seq_length): + input_ids = all_token_ids[i : i + self.seq_length] + if len(input_ids) == self.seq_length: + examples.append(input_ids) + if self.shuffle: + # Shuffle again, otherwise split examples occur in consecutive tensors. + random.shuffle(examples) + for example in examples: + self.current_size += 1 + yield { + "input_ids": torch.LongTensor(example), + "labels": torch.LongTensor(example), + } + + +@dataclass +class RunningMoments: + """ + Calculates the running mean and standard deviation of a data stream. Reference: + https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L75 + """ + + accelerator: Accelerator + mean: float = 0 + std: float = 1 + var: float = 1 + count: float = 1e-24 + + @torch.no_grad() + def update(self, xs: torch.Tensor) -> tuple[float, float]: + """ + Updates running moments from batch's moments computed across ranks + """ + if self.accelerator.use_distributed: + xs_mean, xs_var, xs_count = get_global_statistics(self.accelerator, xs) + else: + xs_count = xs.numel() + xs_var, xs_mean = torch.var_mean(xs, unbiased=False) + xs_mean, xs_var = xs_mean.float(), xs_var.float() + + delta = xs_mean - self.mean + tot_count = self.count + xs_count + + new_sum = xs_var * xs_count + # correct old_sum deviation accounting for the new mean + old_sum = self.var * self.count + delta**2 * self.count * xs_count / tot_count + tot_sum = old_sum + new_sum + + self.mean += (delta * xs_count / tot_count).item() + new_var = tot_sum / tot_count + self.std = (new_var * tot_count / (tot_count - 1)).float().sqrt().item() + self.var = new_var.item() + self.count = tot_count + + return xs_mean.item(), (xs_var * xs_count / (xs_count - 1)).float().sqrt().item() + + def save_to_json(self, json_path: str): + """Save the content of this instance in JSON format inside `json_path`.""" + # save everything except accelerator + if self.accelerator.is_main_process: + save_dict = dataclasses.asdict(self, dict_factory=lambda x: {k: v for (k, v) in x if k != "accelerator"}) + json_string = json.dumps(save_dict, indent=2, sort_keys=True) + "\n" + with open(json_path, "w", encoding="utf-8") as f: + f.write(json_string) + + @classmethod + def load_from_json(cls, accelerator: Accelerator, json_path: str): + """Create an instance from the content of `json_path`.""" + # load everything except accelerator + with open(json_path, encoding="utf-8") as f: + text = f.read() + return cls(accelerator=accelerator, **json.loads(text)) + + +@torch.no_grad() +def get_global_statistics( + accelerator, xs: torch.Tensor, mask=None, device="cpu" +) -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Computes element-wise mean and variance of the tensor across processes. Reference: + https://github.com/OpenLMLab/MOSS-RLHF/blob/40b91eb2f2b71b16919addede0341d2bef70825d/utils.py#L57C1-L73C75 + """ + xs = xs.to(accelerator.device) + sum_and_count = torch.tensor([xs.sum(), (xs.numel() if mask is None else mask.sum())], device=xs.device) + sum_and_count = accelerator.reduce(sum_and_count) + global_sum, count = sum_and_count + global_mean = global_sum / count + + sum_var = torch.sum(((xs - global_mean) ** 2).mul(1 if mask is None else mask)) + sum_var = accelerator.reduce(sum_var) + global_var = sum_var / count + + return global_mean.to(device), global_var.to(device), count.item() + + +def compute_accuracy(eval_pred: EvalPrediction) -> dict[str, float]: + predictions, labels = eval_pred + if predictions.ndim == 3: + # Token classification task. Shapes are (batch_size, seq_len, num_labels) and (batch_size, seq_len) + # Used to compute the accuracy in the prm_trainer. + predictions = np.argmax(predictions, axis=2) + + # Flatten the predictions and labels to remove the ignored tokens. + predictions = np.array( + [p for prediction, label in zip(predictions, labels) for (p, lbl) in zip(prediction, label) if lbl != -100] + ) + labels = np.array([lbl for label in labels for lbl in label if lbl != -100]) + + else: + # Here, predictions is rewards_chosen and rewards_rejected. Shapes are (batch_size, 2) and (batch_size,) + # We want to see how much of the time rewards_chosen > rewards_rejected. + equal_mask = predictions[:, 0] == predictions[:, 1] + equal_predictions_count = int(equal_mask.sum()) + + if equal_predictions_count > 0: + warnings.warn( + f"There are {equal_predictions_count} out of {len(predictions[:, 0])} instances where the predictions " + "for both options are equal. These instances are ignored in the accuracy computation.", + UserWarning, + ) + + # Filter out equal predictions + predictions = predictions[~equal_mask] + labels = labels[~equal_mask] + + # Use the remaining predictions for accuracy calculation + predictions = np.argmax(predictions, axis=1) + + accuracy = np.array(predictions == labels, dtype=float).mean().item() + return {"accuracy": accuracy} + + +def pad_to_length(tensor: torch.Tensor, length: int, pad_value: Union[int, float], dim: int = -1) -> torch.Tensor: + if tensor.size(dim) >= length: + return tensor + else: + pad_size = list(tensor.shape) + pad_size[dim] = length - tensor.size(dim) + return torch.cat( + [ + tensor, + pad_value * torch.ones(*pad_size, dtype=tensor.dtype, device=tensor.device), + ], + dim=dim, + ) + + +def disable_dropout_in_model(model: torch.nn.Module) -> None: + for module in model.modules(): + if isinstance(module, torch.nn.Dropout): + module.p = 0 + + +def exact_div(a, b, custom_error_message=""): + q = a // b + if a != q * b: + raise ValueError(f"{custom_error_message}, inexact division: {a} / {b} = {a / b}") + return q + + +# copied from https://github.com/kvablack/ddpo-pytorch/blob/main/ddpo_pytorch/stat_tracking.py#L5 +class PerPromptStatTracker: + r""" + Class for tracking statistics per prompt. Mainly used to calculate advantage for the DPPO algorithm + + Args: + buffer_size (`int`): + Size of the buffer to keep for each prompt. + min_count (`int`): + Minimum number of samples to keep in the buffer before calculating the mean and std. + """ + + def __init__(self, buffer_size, min_count): + self.buffer_size = buffer_size + self.min_count = min_count + self.stats = {} + + def update(self, prompts, rewards): + prompts = np.array(prompts) + rewards = np.array(rewards) + unique = np.unique(prompts) + advantages = np.empty_like(rewards) + for prompt in unique: + prompt_rewards = rewards[prompts == prompt] + if prompt not in self.stats: + self.stats[prompt] = deque(maxlen=self.buffer_size) + self.stats[prompt].extend(prompt_rewards) + + if len(self.stats[prompt]) < self.min_count: + mean = np.mean(rewards) + std = np.std(rewards) + 1e-6 + else: + mean = np.mean(self.stats[prompt]) + std = np.std(self.stats[prompt]) + 1e-6 + advantages[prompts == prompt] = (prompt_rewards - mean) / std + + return advantages + + def get_stats(self): + return {k: {"mean": np.mean(v), "std": np.std(v), "count": len(v)} for k, v in self.stats.items()} + + +def peft_module_casting_to_bf16(model): + for name, module in model.named_modules(): + if isinstance(module, torch.nn.LayerNorm) or "norm" in name: + module = module.to(torch.float32) + elif any(x in name for x in ["lm_head", "embed_tokens", "wte", "wpe"]): + if hasattr(module, "weight"): + if module.weight.dtype == torch.float32: + module = module.to(torch.bfloat16) + + +def get_quantization_config(model_args: ModelConfig) -> Optional[BitsAndBytesConfig]: + if model_args.load_in_4bit: + quantization_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=model_args.torch_dtype, # For consistency with model weights, we use the same value as `torch_dtype` + bnb_4bit_quant_type=model_args.bnb_4bit_quant_type, + bnb_4bit_use_double_quant=model_args.use_bnb_nested_quant, + bnb_4bit_quant_storage=model_args.torch_dtype, + ) + elif model_args.load_in_8bit: + quantization_config = BitsAndBytesConfig( + load_in_8bit=True, + ) + else: + quantization_config = None + + return quantization_config + + +def get_kbit_device_map() -> Optional[dict[str, int]]: + if torch.cuda.is_available() or is_torch_xpu_available(): + return {"": PartialState().local_process_index} + else: + return None + + +def get_peft_config(model_args: ModelConfig) -> "Optional[PeftConfig]": + if model_args.use_peft is False: + return None + + if not is_peft_available(): + raise ValueError( + "You need to have PEFT library installed in your environment, make sure to install `peft`. " + "Make sure to run `pip install -U peft`." + ) + + peft_config = LoraConfig( + task_type=model_args.lora_task_type, + r=model_args.lora_r, + target_modules=model_args.lora_target_modules, + lora_alpha=model_args.lora_alpha, + lora_dropout=model_args.lora_dropout, + bias="none", + use_rslora=model_args.use_rslora, + use_dora=model_args.use_dora, + modules_to_save=model_args.lora_modules_to_save, + ) + + return peft_config + + +def get_exp_cap(value, decimal=4): + """ + Get the exponent cap of a value. This is used to cap the exponent of a value to avoid overflow. The formula is : + log(value.dtype.max) E.g. + For float32 data type, the maximum exponent value is 88.7228 to 4 decimal points. + + Args: + value (`torch.Tensor`): + The input tensor to obtain the data type + decimal (`int`): + The number of decimal points of the output exponent cap. eg: direct calling exp(log(torch.float32.max)) + will result in inf so we cap the exponent to 88.7228 to avoid overflow. + """ + vdtype_max = torch.zeros([1]).to(value.dtype) + torch.finfo(value.dtype).max + vdtype_log_max = torch.log(vdtype_max).to(value.device) + return torch.floor(vdtype_log_max * 10**decimal) / 10**decimal if decimal > 0 else vdtype_log_max + + +def cap_exp(value, cap=-1): + # Cap the exponent value below the upper-bound to avoid overflow, before calling torch.exp + cap = get_exp_cap(value) if cap < 0 else cap + return torch.exp(torch.clamp(value, max=cap)) + + +def print_rich_table(df: pd.DataFrame) -> None: + if not is_rich_available(): + raise ImportError( + "The function `print_rich_table` requires the `rich` library. Please install it with `pip install rich`." + ) + console = Console() + table = Table(show_lines=True) + for column in df.columns: + table.add_column(column) + for _, row in df.iterrows(): + table.add_row(*row.astype(str).tolist()) + console.print(table) + + +SIMPLE_SFT_CHAT_TEMPLATE = "{% for message in messages %}{{' ' + message['content']}}{% endfor %}{{ eos_token }}" +# SIMPLE_SFT_CHAT_TEMPLATE simply ends things with an EOS token, this helps the SFT model learn to end the completions with EOS tokens + +SIMPLE_CHAT_TEMPLATE = "{% for message in messages %}{{message['role'].capitalize() + ': ' + message['content'] + '\n\n'}}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}" + + +@dataclass +class OnlineTrainerState(TrainerState): + episode: int = 0 + + +@dataclass +class OnPolicyConfig(TrainingArguments): + r""" + Base configuration class for on-policy trainers. + + This class includes only the parameters that are specific to some on-policy training. For a full list of training + arguments, please refer to the [`~transformers.TrainingArguments`] documentation. Note that default values in this + class may differ from those in [`~transformers.TrainingArguments`]. + + Using [`~transformers.HfArgumentParser`] we can turn this class into + [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the + command line. + + Parameters: + run_name (`str` or `None`, *optional*, defaults to `None`): + Name of the run. + dataset_num_proc (`int` or `None`, *optional*, defaults to `None`): + Number of processes to use for processing the dataset. + num_mini_batches (`int`, *optional*, defaults to `1`): + Number of minibatches to split a batch into. + total_episodes (`int` or `None`, *optional*, defaults to `None`): + Total number of episodes in the dataset. + local_rollout_forward_batch_size (`int`, *optional*, defaults to `64`): + Per rank no grad forward pass in the rollout phase. + num_sample_generations (`int`, *optional*, defaults to `10`): + Number of debugging samples generations (i.e., `generate_completions` calls) throughout training. + response_length (`int`, *optional*, defaults to `53`): + Length of the response. + stop_token (`str` or `None`, *optional*, defaults to `None`): + Specifies the stop token to use for text generation. This parameter is mutually exclusive with + `stop_token_id`. + + - `None`: No stop token is applied, unless `stop_token_id` is specified. + - `'eos'`: Uses the tokenizer's `eos_token`. + + stop_token_id (`int` or `None`, *optional*, defaults to `None`): + Specifies the ID of the stop token to use for text generation. If `None`, no stop token ID is applied, + unless `stop_token` is specified. This parameter is mutually exclusive with `stop_token`. + temperature (`float`, *optional*, defaults to `0.7`): + Sampling temperature. + missing_eos_penalty (`float` or `None`, *optional*, defaults to `None`): + Penalty applied to the score when the model fails to generate an EOS token. This is useful to encourage to + generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be a positive + value. + sft_model_path (`str`, *optional*, defaults to `"EleutherAI/pythia-160m"`): + Path to the SFT model. + world_size (`int` or `None`, *optional*, defaults to `None`): + Number of processes (GPUs) to use for the training. + num_total_batches (`int` or `None`, *optional*, defaults to `None`): + Number of total batches to train. + micro_batch_size (`int` or `None`, *optional*, defaults to `None`): + Micro batch size across devices (HF's `per_device_train_batch_size` * `world_size`). + local_batch_size (`int` or `None`, *optional*, defaults to `None`): + Batch size per GPU (HF's `per_device_train_batch_size` * `gradient_accumulation_steps`). + batch_size (`int` or `None`, *optional*, defaults to `None`): + Batch size across devices (HF's `per_device_train_batch_size` * `world_size` * + `gradient_accumulation_steps`). + local_mini_batch_size (`int` or `None`, *optional*, defaults to `None`): + Mini batch size per GPU. + mini_batch_size (`int` or `None`, *optional*, defaults to `None`): + Mini batch size across GPUs. + push_to_hub (`bool`, *optional*, defaults to `False`): + Whether to push the model to the Hub after training. + """ + + # Parameters whose default values are overridden from TrainingArguments + logging_steps: float = field( + default=10, + metadata={ + "help": "Log every X updates steps. Should be an integer or a float in range `[0,1)`. If smaller than 1, " + "will be interpreted as ratio of total training steps." + }, + ) + bf16: Optional[bool] = field( + default=None, + metadata={ + "help": "Whether to use bf16 (mixed) precision instead of 32-bit. Requires Ampere or higher NVIDIA " + "architecture or Intel XPU or using CPU (use_cpu) or Ascend NPU. If not set, it defaults to `True` if " + "`fp16` is not set." + }, + ) + + run_name: Optional[str] = field( + default=None, + metadata={"help": "Name of the run."}, + ) + dataset_num_proc: Optional[int] = field( + default=None, + metadata={"help": "Number of processes to use for processing the dataset."}, + ) + num_mini_batches: int = field( + default=1, + metadata={"help": "Number of minibatches to split a batch into."}, + ) + total_episodes: Optional[int] = field( + default=None, + metadata={"help": "Total number of episodes in the dataset."}, + ) + local_rollout_forward_batch_size: int = field( + default=64, + metadata={"help": "Per rank no grad forward pass in the rollout phase."}, + ) + num_sample_generations: int = field( + default=10, + metadata={ + "help": "Number of debugging samples generations (i.e., `generate_completions` calls) throughout training." + }, + ) + response_length: int = field( + default=53, + metadata={"help": "Length of the response."}, + ) + stop_token: Optional[Literal["eos"]] = field( + default=None, + metadata={ + "help": "Specifies the stop token to use for text generation. This parameter is mutually exclusive with " + "`stop_token_id`." + }, + ) + stop_token_id: Optional[int] = field( + default=None, + metadata={ + "help": "Specifies the ID of the stop token to use for text generation. If `None`, no stop token ID is " + "applied, unless `stop_token` is specified. This parameter is mutually exclusive with `stop_token`." + }, + ) + temperature: float = field( + default=0.7, + metadata={"help": "Sampling temperature."}, + ) + missing_eos_penalty: Optional[float] = field( + default=None, + metadata={ + "help": "Penalty applied to the score when the model fails to generate an EOS token. This is useful to " + "encourage to generate completions shorter than the maximum length (`max_new_tokens`). The penalty must be " + "a positive value." + }, + ) + sft_model_path: str = field( + default="EleutherAI/pythia-160m", + metadata={"help": "Path to the SFT model."}, + ) + world_size: Optional[int] = field( + default=None, + metadata={"help": "Number of processes (GPUs) to use for the training."}, + ) + num_total_batches: Optional[int] = field( + default=None, + metadata={"help": "Number of total batches to train."}, + ) + micro_batch_size: Optional[int] = field( + default=None, + metadata={"help": "Micro batch size across devices (HF's `per_device_train_batch_size` * `world_size`)."}, + ) + local_batch_size: Optional[int] = field( + default=None, + metadata={"help": "Batch size per GPU (HF's `per_device_train_batch_size` * `gradient_accumulation_steps`)."}, + ) + batch_size: Optional[int] = field( + default=None, + metadata={ + "help": "Batch size across devices (HF's `per_device_train_batch_size` * `world_size` * " + "`gradient_accumulation_steps`)." + }, + ) + local_mini_batch_size: Optional[int] = field( + default=None, + metadata={"help": "Mini batch size per GPU."}, + ) + mini_batch_size: Optional[int] = field( + default=None, + metadata={"help": "Mini batch size across GPUs."}, + ) + push_to_hub: bool = field( + default=False, + metadata={"help": "Whether to push the model to the Hub after training."}, + ) + + def __post_init__(self): + self.bf16 = not (self.fp16) if self.bf16 is None else self.bf16 + + super().__post_init__() + + +def first_true_indices(bools: torch.Tensor, dtype=torch.long) -> torch.Tensor: + """ + Takes an N-dimensional bool tensor and returns an (N-1)-dimensional tensor of integers giving the position of the + first True in each "row". + + Returns the length of the rows (bools.size(-1)) if no element is True in a given row. + + Args: + bools (`torch.Tensor`): + An N-dimensional boolean tensor. + dtype (`torch.dtype`, optional): + The desired data type of the output tensor. Defaults to `torch.long`. + + Returns: + `torch.Tensor`: + An (N-1)-dimensional tensor of integers indicating the position of the first True in each row. If no True + value is found in a row, returns the length of the row. + """ + row_len = bools.size(-1) + zero_or_index = row_len * (~bools).type(dtype) + torch.arange(row_len, dtype=dtype, device=bools.device) + return torch.min(zero_or_index, dim=-1).values + + +def get_reward( + model: torch.nn.Module, query_responses: torch.Tensor, pad_token_id: int, context_length: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Computes the reward logits and the rewards for a given model and query responses. + + Args: + model (`torch.nn.Module`): + The model used to compute the reward logits. + query_responses (`torch.Tensor`): + The tensor containing the query responses. + pad_token_id (`int`): + The token ID representing the pad token. + context_length (`int`): + The length of the context in the query responses. + + Returns: + tuple: + - `reward_logits` (`torch.Tensor`): + The logits for the reward model. + - `final_rewards` (`torch.Tensor`): + The final rewards for each query response. + - `sequence_lengths` (`torch.Tensor`): + The lengths of the sequences in the query responses. + """ + attention_mask = query_responses != pad_token_id + position_ids = attention_mask.cumsum(1) - attention_mask.long() # exclusive cumsum + lm_backbone = getattr(model, model.base_model_prefix) + input_ids = torch.masked_fill(query_responses, ~attention_mask, 0) + output = lm_backbone( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + return_dict=True, + output_hidden_states=True, + use_cache=False, # otherwise mistral-based RM would error out + ) + reward_logits = model.score(output.hidden_states[-1]) + sequence_lengths = first_true_indices(query_responses[:, context_length:] == pad_token_id) - 1 + context_length + # https://github.com/huggingface/transformers/blob/dc68a39c8111217683bf49a4912d0c9018bab33d/src/transformers/models/gpt2/modeling_gpt2.py#L1454 + return ( + reward_logits, + reward_logits[ + torch.arange(reward_logits.size(0), device=reward_logits.device), + sequence_lengths, + ].squeeze(-1), + sequence_lengths, + ) + + +def forward( + model: torch.nn.Module, + query_responses: torch.Tensor, + pad_token_id: int, +) -> ModelOutput: + """ + Performs a forward pass through the model with the given query responses and pad token ID. + + Args: + model (`torch.nn.Module`): + The model to perform the forward pass. + query_responses (`torch.Tensor`): + The tensor containing the query responses. + pad_token_id (`int`): + The token ID representing the pad token. + + Returns: + `ModelOutput`: + The output of the model, including hidden states. + """ + attention_mask = query_responses != pad_token_id + position_ids = attention_mask.cumsum(1) - attention_mask.long() + input_ids = torch.masked_fill(query_responses, ~attention_mask, 0) + return model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + return_dict=True, + output_hidden_states=True, + ) + + +def prepare_deepspeed( + model: torch.nn.Module, per_device_train_batch_size: int, fp16: bool = False, bf16: bool = False +) -> torch.nn.Module: + """ + Prepares the model for training with DeepSpeed (both for stage 2 and 3), configuring the appropriate settings based + on the model and batch size. + + Args: + model (`torch.nn.Module`): + The model to be prepared for DeepSpeed training. + per_device_train_batch_size (`int`): + The training batch size per device. + + Returns: + `torch.nn.Module`: + The model initialized and configured with DeepSpeed for training. + """ + import deepspeed + + deepspeed_plugin = AcceleratorState().deepspeed_plugin + config_kwargs = deepspeed_plugin.deepspeed_config + if config_kwargs["zero_optimization"]["stage"] != 3: + config_kwargs["train_micro_batch_size_per_gpu"] = per_device_train_batch_size + config_kwargs = { + "train_micro_batch_size_per_gpu": config_kwargs["train_micro_batch_size_per_gpu"], + "prescale_gradients": False, + "wall_clock_breakdown": False, + } + if bf16: + config_kwargs["bf16"] = {"enabled": True} + elif fp16: + config_kwargs["fp16"] = {"enabled": True} + else: + if hasattr(model, "config"): + hidden_size = ( + max(model.config.hidden_sizes) + if getattr(model.config, "hidden_sizes", None) + else getattr(model.config, "hidden_size", None) + ) + if hidden_size is not None and config_kwargs["zero_optimization"]["stage"] == 3: + # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache @ step 0: expected module 1, but got module 0` + # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081 + config_kwargs.update( + { + "zero_optimization.reduce_bucket_size": hidden_size * hidden_size, + "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size, + "zero_optimization.stage3_prefetch_bucket_size": 0, + } + ) + model, *_ = deepspeed.initialize(model=model, config=config_kwargs) + model.eval() + return model + + +def truncate_response(stop_token_id: int, pad_token_id: int, responses: torch.Tensor) -> torch.Tensor: + """ + Truncates the responses at the first occurrence of the stop token, filling the rest with pad tokens. + + Args: + stop_token_id (`int`): + The token ID representing the stop token where truncation occurs. + pad_token_id (`int`): + The token ID representing the pad token used to fill the truncated responses. + responses (`torch.Tensor`): + The tensor containing the responses to be truncated. + + Returns: + `torch.Tensor`: + The truncated responses tensor with pad tokens filled after the stop token. + """ + trunc_idxs = first_true_indices(responses == stop_token_id).unsqueeze(-1) + new_size = [1] * (len(responses.size()) - 1) + [responses.shape[1]] + idxs = torch.arange(responses.shape[1], device=responses.device).view(*new_size) + postprocessed_responses = torch.masked_fill(responses, idxs > trunc_idxs, pad_token_id) + return postprocessed_responses + + +def generate( + lm_backbone: torch.nn.Module, queries: torch.Tensor, pad_token_id: int, generation_config: GenerationConfig +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Generates sequences from the language model backbone in a way that does not affect padding tokens. + + Args: + lm_backbone (`torch.nn.Module`): + The language model backbone used for generation. + queries (`torch.Tensor`): + The tensor containing the input queries. + pad_token_id (`int`): + The token ID representing the pad token. + generation_config (`GenerationConfig`): + The configuration for the generation process. + + Returns: + tuple: + - `generated_sequences` (`torch.Tensor`): + The concatenated tensor of input queries and generated sequences. + - `logits` (`torch.Tensor`): + The logits output from the generation process. + """ + context_length = queries.shape[1] + attention_mask = queries != pad_token_id + input_ids = torch.masked_fill(queries, ~attention_mask, 0) + output = lm_backbone.generate( + input_ids=input_ids, + attention_mask=attention_mask, + # position_ids=attention_mask.cumsum(1) - attention_mask.long(), # not needed: already adjusted in generations + # https://github.com/huggingface/transformers/blob/ac33aeeeee2a7a89b89c93c2962e6feb90daef0a/src/transformers/models/gpt2/modeling_gpt2.py#L1227-L1250 + generation_config=generation_config, + return_dict_in_generate=True, + output_scores=True, + ) + logits = torch.stack(output.scores, 1) + return torch.cat((queries, output.sequences[:, context_length:]), dim=1), logits + + +@torch.no_grad() +def batch_generation( + model: torch.nn.Module, + queries: torch.Tensor, + local_rollout_forward_batch_size: int, + pad_token_id: int, + generation_config: GenerationConfig, +): + query_responses = [] + logitss = [] + batch_size = queries.shape[0] + for i in range(0, batch_size, local_rollout_forward_batch_size): + query = queries[i : i + local_rollout_forward_batch_size] + query_response, logits = generate( + model, + query, + pad_token_id, + generation_config, + ) + query_responses.append(query_response) + logitss.append(logits) + + # padding tensors + padded_query_responses = pad(query_responses, padding_value=pad_token_id, padding_side="right") + padded_logitss = pad(logitss, padding_value=0, padding_side="right") + + # reshaping + padded_query_responses = padded_query_responses.view(-1, padded_query_responses.shape[-1])[:batch_size] + padded_logitss = padded_logitss.view(-1, *padded_logitss.shape[2:])[:batch_size] + + return padded_query_responses, padded_logitss + + +def add_bos_token_if_needed( + bos_token_id: Optional[int], + prompt_len_input_ids: int, + prompt_tokens: dict[str, list[int]], + chosen_prompt_len_input_ids: int, + chosen_tokens: dict[str, list[int]], + rejected_prompt_len_input_ids: int, + rejected_tokens: dict[str, list[int]], +): + if bos_token_id is not None: + if prompt_len_input_ids == 0 or bos_token_id != prompt_tokens["prompt_input_ids"][0]: + prompt_tokens["prompt_input_ids"] = [bos_token_id] + prompt_tokens["prompt_input_ids"] + prompt_tokens["prompt_attention_mask"] = [1] + prompt_tokens["prompt_attention_mask"] + if chosen_prompt_len_input_ids == 0 or bos_token_id != chosen_tokens["prompt_input_ids"][0]: + chosen_tokens["prompt_input_ids"] = [bos_token_id] + chosen_tokens["prompt_input_ids"] + chosen_tokens["prompt_attention_mask"] = [1] + chosen_tokens["prompt_attention_mask"] + if rejected_prompt_len_input_ids == 0 or bos_token_id != rejected_tokens["prompt_input_ids"][0]: + rejected_tokens["prompt_input_ids"] = [bos_token_id] + rejected_tokens["prompt_input_ids"] + rejected_tokens["prompt_attention_mask"] = [1] + rejected_tokens["prompt_attention_mask"] + return prompt_tokens, chosen_tokens, rejected_tokens + + +def add_eos_token_if_needed( + eos_token_id: int, chosen_tokens: dict[str, list[int]], rejected_tokens: dict[str, list[int]] +): + if len(chosen_tokens["input_ids"]) == 0 or eos_token_id != chosen_tokens["input_ids"][-1]: + chosen_tokens["input_ids"].append(eos_token_id) + chosen_tokens["attention_mask"].append(1) + if len(rejected_tokens["input_ids"]) == 0 or eos_token_id != rejected_tokens["input_ids"][-1]: + rejected_tokens["input_ids"].append(eos_token_id) + rejected_tokens["attention_mask"].append(1) + return chosen_tokens, rejected_tokens + + +def truncate_right( + input_ids: torch.Tensor, stop_token_id: int, pad_token_id: int +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Truncates the input tensor from the right side after the first occurrence of the stop token. + + Args: + input_ids (`torch.Tensor`): + The tensor containing the responses to be truncated + stop_token_id (`int`): + The token ID representing the stop token where truncation occurs + pad_token_id (`int`): + The token ID representing the pad token used to fill the truncated responses + + Returns: + tuple: + - `output_ids` (`torch.Tensor`): + The truncated responses tensor with pad tokens filled after the stop token + - `mask` (`torch.Tensor`): + The mask tensor to indicate the padding tokens + """ + trunc_idxs = first_true_indices(input_ids == stop_token_id).unsqueeze(-1) + new_size = [1] * (len(input_ids.size()) - 1) + [input_ids.shape[1]] + idxs = torch.arange(input_ids.shape[1], device=input_ids.device).view(*new_size) + output_ids = torch.masked_fill(input_ids, idxs > trunc_idxs, pad_token_id) + mask = torch.masked_fill(torch.ones_like(input_ids), idxs > trunc_idxs, 0) + return output_ids, mask + + +def empty_cache() -> None: + """Empties the cache of the available torch device. + + This function checks for the availability of different torch devices (XPU, MLU, NPU, CUDA) and empties the cache of + the first available device it finds. + + If none of the specific devices are available, it defaults to emptying the CUDA cache. + """ + if is_torch_xpu_available(): + torch.xpu.empty_cache() + elif is_torch_mlu_available(): + torch.mlu.empty_cache() + elif is_torch_npu_available(): + torch.npu.empty_cache() + else: + torch.cuda.empty_cache() + + +def decode_and_strip_padding(inputs: torch.Tensor, tokenizer: PreTrainedTokenizerBase) -> list[str]: + """ + Decodes the input tensor and strips the padding tokens. + + Args: + inputs (`torch.Tensor`): + The input tensor to be decoded. + tokenizer (`transformers.PreTrainedTokenizerBase`): + The tokenizer used to decode the input tensor. + + Returns: + `list[str]`: + The list of decoded strings with padding tokens stripped. + """ + decoded = tokenizer.batch_decode(inputs, skip_special_tokens=False) + return [d.replace(tokenizer.pad_token, "") for d in decoded] + + +def generate_model_card( + base_model: Optional[str], + model_name: str, + hub_model_id: str, + dataset_name: Optional[str], + tags: list[str], + wandb_url: Optional[str], + trainer_name: str, + trainer_citation: Optional[str] = None, + paper_title: Optional[str] = None, + paper_id: Optional[str] = None, + comet_url: Optional[str] = None, +) -> ModelCard: + """ + Generate a `ModelCard` from a template. + + Args: + base_model (`str` or `None`): + Base model name. + model_name (`str`): + Model name. + hub_model_id (`str`): + Hub model ID as `username/model_id`. + dataset_name (`str` or `None`): + Dataset name. + tags (`list[str]`): + Tags. + wandb_url (`str` or `None`): + Weights & Biases run URL. + comet_url (`str` or `None`): + Comet experiment URL. + trainer_name (`str`): + Trainer name. + trainer_citation (`str` or `None`, defaults to `None`): + Trainer citation as a BibTeX entry. + paper_title (`str` or `None`, defaults to `None`): + Paper title. + paper_id (`str` or `None`, defaults to `None`): + ArXiv paper ID as `YYMM.NNNNN`. + + Returns: + `ModelCard`: + A ModelCard object. + """ + card_data = ModelCardData( + base_model=base_model, + datasets=dataset_name, + library_name="transformers", + licence="license", + model_name=model_name, + tags=["generated_from_trainer", *tags], + ) + card = ModelCard.from_template( + card_data, + template_path=str(pkg_resources.files("trl").joinpath("templates/lm_model_card.md")), + base_model=base_model, + model_name=model_name, + hub_model_id=hub_model_id, + dataset_name=dataset_name, + wandb_url=wandb_url, + comet_url=comet_url, + trainer_name=trainer_name, + trainer_citation=trainer_citation, + paper_title=paper_title, + paper_id=paper_id, + trl_version=version("trl"), + transformers_version=version("transformers"), + pytorch_version=version("torch"), + datasets_version=version("datasets"), + tokenizers_version=version("tokenizers"), + ) + return card + + +def get_comet_experiment_url() -> Optional[str]: + """ + If Comet integration is enabled, return the URL of the current Comet experiment; otherwise, return `None`. + """ + if not is_comet_available(): + return None + + if comet_ml.get_running_experiment() is not None: + return comet_ml.get_running_experiment().url + + return None + + +def log_table_to_comet_experiment(name: str, table: pd.DataFrame) -> None: + """ + If Comet integration is enabled logs a table to the Comet experiment if it is currently running. + + Args: + name (`str`): + Table name. + table (`pd.DataFrame`): + The Pandas DataFrame containing the table to log. + """ + if not is_comet_available(): + raise ModuleNotFoundError("The comet-ml is not installed. Please install it first: pip install comet-ml") + + experiment = comet_ml.get_running_experiment() + if experiment is not None: + experiment.log_table(tabular_data=table, filename=name) + + +def flush_left(mask: torch.Tensor, *tensors: torch.Tensor) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]: + """ + Shift non-zero elements in the mask and corresponding tensors to the left. + + This function operates on a binary mask and any number of additional tensors with the same dimensions as the mask. + For each row, non-zero values are shifted to the leftmost positions. Then, columns that contain only zeros across + all rows are truncated from the mask and tensors. Visually, this operation can be represented as follows: + + ``` + [[0, 0, x, x, x, x], -> [[x, x, x, x], + [0, x, x, x, 0, 0]] [x, x, x, 0]] + ``` + + Args: + mask (`torch.Tensor`): + 2D tensor (binary mask) with shape `(N, M)`. + *tensors (`torch.Tensor`) + One or more 2D tensors with the same shape as `mask`. These tensors will be processed alongside `mask`, + with non-zero values shifted and excess zero columns truncated in the same manner. + + Returns: + `torch.Tensor`: + Updated binary mask with non-zero values flushed to the left and trailing zero columns removed. + `*torch.Tensor` + Updated tensors, processed in the same way as the mask. + + Example: + ```python + >>> mask = torch.tensor([[0, 0, 1, 1, 1], [0, 1, 1, 0, 0]]) + >>> tensor = torch.tensor([[9, 9, 2, 3, 4], [9, 5, 6, 9, 9]]) + >>> new_mask, new_tensor = flush_left(mask, tensor) + >>> print(new_mask) + tensor([[1, 1, 1], + [1, 1, 0]]) + + >>> print(new_tensor) + tensor([[2, 3, 4], + [5, 6, 0]]) + ``` + """ + _, M = mask.shape + + # Create copy of mask and tensors + mask_copy = mask.clone() + tensors = [t.clone() for t in tensors] + + # Shift non-zero values to the left + first_non_zero = mask_copy.argmax(dim=1) + pos = torch.arange(M, device=mask_copy.device).unsqueeze(0) + idx_roll = (pos + first_non_zero.unsqueeze(1)) % M + mask_roll = mask_copy.gather(1, idx_roll) + rolled_tensors = [t.gather(1, idx_roll) for t in tensors] + + # Truncate trailing columns that are all zeros in mask_roll + col_sums = mask_roll.sum(dim=0) + empty_cols = col_sums == 0 + first_empty_col = int(empty_cols.to(torch.int8).argmax()) if empty_cols.any() else M + flushed_mask = mask_roll[:, :first_empty_col] + flushed_tensors = [t[:, :first_empty_col] for t in rolled_tensors] + + if not flushed_tensors: + return flushed_mask + return flushed_mask, *flushed_tensors + + +def flush_right(mask: torch.Tensor, *tensors: torch.Tensor) -> Union[torch.Tensor, tuple[torch.Tensor, ...]]: + """ + Shift non-zero elements in the mask and corresponding tensors to the right. See `flush_left` for details. + """ + _, M = mask.shape + + # Create copy of mask and tensors + mask_copy = mask.clone() + tensors = [t.clone() for t in tensors] + + # Shift non-zero values to the right + flipped_mask = torch.fliplr(mask_copy) + first_non_zero = flipped_mask.argmax(dim=1) + pos = torch.arange(M, device=mask_copy.device).unsqueeze(0) + idx_roll = (pos - first_non_zero.unsqueeze(1)) % M + mask_roll = mask_copy.gather(1, idx_roll) + rolled_tensors = [t.gather(1, idx_roll) for t in tensors] + + # Truncate leading columns that are all zeros in mask_roll + col_sums = mask_roll.sum(dim=0) + non_empty_cols = col_sums != 0 + first_non_empty_col = int(non_empty_cols.to(torch.int8).argmax()) if non_empty_cols.any() else M + flushed_mask = mask_roll[:, first_non_empty_col:] + flushed_tensors = [t[:, first_non_empty_col:] for t in rolled_tensors] + + if not flushed_tensors: + return flushed_mask + return flushed_mask, *flushed_tensors + + +def selective_log_softmax(logits, index) -> torch.Tensor: + """ + A memory-efficient implementation of the common `log_softmax -> gather` operation. + + This function is equivalent to the following naive implementation: + ```python + logps = torch.gather(logits.log_softmax(-1), dim=-1, index=index.unsqueeze(-1)).squeeze(-1) + ``` + + Args: + logits (`torch.Tensor`): + Logits tensor of shape `(..., num_classes)`. + index (`torch.Tensor`): + Index tensor of shape `(...)`, specifying the positions to gather from the log-softmax output. + + Returns: + `torch.Tensor`: + Gathered log probabilities with the same shape as `index`. + """ + if logits.dtype in [torch.float32, torch.float64]: + selected_logits = torch.gather(logits, dim=-1, index=index.unsqueeze(-1)).squeeze(-1) + # loop to reduce peak mem consumption + logsumexp_values = torch.stack([torch.logsumexp(lg, dim=-1) for lg in logits]) + per_token_logps = selected_logits - logsumexp_values # log_softmax(x_i) = x_i - logsumexp(x) + else: + # logsumexp approach is unstable with bfloat16, fall back to slightly less efficient approach + per_token_logps = [] + for row_logits, row_labels in zip(logits, index): # loop to reduce peak mem consumption + row_logps = F.log_softmax(row_logits, dim=-1) + row_per_token_logps = row_logps.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(-1) + per_token_logps.append(row_per_token_logps) + per_token_logps = torch.stack(per_token_logps) + return per_token_logps + + +def print_prompt_completions_sample( + prompts: list[str], + completions: list[str], + rewards: dict[str, list[float]], + advantages: list[float], + step: int, + num_samples: int = None, +) -> None: + """ + Print out a sample of model completions to the console with multiple reward metrics. + + This function creates a nicely formatted table showing prompt-completion pairs, useful for monitoring model outputs + during training. It requires the `rich` library to be installed. + + Args: + prompts (`list[str]`): + List of prompts. + completions (`list[str]`): + List of completions corresponding to the prompts. + rewards (`dict[str, list[float]]`): + Dictionary where keys are reward names and values are lists of rewards. + advantages (`list[float]`): + List of advantages corresponding to the prompts and completions. + step (`int`): + Current training step number, used in the output title. + num_samples (`int` or `None`, *optional*, defaults to `None`): + Number of random samples to display. If `None` (default), all items will be displayed. + + Example: + ```python + >>> from trl.trainer.utils import print_prompt_completions_sample + + >>> prompts = ["The sky is", "The sun is"] + >>> completions = [" blue.", " in the sky."] + >>> rewards = {"Correctness": [0.123, 0.456], "Format": [0.789, 0.101]} + >>> advantages = [0.987, 0.654] + >>> print_prompt_completions_sample(prompts, completions, rewards, advantages, 42) + ╭──────────────────────────── Step 42 ─────────────────────────────╮ + │ ┏━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━┓ │ + │ ┃ Prompt ┃ Completion ┃ Correctness ┃ Format ┃ Advantage ┃ │ + │ ┡━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━┩ │ + │ │ The sky is │ blue. │ 0.12 │ 0.79 │ 0.99 │ │ + │ ├────────────┼──────────────┼─────────────┼────────┼───────────┤ │ + │ │ The sun is │ in the sky. │ 0.46 │ 0.10 │ 0.65 │ │ + │ └────────────┴──────────────┴─────────────┴────────┴───────────┘ │ + ╰──────────────────────────────────────────────────────────────────╯ + ``` + """ + if not is_rich_available(): + raise ImportError( + "The function `print_prompt_completions_sample` requires the `rich` library. Please install it with " + "`pip install rich`." + ) + console = Console() + table = Table(show_header=True, header_style="bold white", expand=True) + + # Add columns + table.add_column("Prompt", style="bright_yellow") + table.add_column("Completion", style="bright_green") + for reward_name in rewards.keys(): + table.add_column(reward_name, style="bold cyan", justify="right") + table.add_column("Advantage", style="bold magenta", justify="right") + + # Some basic input validation + if num_samples is not None: + if num_samples >= len(prompts): + num_samples = None + elif num_samples <= 0: + return + + # Subsample data if num_samples is specified + if num_samples is not None: + indices = random.sample(range(len(prompts)), num_samples) + prompts = [prompts[i] for i in indices] + completions = [completions[i] for i in indices] + rewards = {key: [val[i] for i in indices] for key, val in rewards.items()} + advantages = [advantages[i] for i in indices] + + for i in range(len(prompts)): + reward_values = [f"{rewards[key][i]:.2f}" for key in rewards.keys()] # 2 decimals + table.add_row(Text(prompts[i]), Text(completions[i]), *reward_values, f"{advantages[i]:.2f}") + table.add_section() # Adds a separator between rows + + panel = Panel(table, expand=False, title=f"Step {step}", border_style="bold white") + console.print(panel) diff --git a/lib/python3.12/site-packages/trl/trainer/xpo_config.py b/lib/python3.12/site-packages/trl/trainer/xpo_config.py new file mode 100644 index 0000000000000000000000000000000000000000..78c1e562128208e0b6a2ba6e78f9becd59e8a1ac --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/xpo_config.py @@ -0,0 +1,44 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass, field + +from trl.trainer.online_dpo_config import OnlineDPOConfig + + +@dataclass +class XPOConfig(OnlineDPOConfig): + r""" + Configuration class for the [`XPOTrainer`]. + + Subclass of [`OnlineDPOConfig`] we can use all its arguments and add the following: + + Parameters: + alpha (`float` or `list[float]`, *optional*, defaults to `1e-5`): + Weight of the XPO loss term. If a list of floats is provided then the alpha is selected for each new epoch + and the last alpha is used for the rest of the epochs. + """ + + alpha: list[float] = field( + default_factory=lambda: [1e-5], + metadata={ + "help": "Weight of the XPO loss term. If a list of floats is provided then the alpha is selected for each " + "new epoch and the last alpha is used for the rest of the epochs." + }, + ) + + def __post_init__(self): + super().__post_init__() + if hasattr(self.alpha, "__len__") and len(self.alpha) == 1: + self.alpha = self.alpha[0] diff --git a/lib/python3.12/site-packages/trl/trainer/xpo_trainer.py b/lib/python3.12/site-packages/trl/trainer/xpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..db557419c619de148623ee6a2f564be198477844 --- /dev/null +++ b/lib/python3.12/site-packages/trl/trainer/xpo_trainer.py @@ -0,0 +1,595 @@ +# Copyright 2020-2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import textwrap +from typing import Any, Callable, Optional, Union + +import jinja2 +import torch +import torch.nn as nn +import torch.nn.functional as F +from datasets import Dataset, IterableDataset +from transformers import ( + BaseImageProcessor, + FeatureExtractionMixin, + PreTrainedModel, + PreTrainedTokenizerBase, + ProcessorMixin, + TrainerCallback, + is_apex_available, + is_wandb_available, +) +from transformers.trainer_utils import EvalPrediction +from transformers.training_args import OptimizerNames +from transformers.utils import is_peft_available + +from ..data_utils import is_conversational, maybe_apply_chat_template +from ..models.utils import unwrap_model_for_generation +from .judges import BasePairwiseJudge +from .online_dpo_trainer import OnlineDPOTrainer +from .utils import ( + SIMPLE_CHAT_TEMPLATE, + empty_cache, + generate_model_card, + get_comet_experiment_url, + get_reward, + selective_log_softmax, + truncate_right, +) +from .xpo_config import XPOConfig + + +if is_apex_available(): + from apex import amp + + +if is_wandb_available(): + import wandb + + +if is_peft_available(): + from peft import PeftModel + + +class XPOTrainer(OnlineDPOTrainer): + r""" + Initialize XPOTrainer as a subclass of [`OnlineDPOConfig`]. + + Args: + model (`transformers.PreTrainedModel`): + The model to train, preferably an `AutoModelForCausalLM`. + ref_model (`PreTrainedModelWrapper`): + Hugging Face transformer model with a casual language modelling head. Used for implicit reward computation + and loss. If no reference model is provided, the trainer will create a reference model with the same + architecture as the model to be optimized. + reward_model (`transformers.PreTrainedModel`): + The reward model to score completions with, preferably an `AutoModelForSequenceClassification`. + judge (`BasePairwiseJudge`): + The judge to use for pairwise comparison of model completions. + args (`XPOConfig`): + The XPO config arguments to use for training. + data_collator (`transformers.DataCollator`): + The data collator to use for training. If None is specified, the default data collator + (`DPODataCollatorWithPadding`) will be used which will pad the sequences to the maximum length of the + sequences in the batch, given a dataset of paired sequences. + train_dataset (`datasets.Dataset`): + The dataset to use for training. + eval_dataset (`datasets.Dataset`): + The dataset to use for evaluation. + processing_class (`PreTrainedTokenizerBase` or `BaseImageProcessor` or `FeatureExtractionMixin` or `ProcessorMixin`, *optional*): + Processing class used to process the data. If provided, will be used to automatically process the inputs + for the model, and it will be saved along the model to make it easier to rerun an interrupted training or + reuse the fine-tuned model. + peft_config (`dict`): + The peft config to use for training. + compute_metrics (`Callable[[EvalPrediction], dict]`, *optional*): + The function to use to compute the metrics. Must take a `EvalPrediction` and return a dictionary string to + metric values. + callbacks (`list[transformers.TrainerCallback]`): + The callbacks to use for training. + optimizers (`tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR]`): + The optimizer and scheduler to use for training. + preprocess_logits_for_metrics (`Callable[[torch.Tensor, torch.Tensor], torch.Tensor]`): + The function to use to preprocess the logits before computing the metrics. + """ + + _tag_names = ["trl", "xpo"] + + def __init__( + self, + model: Union[PreTrainedModel, nn.Module] = None, + ref_model: Union[PreTrainedModel, nn.Module] = None, + reward_model: Optional[nn.Module] = None, + judge: Optional[BasePairwiseJudge] = None, + args: Optional[XPOConfig] = None, + data_collator: Optional[Callable] = None, + train_dataset: Optional[Union[Dataset, IterableDataset]] = None, + eval_dataset: Optional[Union[Dataset, dict[str, Dataset]]] = None, + processing_class: Optional[ + Union[PreTrainedTokenizerBase, BaseImageProcessor, FeatureExtractionMixin, ProcessorMixin] + ] = None, + peft_config: Optional[dict] = None, + compute_metrics: Optional[Callable[[EvalPrediction], dict]] = None, + callbacks: Optional[list[TrainerCallback]] = None, + optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None), + preprocess_logits_for_metrics: Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None, + ) -> None: + super().__init__( + model=model, + ref_model=ref_model, + judge=judge, + reward_model=reward_model, + args=args, + data_collator=data_collator, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + processing_class=processing_class, + reward_processing_class=processing_class, # for now, XPOTrainer can't use any reward model + peft_config=peft_config, + compute_metrics=compute_metrics, + callbacks=callbacks, + optimizers=optimizers, + preprocess_logits_for_metrics=preprocess_logits_for_metrics, + ) + + self._alpha = self.args.alpha + + # Overwrite the stats dictionary to include XPO specific statistics + self.stats = { + # Remove "non_score_reward", "rlhf_reward", "scores" + # Add "loss/dpo", "loss/xpo" + "loss/dpo": [], + "loss/xpo": [], + "objective/kl": [], + "objective/entropy": [], + "rewards/chosen": [], + "rewards/rejected": [], + "rewards/accuracies": [], + "rewards/margins": [], + "logps/chosen": [], + "logps/rejected": [], + # Replace "contain_eos_token" by "model_contain_eos_token" and "ref_contain_eos_token" + "val/model_contain_eos_token": [], + "val/ref_contain_eos_token": [], + "alpha": [], + "beta": [], + } + if self.reward_model is not None: + # Replace "scores" by "model_scores" and "ref_scores" + self.stats["objective/model_scores"] = [] + self.stats["objective/ref_scores"] = [] + self.stats["objective/scores_margin"] = [] + + @property + def alpha(self): + if isinstance(self._alpha, list): + epoch = self.state.epoch + return self._alpha[epoch] if epoch < len(self._alpha) else self._alpha[-1] + else: + return self._alpha + + def _generate_completions(self, prompts, model): + with unwrap_model_for_generation(model, self.accelerator) as unwrapped_policy_model_for_gen: + model_output = unwrapped_policy_model_for_gen.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + actual_model_for_ref_generation: torch.nn.Module + if self.ref_model is None: + unwrapped_main_model_for_ref_logic = self.accelerator.unwrap_model(model) + + if is_peft_available() and isinstance(unwrapped_main_model_for_ref_logic, PeftModel): + actual_model_for_ref_generation = unwrapped_main_model_for_ref_logic.get_base_model() + else: + actual_model_for_ref_generation = unwrapped_main_model_for_ref_logic + else: + actual_model_for_ref_generation = self.accelerator.unwrap_model(self.ref_model) + + with unwrap_model_for_generation(actual_model_for_ref_generation, self.accelerator) as final_ref_model_for_gen: + ref_output = final_ref_model_for_gen.generate( + input_ids=prompts["input_ids"], + attention_mask=prompts["attention_mask"], + generation_config=self.generation_config, + ) + + return model_output, ref_output + + def _process_completions(self, model_output, ref_output, prompts): + context_length = prompts["input_ids"].shape[1] + + # Process model completions + model_completion_ids = model_output[:, context_length:] + model_completion_ids, model_completion_mask = truncate_right( + model_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + model_data = { + "input_ids": torch.cat((prompts["input_ids"], model_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], model_completion_mask), dim=1), + "raw": prompts["raw"], + } + + # Process reference model completions + ref_completion_ids = ref_output[:, context_length:] + ref_completion_ids, ref_completion_mask = truncate_right( + ref_completion_ids, self.processing_class.eos_token_id, self.processing_class.pad_token_id + ) + ref_data = { + "input_ids": torch.cat((prompts["input_ids"], ref_completion_ids), dim=1), + "attention_mask": torch.cat((prompts["attention_mask"], ref_completion_mask), dim=1), + "raw": prompts["raw"], + } + + return model_data, ref_data + + def _compute_rewards(self, model_data, ref_data, context_length): + with torch.no_grad(): + _, model_scores, _ = get_reward( + self.reward_model, model_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + _, ref_scores, _ = get_reward( + self.reward_model, ref_data["input_ids"], self.processing_class.pad_token_id, context_length + ) + + # Apply EOS penalty if needed + if self.args.missing_eos_penalty is not None: + model_contain_eos = torch.any(model_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + ref_contain_eos = torch.any(ref_data["input_ids"] == self.processing_class.eos_token_id, dim=-1) + model_scores[~model_contain_eos] -= self.args.missing_eos_penalty + ref_scores[~ref_contain_eos] -= self.args.missing_eos_penalty + + return model_scores, ref_scores + + def _compute_judge(self, model_data, ref_data, context_length): + prompts = model_data["raw"] + model_data_completions = self.processing_class.batch_decode( + model_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + model_data_completions = [completion.strip() for completion in model_data_completions] + + ref_data_completions = self.processing_class.batch_decode( + ref_data["input_ids"][:, context_length:], skip_special_tokens=True + ) + ref_data_completions = [completion.strip() for completion in ref_data_completions] + + if is_conversational({"prompt": prompts[0]}): + model_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in model_data_completions + ] + environment = jinja2.Environment() + template = environment.from_string(SIMPLE_CHAT_TEMPLATE) + prompts = [template.render(messages=message) for message in prompts] + model_data_completions = [template.render(messages=completion) for completion in model_data_completions] + + ref_data_completions = [ + [{"role": "assistant", "content": completion}] for completion in ref_data_completions + ] + ref_data_completions = [template.render(messages=completion) for completion in ref_data_completions] + + ranks_of_first_completion = self.judge.judge( + prompts, + list(zip(model_data_completions, ref_data_completions)), + ) + # convert ranks to a True/False mask: + # when rank == 0, it means the first completion is the best + # when rank == 1, it means the second completion is the best + return torch.tensor([rank == 0 for rank in ranks_of_first_completion], device=model_data["input_ids"].device) + + def _compute_logprobs(self, model, model_data, ref_data, context_length): + def compute_logprobs_for_data(m, data): + output = m(data["input_ids"], attention_mask=data["attention_mask"]) + logits = output.logits[:, context_length - 1 : -1] + token_logprobs = selective_log_softmax(logits, data["input_ids"][:, context_length:]) + return token_logprobs + + # Compute logprobs for model completions + model_logprobs_model_data = compute_logprobs_for_data(model, model_data) + # Compute logprobs for model on reference completions (for XPO loss) + model_logprobs_ref_data = compute_logprobs_for_data(model, ref_data) + + # Compute logprobs for reference model completions + with torch.no_grad(): + if self.ref_model is None: + with model.disable_adapter(): + ref_logprobs_model_data = compute_logprobs_for_data(model, model_data) + ref_logprobs_ref_data = compute_logprobs_for_data(model, ref_data) + else: + ref_logprobs_model_data = compute_logprobs_for_data(self.ref_model, model_data) + ref_logprobs_ref_data = compute_logprobs_for_data(self.ref_model, ref_data) + + # Mask padding tokens + model_padding_mask = model_data["attention_mask"][:, context_length:] == 0 + ref_padding_mask = ref_data["attention_mask"][:, context_length:] == 0 + model_logprobs_model_data = model_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + model_logprobs_ref_data = model_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0) + ref_logprobs_ref_data = ref_logprobs_ref_data.masked_fill(ref_padding_mask, 0.0) + ref_logprobs_model_data = ref_logprobs_model_data.masked_fill(model_padding_mask, 0.0) + + return model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data + + def _compute_losses( + self, + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + ): + # Compute log probs + model_logprobs_model_data_sum = model_logprobs_model_data.sum(1) + model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1) + ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1) + ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1) + + chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs + + rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs + + # Compute logits as the difference between chosen and rejected log ratios + logits = chosen_log_ratios - rejected_log_ratios + + if self.args.loss_type == "sigmoid": + dpo_losses = -F.logsigmoid(self.beta * logits) + elif self.args.loss_type == "ipo": + dpo_losses = (logits - 1 / (2 * self.beta)) ** 2 + else: + raise NotImplementedError(f"invalid loss type {self.args.loss_type}") + + # Compute XPO specific loss + xpo_losses = self.alpha * model_logprobs_ref_data_sum + + # Total loss + loss = (dpo_losses + xpo_losses).mean() + + return loss, dpo_losses, xpo_losses + + def _log_statistics( + self, + model_data, + ref_data, + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + dpo_losses, + xpo_losses, + context_length, + model_scores=None, + ref_scores=None, + ): + # Helper function to gather and compute mean + def gather_mean(tensor): + return self.accelerator.gather_for_metrics(tensor).mean().item() + + # Log losses + self.stats["loss/dpo"].append(gather_mean(dpo_losses)) + self.stats["loss/xpo"].append(gather_mean(xpo_losses)) + + # Log scores + if self.reward_model is not None: + self.stats["objective/model_scores"].append(gather_mean(model_scores)) + self.stats["objective/ref_scores"].append(gather_mean(ref_scores)) + self.stats["objective/scores_margin"].append(gather_mean(model_scores - ref_scores)) + + # Log logprobs + model_logprobs_model_data_sum = model_logprobs_model_data.sum(1) + model_logprobs_ref_data_sum = model_logprobs_ref_data.sum(1) + ref_logprobs_ref_data_sum = ref_logprobs_ref_data.sum(1) + ref_logprobs_model_data_sum = ref_logprobs_model_data.sum(1) + + chosen_model_logprobs = torch.where(chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + chosen_ref_logprobs = torch.where(chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + chosen_log_ratios = chosen_model_logprobs - chosen_ref_logprobs + + rejected_model_logprobs = torch.where(~chosen_mask, model_logprobs_model_data_sum, model_logprobs_ref_data_sum) + rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum) + rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs + + self.stats["logps/chosen"].append(gather_mean(chosen_model_logprobs.mean() + chosen_ref_logprobs.mean())) + self.stats["logps/rejected"].append(gather_mean(rejected_model_logprobs.mean() + rejected_ref_logprobs.mean())) + + # Log rewards + # Compute various statistics + chosen_rewards = chosen_log_ratios * self.beta + rejected_rewards = rejected_log_ratios * self.beta + self.stats["rewards/chosen"].append(gather_mean(chosen_rewards.mean())) + self.stats["rewards/rejected"].append(gather_mean(rejected_rewards.mean())) + + # Calculate KL divergence for model and ref data + kl_model_data = model_logprobs_model_data - ref_logprobs_model_data + kl_ref_data = model_logprobs_ref_data - ref_logprobs_ref_data + mean_kl = (kl_model_data.sum(1) + kl_ref_data.sum(1)).mean() / 2 + self.stats["objective/kl"].append(gather_mean(mean_kl)) + + # Calculate entropy for model and ref data + entropy_model_data = -model_logprobs_model_data.sum(1) + entropy_ref_data = -model_logprobs_ref_data.sum(1) + mean_entropy = (entropy_model_data.mean() + entropy_ref_data.mean()) / 2 + self.stats["objective/entropy"].append(gather_mean(mean_entropy)) + + # Calculate margins + margin = chosen_rewards - rejected_rewards + self.stats["rewards/margins"].append(gather_mean(margin.mean())) + + # Calculate accuracy + accuracy = (margin > 0).float() + self.stats["rewards/accuracies"].append(gather_mean(accuracy.mean())) + + # Log EOS token statistics + model_eos = (model_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + ref_eos = (ref_data["input_ids"][:, context_length:] == self.processing_class.eos_token_id).any(dim=1) + self.stats["val/model_contain_eos_token"].append(gather_mean(model_eos.float())) + self.stats["val/ref_contain_eos_token"].append(gather_mean(ref_eos.float())) + + # Log alpha and beta + self.stats["alpha"].append(self.alpha) + self.stats["beta"].append(self.beta) + + def training_step( + self, model: nn.Module, inputs: dict[str, Union[torch.Tensor, Any]], num_items_in_batch: Optional[int] = None + ) -> torch.Tensor: + model.train() + + # Apply chat template and tokenize the input + batch_size = len(next(iter(inputs.values()))) + prompts = inputs["prompt"] + inputs = [{k: v[i] for k, v in inputs.items()} for i in range(batch_size)] + inputs = [maybe_apply_chat_template(x, self.processing_class) for x in inputs] + inputs = [self.tokenize_row(x, self.model.config.is_encoder_decoder, self.processing_class) for x in inputs] + inputs = self.data_collator(inputs) + + # need the prompt_ only + inputs = self._prepare_inputs(inputs) + context_length = inputs["prompt_input_ids"].shape[1] + prompts = { + "input_ids": inputs["prompt_input_ids"], + "attention_mask": inputs["prompt_attention_mask"], + "raw": prompts, + } + del inputs + + # Sample completions from both the model and the reference model + model_output, ref_output = self._generate_completions(prompts, model) + + # Process model completions + model_data, ref_data = self._process_completions(model_output, ref_output, prompts) + + # Compute rewards + if self.reward_model is not None: + model_scores, ref_scores = self._compute_rewards(model_data, ref_data, context_length) + chosen_mask = model_scores >= ref_scores + else: + model_scores, ref_scores = None, None + chosen_mask = self._compute_judge(model_data, ref_data, context_length) + + # Compute logprobs + model_logprobs_model_data, model_logprobs_ref_data, ref_logprobs_ref_data, ref_logprobs_model_data = ( + self._compute_logprobs(model, model_data, ref_data, context_length) + ) + + # Compute loss + loss, dpo_losses, xpo_losses = self._compute_losses( + model_logprobs_model_data, + model_logprobs_ref_data, + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + ) + + # Log everything + self._log_statistics( + model_data, + ref_data, + model_logprobs_model_data.detach(), + model_logprobs_ref_data.detach(), + ref_logprobs_ref_data, + ref_logprobs_model_data, + chosen_mask, + dpo_losses.detach(), + xpo_losses.detach(), + context_length, + model_scores, + ref_scores, + ) + + if ( + self.args.torch_empty_cache_steps is not None + and self.state.global_step % self.args.torch_empty_cache_steps == 0 + ): + empty_cache() + + kwargs = {} + # For LOMO optimizers you need to explicitly use the learning rate + if self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: + kwargs["learning_rate"] = self._get_learning_rate() + + if self.args.n_gpu > 1: + loss = loss.mean() # mean() to average on multi-gpu parallel training + + if self.use_apex: + with amp.scale_loss(loss, self.optimizer) as scaled_loss: + scaled_loss.backward() + else: + self.accelerator.backward(loss, **kwargs) + + return loss.detach() / self.args.gradient_accumulation_steps + + def create_model_card( + self, + model_name: Optional[str] = None, + dataset_name: Optional[str] = None, + tags: Union[str, list[str], None] = None, + ): + """ + Creates a draft of a model card using the information available to the `Trainer`. + + Args: + model_name (`str` or `None`, *optional*, defaults to `None`): + Name of the model. + dataset_name (`str` or `None`, *optional*, defaults to `None`): + Name of the dataset used for training. + tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): + Tags to be associated with the model card. + """ + if not self.is_world_process_zero(): + return + + if hasattr(self.model.config, "_name_or_path") and not os.path.isdir(self.model.config._name_or_path): + base_model = self.model.config._name_or_path + else: + base_model = None + + # normalize `tags` to a mutable set + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + else: + tags = set(tags) + + if hasattr(self.model.config, "unsloth_version"): + tags.add("unsloth") + + tags.update(self._tag_names) + + citation = textwrap.dedent("""\ + @article{jung2024binary, + title = {{Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF}}, + author = {Tengyang Xie and Dylan J. Foster and Akshay Krishnamurthy and Corby Rosset and Ahmed Awadallah and Alexander Rakhlin}, + year = 2024, + eprint = {arXiv:2405.21046} + }""") + + model_card = generate_model_card( + base_model=base_model, + model_name=model_name, + hub_model_id=self.hub_model_id, + dataset_name=dataset_name, + tags=tags, + wandb_url=wandb.run.get_url() if is_wandb_available() and wandb.run is not None else None, + comet_url=get_comet_experiment_url(), + trainer_name="XPO", + trainer_citation=citation, + paper_title="Exploratory Preference Optimization: Harnessing Implicit Q*-Approximation for Sample-Efficient RLHF", + paper_id="2405.21046", + ) + + model_card.save(os.path.join(self.args.output_dir, "README.md")) diff --git a/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/INSTALLER b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/METADATA b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..1640d911c8a23ed7dca476ed58a962b1923e4473 --- /dev/null +++ b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/METADATA @@ -0,0 +1,70 @@ +Metadata-Version: 2.1 +Name: xgrammar +Version: 0.1.18 +Summary: Efficient, Flexible and Portable Structured Generation +Keywords: machine learning,inference +Author: MLC Team +License: Apache 2.0 +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Development Status :: 4 - Beta +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Education +Classifier: Intended Audience :: Science/Research +Project-URL: Homepage, https://xgrammar.mlc.ai/ +Project-URL: GitHub, https://github.com/mlc-ai/xgrammar +Requires-Python: <4,>=3.8 +Requires-Dist: pydantic +Requires-Dist: sentencepiece +Requires-Dist: tiktoken +Requires-Dist: torch>=1.10.0 +Requires-Dist: transformers>=4.38.0 +Requires-Dist: triton; platform_system == "Linux" and platform_machine == "x86_64" +Requires-Dist: mlx-lm; platform_system == "Darwin" and platform_machine == "arm64" +Requires-Dist: ninja +Provides-Extra: test +Requires-Dist: pytest; extra == "test" +Requires-Dist: protobuf; extra == "test" +Requires-Dist: huggingface-hub[cli]; extra == "test" +Requires-Dist: transformers<4.50.0; platform_system == "Darwin" and extra == "test" +Description-Content-Type: text/markdown + +
+ +# XGrammar + +[![Documentation](https://img.shields.io/badge/docs-latest-green)](https://xgrammar.mlc.ai/docs/) +[![License](https://img.shields.io/badge/license-apache_2-blue)](https://github.com/mlc-ai/xgrammar/blob/main/LICENSE) +[![PyPI](https://img.shields.io/pypi/v/xgrammar)](https://pypi.org/project/xgrammar) +[![PyPI Downloads](https://static.pepy.tech/badge/xgrammar)](https://pepy.tech/projects/xgrammar) + +**Efficient, Flexible and Portable Structured Generation** + + +[Get Started](#get-started) | [Documentation](https://xgrammar.mlc.ai/docs/) | [Blogpost](https://blog.mlc.ai/2024/11/22/achieving-efficient-flexible-portable-structured-generation-with-xgrammar) | [Technical Report](https://arxiv.org/abs/2411.15100) + +
+ +## News +- [2025/02] XGrammar has been officially integrated into [Modular's MAX](https://docs.modular.com/max/serve/structured-output) +- [2025/01] XGrammar has been officially integrated into [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM). +- [2024/12] XGrammar has been officially integrated into [vLLM](https://github.com/vllm-project/vllm). +- [2024/12] We presented research talks on XGrammar at CMU Catalyst, Berkeley SkyLab, MIT HANLAB, THU IIIS, SJTU, Ant Group, SGLang Meetup, Qingke AI, Camel AI. The slides can be found [here](https://docs.google.com/presentation/d/1iS7tu2EV4IKRWDaR0F3YD7ubrNqtGYUStSskceneelc/edit?usp=sharing). +- [2024/11] XGrammar has been officially integrated into [SGLang](https://github.com/sgl-project/sglang). +- [2024/11] XGrammar has been officially integrated into [MLC-LLM](https://github.com/mlc-ai/mlc-llm). +- [2024/11] We officially released XGrammar v0.1.0! + +## Overview + + +XGrammar is an open-source library for efficient, flexible, and portable structured generation. +It supports general context-free grammar to enable a broad range of structures while bringing careful system optimizations to enable fast executions. +XGrammar features a minimal and portable C++ backend that can be easily integrated into multiple environments and frameworks, +and is co-designed with the LLM inference engine and enables zero-overhead structured generation in LLM inference. + + + +## Get Started + +Please visit our [documentation](https://xgrammar.mlc.ai/docs/) to get started with XGrammar. +- [Installation](https://xgrammar.mlc.ai/docs/start/install) +- [Quick start](https://xgrammar.mlc.ai/docs/start/quick_start) diff --git a/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/RECORD b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..f9da594f6c71b94801b9b52664d3735080b6c4c1 --- /dev/null +++ b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/RECORD @@ -0,0 +1,47 @@ +xgrammar-0.1.18.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +xgrammar-0.1.18.dist-info/METADATA,sha256=gknu2i0vD8g7MaC-gSTwOdDxJcyjPDRFVKjel1izUq0,3584 +xgrammar-0.1.18.dist-info/RECORD,, +xgrammar-0.1.18.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xgrammar-0.1.18.dist-info/WHEEL,sha256=pb4lHaPmHh87DTkQaQ89Cq9PHCFaV8glEyeKll_RPs4,156 +xgrammar-0.1.18.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357 +xgrammar-0.1.18.dist-info/licenses/NOTICE,sha256=dpxpBwc8ZYAb_P484cDoC_CuThJOJ2nU8CbAOK1YPA8,54 +xgrammar/__init__.py,sha256=vr-8u4CcpaFRx89GPv1Y2Pyjzudo6zDuwYs9Dz0ZVP8,375 +xgrammar/__pycache__/__init__.cpython-312.pyc,, +xgrammar/__pycache__/base.cpython-312.pyc,, +xgrammar/__pycache__/compiler.cpython-312.pyc,, +xgrammar/__pycache__/grammar.cpython-312.pyc,, +xgrammar/__pycache__/matcher.cpython-312.pyc,, +xgrammar/__pycache__/testing.cpython-312.pyc,, +xgrammar/__pycache__/tokenizer_info.cpython-312.pyc,, +xgrammar/__pycache__/version.cpython-312.pyc,, +xgrammar/base.py,sha256=P1vjpdTGijQGy25bKiawFCqMcxYm_HqAO7WFUjiiuJ8,2360 +xgrammar/compiler.py,sha256=UCb9Fol_NBCFTDXhimc-J3P2X8kgkNrNWVflBoL1TqA,7728 +xgrammar/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xgrammar/contrib/__pycache__/__init__.cpython-312.pyc,, +xgrammar/contrib/__pycache__/hf.cpython-312.pyc,, +xgrammar/contrib/__pycache__/mlxlm.cpython-312.pyc,, +xgrammar/contrib/hf.py,sha256=nAoqTdxWhP7uWEba1HAsSyrn1MYZYvKRudDkrerIH-k,4614 +xgrammar/contrib/mlxlm.py,sha256=eoZ5baxaIJplP0ohiBY6AjofHbjfn6J5v9sFWi-lwpw,2812 +xgrammar/grammar.py,sha256=VgyNnAZFhLHHejO2W9OgWtn7ON5l7iitpBVxyestL3Y,13072 +xgrammar/kernels/__init__.py,sha256=Rz8sBcZMbQeYZiKoxvG1u7jteQj0S7gOJ3MCOfJWRX0,259 +xgrammar/kernels/__pycache__/__init__.cpython-312.pyc,, +xgrammar/kernels/__pycache__/apply_token_bitmask_inplace_cpu.cpython-312.pyc,, +xgrammar/kernels/__pycache__/apply_token_bitmask_inplace_cuda.cpython-312.pyc,, +xgrammar/kernels/__pycache__/apply_token_bitmask_inplace_torch_compile.cpython-312.pyc,, +xgrammar/kernels/__pycache__/apply_token_bitmask_inplace_triton.cpython-312.pyc,, +xgrammar/kernels/__pycache__/apply_token_bitmask_mlx.cpython-312.pyc,, +xgrammar/kernels/apply_token_bitmask_inplace_cpu.py,sha256=HOxSWkD4IZ3ikKDYKUPsaz-DL2LU8-69y_msX5tpf9k,1512 +xgrammar/kernels/apply_token_bitmask_inplace_cuda.cu,sha256=jEEl3KqavKcPDsRkf1AcPVwlaTUeVO8oOT7EX7rdQbU,10015 +xgrammar/kernels/apply_token_bitmask_inplace_cuda.py,sha256=m5bLvWDtXQG0t4FTY8akh_6NOD0hDTyp7ENN1Smrvk8,4049 +xgrammar/kernels/apply_token_bitmask_inplace_torch_compile.py,sha256=m2-wTBVvSyi-sebEgSKRfQIhg8_0jt5ZhgW-dPKWjs4,2138 +xgrammar/kernels/apply_token_bitmask_inplace_triton.py,sha256=a4STzJf7CKLZayqsz_Hqyo-zlqXkdk2Fmuqc6BZi73A,3810 +xgrammar/kernels/apply_token_bitmask_mlx.py,sha256=N1hbWiCyc0eOuK0R2Cj9LHWVi1DnOqZjWVP5TvQlgsU,883 +xgrammar/matcher.py,sha256=HrVDW_IDIl2OTV8N7TKED0Ez9uc4lTIEfs7jfBxP4Rg,13017 +xgrammar/support/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +xgrammar/support/__pycache__/__init__.cpython-312.pyc,, +xgrammar/support/__pycache__/logging.cpython-312.pyc,, +xgrammar/support/logging.py,sha256=8jmO9FAVVJpgbLOwclXpAowOe2vjdbtjk3Ryitphl9U,594 +xgrammar/testing.py,sha256=bOYvVYCfVBbAYWWmcc9SJF8zgy7UiEL72B7juaDxhDI,10867 +xgrammar/tokenizer_info.py,sha256=348J7MkyEBtvTUcyE4diby0xPM0U-kGeN9glu7M787Q,15734 +xgrammar/version.py,sha256=7VtgdhF7CEhyWBFHSaTl_LXmNFhWTvlSu3U4f222M_0,4296 +xgrammar/xgrammar_bindings.cpython-312-x86_64-linux-gnu.so,sha256=VQ7yxD7_cDDw6eoZ2AW8Nam6EbuyXn7JkAvDYMicBy8,17258920 diff --git a/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/REQUESTED b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/WHEEL b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..d7683c1fb81a50a37ee5efebca05f7cb2decd3a9 --- /dev/null +++ b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/WHEEL @@ -0,0 +1,6 @@ +Wheel-Version: 1.0 +Generator: scikit-build-core 0.11.1 +Root-Is-Purelib: false +Tag: cp312-cp312-manylinux_2_17_x86_64 +Tag: cp312-cp312-manylinux2014_x86_64 + diff --git a/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/licenses/LICENSE b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/licenses/LICENSE @@ -0,0 +1,201 @@ + 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. diff --git a/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/licenses/NOTICE b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/licenses/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..2d56df4754f7ebeba356393e69df5e2eef265232 --- /dev/null +++ b/lib/python3.12/site-packages/xgrammar-0.1.18.dist-info/licenses/NOTICE @@ -0,0 +1,3 @@ +XGrammar + +Copyright (c) 2024 by XGrammar Contributors