From d199d6c261cabaab32bdcfbefa04671278b53698 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Sat, 27 May 2023 11:51:01 -0400
Subject: [PATCH] automated testing in github actions

---
 .github/workflows/tests.yml | 28 +++++++++++++++++++++
 requirements-tests.txt      | 11 +++++++++
 tests/test_prompters.py     | 49 +++++++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+)
 create mode 100644 .github/workflows/tests.yml
 create mode 100644 requirements-tests.txt
 create mode 100644 tests/test_prompters.py

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
new file mode 100644
index 000000000..3543c0b65
--- /dev/null
+++ b/.github/workflows/tests.yml
@@ -0,0 +1,28 @@
+name: PyTest
+on: push
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+
+    steps:
+      - name: Check out repository code
+        uses: actions/checkout@v2
+
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.9"
+
+      - name: Install pytest
+        run: |
+          pip install --upgrade pytest
+
+      - name: Install dependencies
+        run: |
+          pip install -r requirements-tests.txt
+
+      - name: Run tests
+        run: |
+          pytest tests/
diff --git a/requirements-tests.txt b/requirements-tests.txt
new file mode 100644
index 000000000..e58c194df
--- /dev/null
+++ b/requirements-tests.txt
@@ -0,0 +1,11 @@
+peft @ git+https://github.com/huggingface/peft.git
+transformers @ git+https://github.com/huggingface/transformers.git
+bitsandbytes>=0.39.0
+attrdict
+fire
+PyYAML==6.0
+black
+datasets
+accelerate>=0.19.0
+sentencepiece
+wandb
diff --git a/tests/test_prompters.py b/tests/test_prompters.py
new file mode 100644
index 000000000..1c3c13852
--- /dev/null
+++ b/tests/test_prompters.py
@@ -0,0 +1,49 @@
+import unittest
+
+from axolotl.prompters import AlpacaPrompter, PromptStyle
+
+
+class AlpacaPrompterTest(unittest.TestCase):
+    def test_prompt_style_w_none(self):
+        prompter = AlpacaPrompter(prompt_style=None)
+        res = next(prompter.build_prompt("tell me a joke"))
+        # just testing that it uses instruct style
+        assert "### Instruction:" in res
+
+    def test_prompt_style_w_instruct(self):
+        prompter = AlpacaPrompter(prompt_style=PromptStyle.instruct.value)
+        res = next(prompter.build_prompt("tell me a joke about the following", "alpacas"))
+        assert "Below is an instruction" in res
+        assert "### Instruction:" in res
+        assert "### Input:" in res
+        assert "alpacas" in res
+        assert "### Response:" in res
+        assert "USER:" not in res
+        assert "ASSISTANT:" not in res
+        res = next(prompter.build_prompt("tell me a joke about the following"))
+        assert "Below is an instruction" in res
+        assert "### Instruction:" in res
+        assert "### Input:" not in res
+        assert "### Response:" in res
+        assert "USER:" not in res
+        assert "ASSISTANT:" not in res
+
+    def test_prompt_style_w_chat(self):
+        prompter = AlpacaPrompter(prompt_style=PromptStyle.chat.value)
+        res = next(prompter.build_prompt("tell me a joke about the following", "alpacas"))
+        assert "Below is an instruction" in res
+        assert "### Instruction:" not in res
+        assert "### Input:" not in res
+        assert "alpacas" in res
+        assert "### Response:" not in res
+        assert "USER:" in res
+        assert "ASSISTANT:" in res
+        res = next(prompter.build_prompt("tell me a joke about the following"))
+        assert "Below is an instruction" in res
+        assert "### Instruction:" not in res
+        assert "### Input:" not in res
+        assert "### Response:" not in res
+        assert "USER:" in res
+        assert "ASSISTANT:" in res
+
+