fix pylint

This commit is contained in:
Johan Hansson
2024-01-11 16:54:03 +01:00
parent 4fa557889c
commit cd2cda3cda
2 changed files with 4 additions and 44 deletions

View File

@@ -3,12 +3,11 @@ Test module for alpacha integration w chatml
"""
import pytest
from datasets import Dataset
from tokenizers import AddedToken
from transformers import AutoTokenizer
from axolotl.datasets import TokenizedPromptDataset
from axolotl.prompt_tokenizers import AlpacaPromptTokenizingStrategy
from axolotl.prompters import AlpacaPrompter
from utils import fixture_tokenizer
@pytest.fixture(name="alpacha_dataset")
@@ -24,25 +23,6 @@ def fixture_alpacha_dataset():
)
@pytest.fixture(name="tokenizer")
def fixture_tokenizer():
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
tokenizer.add_special_tokens(
{
"eos_token": AddedToken(
"<|im_end|>", rstrip=False, lstrip=False, normalized=False
)
}
)
tokenizer.add_tokens(
[
AddedToken("<|im_start|>", rstrip=False, lstrip=False, normalized=False),
]
)
return tokenizer
class TestAlpacha:
"""
Test class for alpacha prompter