138 lines
4.3 KiB
Python
138 lines
4.3 KiB
Python
"""Tests for strategy module."""
|
||
|
||
import pytest
|
||
from pipeline.strategies import (
|
||
PromptStrategy,
|
||
ChatMLStrategy,
|
||
AlpacaStrategy,
|
||
StrategyFactory,
|
||
)
|
||
|
||
|
||
class DummyTokenizer:
|
||
def __init__(self):
|
||
self._special_token_map = {}
|
||
self._chat_template = None
|
||
|
||
def encode(self, text: str, add_special_tokens: bool = False):
|
||
return [ord(c) for c in text]
|
||
|
||
def decode(self, tokens, skip_special_tokens=True):
|
||
return "".join(chr(t) for t in tokens)
|
||
|
||
def token_to_id(self, token: str):
|
||
return ord(token)
|
||
|
||
def set_chat_template(self, template):
|
||
self._chat_template = template
|
||
|
||
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=True):
|
||
text = ""
|
||
for m in messages:
|
||
text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n"
|
||
if add_generation_prompt:
|
||
text += "<｜im▁start｜>assistant\n"
|
||
return self.encode(text) if tokenize else text
|
||
|
||
|
||
class DummyStrategy(PromptStrategy):
|
||
def __init__(self, tokenizer):
|
||
super().__init__(tokenizer)
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
return "dummy"
|
||
|
||
def assemble_prompt(self, query_tokens):
|
||
prefix = self._encode_format("Q: ")
|
||
return prefix + query_tokens
|
||
|
||
def assemble_response(self, response_tokens):
|
||
suffix = self._encode_format("<｜end▁of▁sentence｜>")
|
||
return response_tokens + suffix
|
||
|
||
|
||
def _decode(tokens):
|
||
return "".join(chr(t) for t in tokens)
|
||
|
||
|
||
class TestChatMLStrategy:
|
||
def test_name(self):
|
||
assert ChatMLStrategy(DummyTokenizer()).name == "chatml"
|
||
|
||
def test_assemble_prompt(self):
|
||
tk = DummyTokenizer()
|
||
strategy = ChatMLStrategy(tk)
|
||
query_tokens = tk.encode("hello")
|
||
prompt = strategy.assemble_prompt(query_tokens)
|
||
text = _decode(prompt)
|
||
assert "<｜im▁start｜>user" in text
|
||
assert "hello" in text
|
||
assert "<｜im▁start｜>assistant" in text
|
||
|
||
def test_assemble_response(self):
|
||
tk = DummyTokenizer()
|
||
strategy = ChatMLStrategy(tk)
|
||
response_tokens = tk.encode("world")
|
||
response = strategy.assemble_response(response_tokens)
|
||
text = _decode(response)
|
||
assert "world" in text
|
||
assert "<｜im▁end｜>" in text
|
||
|
||
def test_prompt_ends_with_assistant_start(self):
|
||
tk = DummyTokenizer()
|
||
strategy = ChatMLStrategy(tk)
|
||
prompt = strategy.assemble_prompt(tk.encode("hi"))
|
||
assistant_start = tk.encode("<｜im▁start｜>assistant\n")
|
||
assert prompt[-len(assistant_start):] == assistant_start
|
||
|
||
|
||
class TestAlpacaStrategy:
|
||
def test_name(self):
|
||
assert AlpacaStrategy(DummyTokenizer()).name == "alpaca"
|
||
|
||
def test_assemble_prompt(self):
|
||
tk = DummyTokenizer()
|
||
strategy = AlpacaStrategy(tk)
|
||
query_tokens = tk.encode("hello")
|
||
prompt = strategy.assemble_prompt(query_tokens)
|
||
text = _decode(prompt)
|
||
assert "### Instruction:" in text
|
||
assert "hello" in text
|
||
assert "### Response:" in text
|
||
|
||
def test_assemble_response(self):
|
||
tk = DummyTokenizer()
|
||
strategy = AlpacaStrategy(tk)
|
||
response_tokens = tk.encode("world")
|
||
response = strategy.assemble_response(response_tokens)
|
||
text = _decode(response)
|
||
assert "world" in text
|
||
assert "<｜end▁of▁sentence｜>" in text
|
||
|
||
|
||
class TestStrategyFactory:
|
||
def test_create_chatml(self):
|
||
tk = DummyTokenizer()
|
||
assert isinstance(StrategyFactory.create("chatml", tk), ChatMLStrategy)
|
||
|
||
def test_create_alpaca(self):
|
||
tk = DummyTokenizer()
|
||
assert isinstance(StrategyFactory.create("alpaca", tk), AlpacaStrategy)
|
||
|
||
def test_create_invalid_raises_error(self):
|
||
with pytest.raises(ValueError, match="Unknown strategy"):
|
||
StrategyFactory.create("invalid_strategy", DummyTokenizer())
|
||
|
||
def test_register_and_create(self):
|
||
StrategyFactory.register("dummy")(DummyStrategy)
|
||
tk = DummyTokenizer()
|
||
strategy = StrategyFactory.create("dummy", tk)
|
||
assert isinstance(strategy, DummyStrategy)
|
||
assert strategy.name == "dummy"
|
||
|
||
def test_available_strategies(self):
|
||
strategies = StrategyFactory.available_types()
|
||
assert "chatml" in strategies
|
||
assert "alpaca" in strategies
|