63 lines
2.1 KiB
Python
63 lines
2.1 KiB
Python
"""
|
|
tests for chat_template prompt strategy
|
|
"""
|
|
|
|
import datasets
|
|
import pytest
|
|
from datasets import Dataset
|
|
from transformers import AutoTokenizer
|
|
|
|
from axolotl.datasets import TokenizedPromptDataset
|
|
from axolotl.prompt_strategies.stepwise_supervised import (
|
|
StepwiseSupervisedPromptTokenizingStrategy,
|
|
)
|
|
|
|
|
|
class TestStepWiseSupervisedPromptTokenizingStrategy:
|
|
"""
|
|
Test class for stepwise supervised prompt strategy
|
|
"""
|
|
|
|
@pytest.fixture()
|
|
def stepwise_supervised_dataset(self):
|
|
return Dataset.from_list(
|
|
[
|
|
{
|
|
"prompt": "Which number is larger, 9.8 or 9.11?",
|
|
"completions": [
|
|
"The fractional part of 9.8 is 0.8, while the fractional part of 9.11 is 0.11.",
|
|
"Since 0.11 is greater than 0.8, the number 9.11 is larger than 9.8.",
|
|
"Actually, this is incorrect. In decimal numbers, 0.8 is equal to 0.80, which is larger than 0.11. Therefore, 9.8 is larger than 9.11.",
|
|
],
|
|
"labels": [True, False, False],
|
|
}
|
|
]
|
|
)
|
|
|
|
@pytest.fixture()
|
|
def tokenizer(self):
|
|
return AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
|
|
|
|
def test_stepwise_supervised_dataset(self, tokenizer, stepwise_supervised_dataset):
|
|
strategy = StepwiseSupervisedPromptTokenizingStrategy(
|
|
tokenizer,
|
|
sequence_len=2048,
|
|
step_separator="\n",
|
|
)
|
|
stepwise_supervised_dataset = stepwise_supervised_dataset.cast_column(
|
|
"labels", datasets.Sequence(datasets.Value("int64"))
|
|
)
|
|
dataset_wrapper = TokenizedPromptDataset(
|
|
strategy,
|
|
stepwise_supervised_dataset,
|
|
process_count=1,
|
|
)
|
|
labels = dataset_wrapper[0]["labels"]
|
|
# expected labels is:
|
|
# the prompt + first step are ignored, followed by the label for step 1 (True)
|
|
# the second step, and its label (False)
|
|
# the third step, and its label (False)
|
|
expected = [-100] * 47 + [1] + [-100] * 29 + [0] + [-100] * 48 + [0]
|
|
|
|
assert labels == expected
|