From 7eb33a77dde99a5434eab2c9ddd4c66acd3cee05 Mon Sep 17 00:00:00 2001 From: NanoCode012 Date: Mon, 29 May 2023 14:02:43 +0900 Subject: [PATCH] Lint test_prompters --- tests/test_prompters.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_prompters.py b/tests/test_prompters.py index b4a34c6c0..11610ccc5 100644 --- a/tests/test_prompters.py +++ b/tests/test_prompters.py @@ -1,9 +1,15 @@ +"""Module testing prompters""" + import unittest from axolotl.prompters import AlpacaPrompter, PromptStyle class AlpacaPrompterTest(unittest.TestCase): + """ + Test AlpacaPrompter + """ + def test_prompt_style_w_none(self): prompter = AlpacaPrompter(prompt_style=None) res = next(prompter.build_prompt("tell me a joke")) @@ -11,7 +17,7 @@ class AlpacaPrompterTest(unittest.TestCase): assert "### Instruction:" in res def test_prompt_style_w_instruct(self): - prompter = AlpacaPrompter(prompt_style=PromptStyle.instruct.value) + prompter = AlpacaPrompter(prompt_style=PromptStyle.INSTRUCT.value) res = next( prompter.build_prompt("tell me a joke about the following", "alpacas") ) @@ -31,7 +37,7 @@ class AlpacaPrompterTest(unittest.TestCase): assert "ASSISTANT:" not in res def test_prompt_style_w_chat(self): - prompter = AlpacaPrompter(prompt_style=PromptStyle.chat.value) + prompter = AlpacaPrompter(prompt_style=PromptStyle.CHAT.value) res = next( prompter.build_prompt("tell me a joke about the following", "alpacas") )