diff --git a/pkgs/development/python-modules/trl/default.nix b/pkgs/development/python-modules/trl/default.nix new file mode 100644 index 000000000000..1c75339a5fae --- /dev/null +++ b/pkgs/development/python-modules/trl/default.nix @@ -0,0 +1,53 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + + # build-system + setuptools, + setuptools-scm, + + # dependencies + accelerate, + datasets, + rich, + transformers, +}: + +buildPythonPackage rec { + pname = "trl"; + version = "0.15.2"; + pyproject = true; + + src = fetchFromGitHub { + owner = "huggingface"; + repo = "trl"; + tag = "v${version}"; + hash = "sha256-HsSmFXFqDOWVLa6VXdPZVS9C3bjYcsliR0TwNpPiQx4="; + }; + + build-system = [ + setuptools + setuptools-scm + ]; + + dependencies = [ + accelerate + datasets + rich + transformers + ]; + + # Many tests require internet access. + doCheck = false; + + pythonImportsCheck = [ "trl" ]; + + meta = { + description = "Train transformer language models with reinforcement learning"; + homepage = "https://github.com/huggingface/trl"; + changelog = "https://github.com/huggingface/trl/releases/tag/v${version}"; + license = lib.licenses.asl20; + maintainers = with lib.maintainers; [ hoh ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index dfce50496beb..7bb745602db0 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -17667,6 +17667,8 @@ self: super: with self; { tritonclient = callPackage ../development/python-modules/tritonclient { }; + trl = callPackage ../development/python-modules/trl { }; + trlib = toPythonModule ( pkgs.trlib.override { pythonSupport = true;