google-tunix 0.0.2__tar.gz → 0.1.0.dev0__tar.gz

{google_tunix-0.0.2 → google_tunix-0.1.0.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-tunix
-Version: 0.0.2
+Version: 0.1.0.dev0
 Summary: A lightweight JAX-native LLM post-training framework.
 Author-email: Tunix Developers <tunix-dev@google.com>
 License-Expression: Apache-2.0
@@ -16,7 +16,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: datasets
-Requires-Dist: flax
+Requires-Dist: gcsfs
 Requires-Dist: grain
 Requires-Dist: huggingface_hub
 Requires-Dist: jax
@@ -30,6 +30,7 @@ Requires-Dist: tensorboardX
 Requires-Dist: tensorflow_datasets
 Requires-Dist: tqdm
 Requires-Dist: transformers
+Requires-Dist: python-dotenv
 Provides-Extra: docs
 Requires-Dist: sphinx>=8.2.3; extra == "docs"
 Requires-Dist: sphinx-book-theme>=1.1.4; extra == "docs"
@@ -40,6 +41,10 @@ Requires-Dist: matplotlib>=3.10.0; extra == "docs"
 Requires-Dist: sphinx-gallery>=0.19.0; extra == "docs"
 Requires-Dist: sphinx-collections>=0.0.1; extra == "docs"
 Requires-Dist: sphinx_contributors; extra == "docs"
+Provides-Extra: prod
+Requires-Dist: flax>=0.12.0; extra == "prod"
+Provides-Extra: dev
+Requires-Dist: flax>=0.11.2; extra == "dev"
 Dynamic: license-file
 
 # Tunix: A JAX-native LLM Post-Training Library
@@ -106,13 +111,30 @@ Tunix is still under development, here's a glimpse of the current features:
 
 ## Installation
 
-Tunix doesn't have a PyPI package yet. To use Tunix, you need to install from
-GitHub directly.
+You can install Tunix in several ways:
+
+1. From PyPI (recommended):
+
+```sh
+pip install "tunix[prod]"
+```
+
+2. Directly from GitHub (latest main branch)
 
 ```sh
 pip install git+https://github.com/google/tunix
 ```
 
+3. From source (editable install) If you plan to modify the codebase and run it
+   in development mode:
+
+```sh
+git clone https://github.com/google/tunix.git
+cd tunix
+pip install -e ".[dev]"
+
+```
+
 ## Getting Started
 
 To get started, we have a bunch of detailed examples and tutorials.
@@ -155,3 +177,9 @@ cutting-edge research and easy reproducibility.
 Thank you for your interest in Tunix. We're working hard to bring you a powerful
 and efficient library for LLM post-training. Please follow our progress and
 check back for updates!
+
+## Acknowledgements
+
+Thank you to all our wonderful contributors!
+
+[![Contributors](https://contrib.rocks/image?repo=google/tunix)](https://github.com/google/tunix/graphs/contributors)

{google_tunix-0.0.2 → google_tunix-0.1.0.dev0}/README.md

@@ -62,13 +62,30 @@ Tunix is still under development, here's a glimpse of the current features:
 
 ## Installation
 
-Tunix doesn't have a PyPI package yet. To use Tunix, you need to install from
-GitHub directly.
+You can install Tunix in several ways:
+
+1. From PyPI (recommended):
+
+```sh
+pip install "tunix[prod]"
+```
+
+2. Directly from GitHub (latest main branch)
 
 ```sh
 pip install git+https://github.com/google/tunix
 ```
 
+3. From source (editable install) If you plan to modify the codebase and run it
+   in development mode:
+
+```sh
+git clone https://github.com/google/tunix.git
+cd tunix
+pip install -e ".[dev]"
+
+```
+
 ## Getting Started
 
 To get started, we have a bunch of detailed examples and tutorials.
@@ -111,3 +128,9 @@ cutting-edge research and easy reproducibility.
 Thank you for your interest in Tunix. We're working hard to bring you a powerful
 and efficient library for LLM post-training. Please follow our progress and
 check back for updates!
+
+## Acknowledgements
+
+Thank you to all our wonderful contributors!
+
+[![Contributors](https://contrib.rocks/image?repo=google/tunix)](https://github.com/google/tunix/graphs/contributors)

{google_tunix-0.0.2 → google_tunix-0.1.0.dev0}/google_tunix.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-tunix
-Version: 0.0.2
+Version: 0.1.0.dev0
 Summary: A lightweight JAX-native LLM post-training framework.
 Author-email: Tunix Developers <tunix-dev@google.com>
 License-Expression: Apache-2.0
@@ -16,7 +16,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: datasets
-Requires-Dist: flax
+Requires-Dist: gcsfs
 Requires-Dist: grain
 Requires-Dist: huggingface_hub
 Requires-Dist: jax
@@ -30,6 +30,7 @@ Requires-Dist: tensorboardX
 Requires-Dist: tensorflow_datasets
 Requires-Dist: tqdm
 Requires-Dist: transformers
+Requires-Dist: python-dotenv
 Provides-Extra: docs
 Requires-Dist: sphinx>=8.2.3; extra == "docs"
 Requires-Dist: sphinx-book-theme>=1.1.4; extra == "docs"
@@ -40,6 +41,10 @@ Requires-Dist: matplotlib>=3.10.0; extra == "docs"
 Requires-Dist: sphinx-gallery>=0.19.0; extra == "docs"
 Requires-Dist: sphinx-collections>=0.0.1; extra == "docs"
 Requires-Dist: sphinx_contributors; extra == "docs"
+Provides-Extra: prod
+Requires-Dist: flax>=0.12.0; extra == "prod"
+Provides-Extra: dev
+Requires-Dist: flax>=0.11.2; extra == "dev"
 Dynamic: license-file
 
 # Tunix: A JAX-native LLM Post-Training Library
@@ -106,13 +111,30 @@ Tunix is still under development, here's a glimpse of the current features:
 
 ## Installation
 
-Tunix doesn't have a PyPI package yet. To use Tunix, you need to install from
-GitHub directly.
+You can install Tunix in several ways:
+
+1. From PyPI (recommended):
+
+```sh
+pip install "tunix[prod]"
+```
+
+2. Directly from GitHub (latest main branch)
 
 ```sh
 pip install git+https://github.com/google/tunix
 ```
 
+3. From source (editable install) If you plan to modify the codebase and run it
+   in development mode:
+
+```sh
+git clone https://github.com/google/tunix.git
+cd tunix
+pip install -e ".[dev]"
+
+```
+
 ## Getting Started
 
 To get started, we have a bunch of detailed examples and tutorials.
@@ -155,3 +177,9 @@ cutting-edge research and easy reproducibility.
 Thank you for your interest in Tunix. We're working hard to bring you a powerful
 and efficient library for LLM post-training. Please follow our progress and
 check back for updates!
+
+## Acknowledgements
+
+Thank you to all our wonderful contributors!
+
+[![Contributors](https://contrib.rocks/image?repo=google/tunix)](https://github.com/google/tunix/graphs/contributors)

{google_tunix-0.0.2 → google_tunix-0.1.0.dev0}/google_tunix.egg-info/SOURCES.txt

@@ -7,6 +7,11 @@ google_tunix.egg-info/dependency_links.txt
 google_tunix.egg-info/requires.txt
 google_tunix.egg-info/top_level.txt
 tunix/__init__.py
+tunix/cli/config.py
+tunix/cli/grpo_main.py
+tunix/cli/peft_main.py
+tunix/cli/utils/model.py
+tunix/cli/utils/reward.py
 tunix/distillation/__init__.py
 tunix/distillation/distillation_trainer.py
 tunix/distillation/feature_extraction/__init__.py
@@ -19,6 +24,7 @@ tunix/distillation/strategies/base_strategy.py
 tunix/distillation/strategies/feature_pooling.py
 tunix/distillation/strategies/feature_projection.py
 tunix/distillation/strategies/logit.py
+tunix/examples/data/math_dataset.py
 tunix/examples/data/translation_dataset.py
 tunix/generate/base_sampler.py
 tunix/generate/beam_search.py
@@ -27,7 +33,7 @@ tunix/generate/tokenizer_adapter.py
 tunix/generate/utils.py
 tunix/generate/vllm_sampler.py
 tunix/models/safetensors_loader.py
-tunix/models/gemma/gemma.py
+tunix/models/gemma/model.py
 tunix/models/gemma/params.py
 tunix/models/gemma/params_safetensors.py
 tunix/models/gemma/sampler.py
@@ -73,11 +79,9 @@ tunix/rl/rollout/base_rollout.py
 tunix/rl/rollout/vanilla_rollout.py
 tunix/rl/rollout/vllm_rollout.py
 tunix/sft/checkpoint_manager.py
-tunix/sft/config.py
 tunix/sft/hooks.py
 tunix/sft/inflight_throttler.py
 tunix/sft/metrics_logger.py
-tunix/sft/peft_main.py
 tunix/sft/peft_trainer.py
 tunix/sft/profiler.py
 tunix/sft/progress_bar.py

{google_tunix-0.0.2 → google_tunix-0.1.0.dev0}/google_tunix.egg-info/requires.txt

@@ -1,5 +1,5 @@
 datasets
-flax
+gcsfs
 grain
 huggingface_hub
 jax
@@ -13,6 +13,10 @@ tensorboardX
 tensorflow_datasets
 tqdm
 transformers
+python-dotenv
+
+[dev]
+flax>=0.11.2
 
 [docs]
 sphinx>=8.2.3
@@ -24,3 +28,6 @@ matplotlib>=3.10.0
 sphinx-gallery>=0.19.0
 sphinx-collections>=0.0.1
 sphinx_contributors
+
+[prod]
+flax>=0.12.0

{google_tunix-0.0.2 → google_tunix-0.1.0.dev0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "google-tunix"
-version = "0.0.2"
+version = "0.1.0.dev0"
 authors = [
   { name = "Tunix Developers", email = "tunix-dev@google.com" },
 ]
@@ -19,7 +19,7 @@ classifiers = [
 ]
 dependencies = [
   "datasets",
-  "flax",
+  "gcsfs",
   "grain",
   "huggingface_hub",
   "jax",
@@ -33,6 +33,7 @@ dependencies = [
   "tensorflow_datasets",
   "tqdm",
   "transformers",
+  "python-dotenv",
 ]
 
 [project.optional-dependencies]
@@ -47,6 +48,12 @@ docs = [
     "sphinx-collections>=0.0.1",
     "sphinx_contributors",
 ]
+prod = [
+    "flax>=0.12.0",
+]
+dev = [
+    "flax>=0.11.2",
+]
 
 [project.urls]
 Source = "https://github.com/google/tunix"

{google_tunix-0.0.2 → google_tunix-0.1.0.dev0}/tunix/__init__.py

@@ -14,21 +14,39 @@
 
 """Tunix API."""
 
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version
+
+try:
+  __version__ = version("google-tunix")  # match the name in pyproject.toml
+except PackageNotFoundError:
+  __version__ = "0.0.0.dev0"  # fallback for editable installs
+
+
 # pylint: disable=g-multiple-import, g-importing-member
 
 from tunix.distillation.distillation_trainer import DistillationTrainer
 from tunix.distillation.distillation_trainer import TrainingConfig as DistillationTrainingConfig
 from tunix.generate.sampler import CacheConfig
 from tunix.generate.sampler import Sampler
+from tunix.rl.grpo.grpo_learner import GRPOConfig
 from tunix.rl.grpo.grpo_learner import GrpoConfig
+from tunix.rl.grpo.grpo_learner import GRPOLearner
 from tunix.rl.grpo.grpo_learner import GrpoLearner
 from tunix.rl.grpo.grpo_learner import RewardFn
+from tunix.rl.ppo.ppo_learner import PPOConfig
+from tunix.rl.ppo.ppo_learner import PpoConfig
+from tunix.rl.ppo.ppo_learner import PPOLearner
+from tunix.rl.ppo.ppo_learner import PpoLearner
 from tunix.rl.rl_cluster import ClusterConfig
+from tunix.rl.rl_cluster import MetricsBuffer
 from tunix.rl.rl_cluster import RLCluster
 from tunix.rl.rl_cluster import RLTrainingConfig
 from tunix.rl.rl_cluster import Role
 from tunix.rl.rollout.base_rollout import RolloutConfig
+from tunix.sft.dpo.dpo_trainer import DPOTrainer
 from tunix.sft.dpo.dpo_trainer import DpoTrainer
+from tunix.sft.dpo.dpo_trainer import DPOTrainingConfig
 from tunix.sft.dpo.dpo_trainer import DpoTrainingConfig
 from tunix.sft.metrics_logger import MetricsLogger
 from tunix.sft.metrics_logger import MetricsLoggerOptions