| | from unittest import mock |
| | from unittest.mock import MagicMock, patch |
| |
|
| | import numpy as np |
| | import pandas as pd |
| | import pytest |
| |
|
| | from llm_studio.app_utils.default_datasets import ( |
| | prepare_default_dataset_causal_language_modeling, |
| | ) |
| | from llm_studio.python_configs.text_causal_language_modeling_config import ( |
| | ConfigNLPCausalLMDataset, |
| | ConfigNLPCausalLMTokenizer, |
| | ConfigProblemBase, |
| | ) |
| | from llm_studio.src.datasets.text_causal_language_modeling_ds import CustomDataset |
| |
|
| |
|
| | def test_prepare_default_dataset(tmp_path): |
| | df = prepare_default_dataset_causal_language_modeling(tmp_path) |
| | assert isinstance(df, pd.DataFrame) |
| | assert set(df.keys()) == set( |
| | ["instruction", "output", "id", "parent_id", "lang", "rank"] |
| | ) |
| | assert df.shape == (13026, 6) |
| |
|
| |
|
| | def test_clean_output(): |
| | output = { |
| | "predicted_text": np.array( |
| | [ |
| | "This is a test", |
| | "This is a test <stop> This is a test", |
| | "This is a test <stop2> This is a test", |
| | "This is a test <stop3> <stop> This is a test", |
| | "<stop2> <stop> This is a test", |
| | "This is a test <stop>", |
| | ] |
| | ) |
| | } |
| |
|
| | cfg = mock.MagicMock() |
| | cfg.tokenizer._stop_words = ["<stop>", "<stop2>", "<stop3>"] |
| |
|
| | predicted_text_clean = CustomDataset.clean_output(output=output, cfg=cfg)[ |
| | "predicted_text" |
| | ] |
| | assert predicted_text_clean == [ |
| | "This is a test", |
| | "This is a test", |
| | "This is a test", |
| | "This is a test", |
| | "", |
| | "This is a test", |
| | ] |
| |
|
| |
|
| | def test_sanity_check_raises_error(): |
| | mock_config = MagicMock() |
| | mock_config.dataset.parent_id_column = "parent_id" |
| | mock_config.dataset.answer_column = "answer" |
| |
|
| | df_1 = pd.DataFrame( |
| | { |
| | "id": [1, 2, 3, 4], |
| | "parent_id": [2, None, 4, 1], |
| | "answer": ["a", "b", "c", "d"], |
| | "other_data": ["a", "b", "c", "d"], |
| | } |
| | ) |
| | CustomDataset.sanity_check(df_1, mock_config) |
| |
|
| | df_2 = pd.DataFrame( |
| | { |
| | "id": [1, 2, 3, 4], |
| | "parent_id": [None, None, None, None], |
| | "answer": ["a", "b", "c", "d"], |
| | "other_data": ["a", "b", "c", "d"], |
| | } |
| | ) |
| | CustomDataset.sanity_check(df_2, mock_config) |
| |
|
| | invalid_df_1 = pd.DataFrame( |
| | { |
| | "id": [1, 2, 3, 4], |
| | "parent_id": [1, 2, 3, 4], |
| | "answer": ["a", "b", "c", "d"], |
| | "other_data": ["a", "b", "c", "d"], |
| | } |
| | ) |
| | with pytest.raises( |
| | AssertionError, match="Parent id column is the same as id column for some rows" |
| | ): |
| | CustomDataset.sanity_check(invalid_df_1, mock_config) |
| |
|
| | invalid_df_2 = pd.DataFrame( |
| | { |
| | "id": [1, 2, 3, 4], |
| | "parent_id": [2, 3, 4, 1], |
| | "other_data": ["a", "b", "c", "d"], |
| | } |
| | ) |
| | with pytest.raises( |
| | AssertionError, |
| | match="Did not find any conversation start. " |
| | "Please ensure that some parent ids are empty.", |
| | ): |
| | CustomDataset.sanity_check(invalid_df_2, mock_config) |
| |
|
| |
|
| | @pytest.fixture |
| | def mock_auto_tokenizer(): |
| | |
| | |
| | with patch( |
| | "transformers.AutoTokenizer.from_pretrained", autospec=True |
| | ) as mock_from_pretrained: |
| | yield mock_from_pretrained |
| |
|
| |
|
| | def test_init(mock_auto_tokenizer): |
| | df = pd.DataFrame( |
| | { |
| | "col_A": [1, 2, 3], |
| | "col_B": [4, 5, 6], |
| | } |
| | ) |
| | cfg = mock.MagicMock() |
| | cfg.dataset.prompt_column = "col_A" |
| | cfg.dataset.answer_column = "col_B" |
| | cfg.dataset.parent_id_column = "None" |
| | cfg.dataset.system_column = "None" |
| |
|
| | cfg.dataset.text_system_start = "" |
| | cfg.dataset.text_prompt_start = "" |
| | cfg.dataset.text_answer_separator = "" |
| |
|
| | dataset = CustomDataset(df, cfg) |
| |
|
| | assert dataset.df.equals(df) |
| | assert dataset.mode == "train" |
| |
|
| |
|
| | def test_getitem(): |
| | df = pd.DataFrame( |
| | { |
| | "prompt": ["prompt 1", "prompt 2", "prompt 3"], |
| | "answer": ["answer 1", "answer 2", "answer 3"], |
| | "parent_id": [None, 0, 1], |
| | "system": ["system 1", "system 2", "system 3"], |
| | "id": [0, 1, 2], |
| | } |
| | ) |
| |
|
| | cfg = ConfigProblemBase( |
| | dataset=ConfigNLPCausalLMDataset( |
| | prompt_column=("prompt",), |
| | answer_column="answer", |
| | parent_id_column="parent_id", |
| | system_column="system", |
| | text_system_start="System:", |
| | text_prompt_start="Prompt:", |
| | text_answer_separator="Answer:", |
| | add_eos_token_to_answer=True, |
| | limit_chained_samples=True, |
| | ), |
| | tokenizer=ConfigNLPCausalLMTokenizer(max_length=513), |
| | ) |
| |
|
| | cfg.llm_backbone = "EleutherAI/pythia-2.8b-deduped" |
| |
|
| | dataset = CustomDataset(df, cfg) |
| | assert len(dataset) == 1 |
| |
|
| | result = dataset[0] |
| | assert isinstance(result, dict) |
| | assert set(result.keys()) == { |
| | "labels", |
| | "input_ids", |
| | "attention_mask", |
| | "prompt_input_ids", |
| | "prompt_attention_mask", |
| | "answer_input_ids", |
| | "answer_attention_mask", |
| | } |
| |
|
| | assert ( |
| | dataset.tokenizer.decode(result["input_ids"], skip_special_tokens=True) |
| | == "System:system 1" |
| | "Prompt:prompt 1" |
| | "Answer:answer 1" |
| | "Prompt:prompt 2" |
| | "Answer:answer 2" |
| | "Prompt:prompt 3" |
| | "Answer:answer 3" |
| | ) |
| |
|
| | assert ( |
| | dataset.tokenizer.decode(result["prompt_input_ids"], skip_special_tokens=True) |
| | == "System:system 1" |
| | "Prompt:prompt 1" |
| | "Answer:answer 1" |
| | "Prompt:prompt 2" |
| | "Answer:answer 2" |
| | "Prompt:prompt 3" |
| | "Answer:" |
| | ) |
| |
|
| | assert ( |
| | dataset.tokenizer.decode(result["input_ids"], skip_special_tokens=False) |
| | == "<|endoftext|>" * 475 + "System:system 1" |
| | "<|endoftext|>" |
| | "Prompt:prompt 1" |
| | "<|endoftext|>" |
| | "Answer:answer 1" |
| | "<|endoftext|>" |
| | "Prompt:prompt 2" |
| | "<|endoftext|>" |
| | "Answer:answer 2" |
| | "<|endoftext|>" |
| | "Prompt:prompt 3" |
| | "<|endoftext|>" |
| | "Answer:answer 3" |
| | "<|endoftext|>" |
| | ) |
| |
|
| | assert result["input_ids"].shape == (513,) |
| | assert result["prompt_input_ids"].shape == (513,) |
| |
|
| |
|
| | def test_getitem_no_chaining(): |
| | df = pd.DataFrame( |
| | { |
| | "prompt": ["prompt 1", "prompt 2", "prompt 3"], |
| | "answer": ["answer 1", "answer 2", "answer 3"], |
| | "parent_id": [None, 0, 1], |
| | "system": ["system 1", "system 2", "system 3"], |
| | "id": [0, 1, 2], |
| | } |
| | ) |
| |
|
| | cfg = ConfigProblemBase( |
| | dataset=ConfigNLPCausalLMDataset( |
| | prompt_column=("prompt",), |
| | answer_column="answer", |
| | parent_id_column="None", |
| | system_column="system", |
| | text_system_start="System:", |
| | text_prompt_start="Prompt:", |
| | text_answer_separator="Answer:", |
| | add_eos_token_to_answer=True, |
| | ), |
| | tokenizer=ConfigNLPCausalLMTokenizer(max_length=513), |
| | ) |
| |
|
| | cfg.llm_backbone = "EleutherAI/pythia-2.8b-deduped" |
| |
|
| | dataset = CustomDataset(df, cfg) |
| | assert len(dataset) == 3 |
| |
|
| | for i in range(3): |
| | result = dataset[i] |
| | assert isinstance(result, dict) |
| | assert set(result.keys()) == { |
| | "labels", |
| | "input_ids", |
| | "attention_mask", |
| | "prompt_input_ids", |
| | "prompt_attention_mask", |
| | "answer_input_ids", |
| | "answer_attention_mask", |
| | } |
| |
|
| | assert ( |
| | dataset.tokenizer.decode(result["input_ids"], skip_special_tokens=True) |
| | == f"System:system {i+1}" |
| | f"Prompt:prompt {i+1}" |
| | f"Answer:answer {i+1}" |
| | ) |
| |
|
| | assert ( |
| | dataset.tokenizer.decode( |
| | result["prompt_input_ids"], skip_special_tokens=True |
| | ) |
| | == f"System:system {i+1}" |
| | f"Prompt:prompt {i+1}" |
| | "Answer:" |
| | ) |
| |
|
