| import os |
| import pathlib |
| import random |
| import unittest |
| from unittest.mock import MagicMock |
|
|
| import pandas as pd |
| import pytest |
|
|
| from llm_studio.app_utils.default_datasets import ( |
| prepare_default_dataset_causal_language_modeling, |
| ) |
| from llm_studio.src.datasets.conversation_chain_handler import ConversationChainHandler |
| from llm_studio.src.utils.data_utils import load_train_valid_data |
|
|
|
|
| @pytest.fixture |
| def cfg_mock(): |
| cfg = MagicMock() |
| cfg.dataset.train_dataframe = "/path/to/train/data" |
| cfg.dataset.validation_dataframe = "/path/to/validation/data" |
|
|
| cfg.dataset.system_column = "None" |
| cfg.dataset.prompt_column = "prompt" |
| cfg.dataset.answer_column = "answer" |
|
|
| cfg.dataset.validation_size = 0.2 |
| return cfg |
|
|
|
|
| @pytest.fixture |
| def read_dataframe_drop_missing_labels_mock(monkeypatch): |
| data = { |
| "prompt": [f"Prompt{i}" for i in range(100)], |
| "answer": [f"Answer{i}" for i in range(100)], |
| "id": list(range(100)), |
| } |
| df = pd.DataFrame(data) |
| mock = MagicMock(return_value=df) |
| monkeypatch.setattr( |
| "llm_studio.src.utils.data_utils.read_dataframe_drop_missing_labels", mock |
| ) |
| return mock |
|
|
|
|
| numbers = list(range(100)) |
| random.shuffle( |
| numbers, |
| ) |
| groups = [numbers[n::13] for n in range(13)] |
|
|
|
|
| @pytest.fixture |
| def conversation_chain_ids_mock(monkeypatch): |
| def mocked_init(self, *args, **kwargs): |
| self.conversation_chain_ids = groups |
|
|
| with unittest.mock.patch.object( |
| ConversationChainHandler, "__init__", new=mocked_init |
| ): |
| yield |
|
|
|
|
| def test_get_data_custom_validation_strategy( |
| cfg_mock, read_dataframe_drop_missing_labels_mock |
| ): |
| cfg_mock.dataset.validation_strategy = "custom" |
| train_df, val_df = load_train_valid_data(cfg_mock) |
| assert len(train_df), len(val_df) == 100 |
|
|
|
|
| def test_get_data_automatic_split( |
| cfg_mock, read_dataframe_drop_missing_labels_mock, conversation_chain_ids_mock |
| ): |
| cfg_mock.dataset.validation_strategy = "automatic" |
| train_df, val_df = load_train_valid_data(cfg_mock) |
| train_ids = set(train_df["id"].tolist()) |
| val_ids = set(val_df["id"].tolist()) |
|
|
| assert len(train_ids.intersection(val_ids)) == 0 |
| assert len(train_ids) + len(val_ids) == 100 |
|
|
| shared_groups = [ |
| i for i in groups if not train_ids.isdisjoint(i) and not val_ids.isdisjoint(i) |
| ] |
| assert len(shared_groups) == 0 |
|
|
|
|
| def test_oasst_data_automatic_split(tmp_path: pathlib.Path): |
| prepare_default_dataset_causal_language_modeling(tmp_path) |
| assert len(os.listdir(tmp_path)) > 0, tmp_path |
| cfg_mock = MagicMock() |
| for file in os.listdir(tmp_path): |
| if file.endswith(".pq"): |
| cfg_mock.dataset.train_dataframe = os.path.join(tmp_path, file) |
|
|
| cfg_mock.dataset.system_column = "None" |
| cfg_mock.dataset.prompt_column = ("instruction",) |
| cfg_mock.dataset.answer_column = "output" |
| cfg_mock.dataset.parent_id_column = "parent_id" |
|
|
| cfg_mock.dataset.validation_strategy = "automatic" |
|
|
| for validation_size in [0.05, 0.1, 0.2, 0.3, 0.4, 0.5]: |
| cfg_mock.dataset.validation_size = validation_size |
|
|
| train_df, val_df = load_train_valid_data(cfg_mock) |
| assert set(train_df["parent_id"].dropna().values).isdisjoint( |
| set(val_df["id"].dropna().values) |
| ) |
| assert set(val_df["parent_id"].dropna().values).isdisjoint( |
| set(train_df["id"].dropna().values) |
| ) |
| assert (len(val_df) / (len(train_df) + len(val_df))) == pytest.approx( |
| validation_size, 0.05 |
| ) |
|
|
