Adds github actions and dummy test

.github/workflows/tests.yml

@@ -0,0 +1,29 @@
+name: tests
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+    branches: [ master ]
+
+jobs:
+  tests:
+    strategy:
+      matrix:
+        python-version: [3.7, 3.8, 3.9]
+        os: ['ubuntu-latest', 'macos-latest', 'windows-latest']
+    runs-on: ${{ matrix.os }}
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        # cpu version of pytorch
+        pip install .[test]
+    - name: Test with pytest
+      run: |
+        make test

Makefile

@@ -14,7 +14,11 @@ docs: $(SRC)
 	touch docs
 
 test:
-	nbdev_test_nbs
+	pytest tests
+
+format:
+	black --line-length 119 --target-version py36 tests trl examples 
+	isort tests trl examples 
 
 release: pypi
 	nbdev_bump_version

settings.ini

@@ -14,7 +14,7 @@ language = English
 custom_sidebar = False
 license = apache2
 status = 2
-requirements = torch>=1.4.0 transformers==4.3.2 numpy>=1.18.2
+requirements = torch>=1.4.0 transformers>=4.18.0 numpy>=1.18.2
 nbs_path = ./nbs/
 doc_path = docs
 doc_host = https://lvwerra.github.io

setup.py

@@ -21,6 +21,10 @@
 py_versions = '2.0 2.1 2.2 2.3 2.4 2.5 2.6 2.7 3.0 3.1 3.2 3.3 3.4 3.5 3.6 3.7 3.8'.split()
 
 requirements = cfg.get('requirements','').split()
+extras = {
+    "test" : ["pytest","pytest-xdist",],
+    "dev" : ["pytest","pytest-xdist", "black", "isort", "flake8>=3.8.3"],
+}
 lic = licenses[cfg['license']]
 min_python = cfg['min_python']
 
@@ -37,6 +41,7 @@
     packages = setuptools.find_packages(),
     include_package_data = True,
     install_requires = requirements,
+    extras_require = extras,
     python_requires  = '>=' + cfg['min_python'],
     long_description = open('README.md').read(),
     long_description_content_type = 'text/markdown',

tests/test_gpt2_model.py

@@ -0,0 +1,69 @@
+# imports
+import pytest
+import torch
+from transformers import GPT2Tokenizer
+from trl.gpt2 import GPT2HeadWithValueModel, respond_to_batch
+from trl.ppo import PPOTrainer
+
+
+def test_gpt2_model():
+    # get models
+    gpt2_model = GPT2HeadWithValueModel.from_pretrained("gpt2")
+    gpt2_model_ref = GPT2HeadWithValueModel.from_pretrained("gpt2")
+    gpt2_tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+
+    # initialize trainer
+    ppo_config = {"batch_size": 1, "forward_batch_size": 1}
+    ppo_trainer = PPOTrainer(gpt2_model, gpt2_model_ref, gpt2_tokenizer, **ppo_config)
+
+    # encode a query
+    query_txt = "This morning I went to the "
+    query_tensor = gpt2_tokenizer.encode(query_txt, return_tensors="pt")
+    assert query_tensor.shape == (1, 7)
+    # get model response
+    response_tensor = respond_to_batch(gpt2_model, query_tensor)
+    assert response_tensor.shape == (1, 20)
+    response_txt = gpt2_tokenizer.decode(response_tensor[0, :])
+
+    # define a reward for response
+    # (this could be any reward such as human feedback or output from another model)
+    reward = [torch.tensor(1.0)]
+
+    # train model with ppo
+    train_stats = ppo_trainer.step([query_tensor[0]], [response_tensor[0]], reward)
+
+    EXPECTED_STATS = [
+        "objective/kl",
+        "objective/kl_dist",
+        "objective/logprobs",
+        "objective/ref_logprobs",
+        "objective/kl_coef",
+        "objective/entropy",
+        "ppo/mean_non_score_reward",
+        "ppo/loss/policy",
+        "ppo/loss/value",
+        "ppo/loss/total",
+        "ppo/policy/entropy",
+        "ppo/policy/approxkl",
+        "ppo/policy/policykl",
+        "ppo/policy/clipfrac",
+        "ppo/policy/advantages",
+        "ppo/policy/advantages_mean",
+        "ppo/policy/ratio",
+        "ppo/returns/mean",
+        "ppo/returns/var",
+        "ppo/val/vpred",
+        "ppo/val/error",
+        "ppo/val/clipfrac",
+        "ppo/val/mean",
+        "ppo/val/var",
+        "ppo/val/var_explained",
+        "time/ppo/forward_pass",
+        "time/ppo/compute_rewards",
+        "time/ppo/optimize_step",
+        "time/ppo/calc_stats",
+        "time/ppo/total",
+    ]
+
+    for stat in EXPECTED_STATS:
+        assert stat in train_stats.keys()