unit tests

Author: benjaminye

test_utils/__init__.py

@@ -0,0 +1,73 @@
+from llmtune.pydantic_models.config_model import (
+    AblationConfig,
+    BitsAndBytesConfig,
+    Config,
+    DataConfig,
+    InferenceConfig,
+    LoraConfig,
+    ModelConfig,
+    SftArgs,
+    TrainingArgs,
+    TrainingConfig,
+)
+
+
+def get_sample_config():
+    """Function to return a comprehensive Config object for testing."""
+    return Config(
+        save_dir="./test",
+        ablation=AblationConfig(
+            use_ablate=False,
+        ),
+        model=ModelConfig(
+            hf_model_ckpt="NousResearch/Llama-2-7b-hf",
+            device_map="auto",
+            torch_dtype="auto",
+            quantize=False,
+            bitsandbytes=BitsAndBytesConfig(
+                load_in_8bit=False,
+                load_in_4bit=False,
+                bnb_4bit_compute_dtype="float32",
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+            ),
+        ),
+        lora=LoraConfig(
+            r=8,
+            task_type="CAUSAL_LM",
+            lora_alpha=16,
+            bias="none",
+            lora_dropout=0.1,
+            target_modules=None,
+            fan_in_fan_out=False,
+        ),
+        training=TrainingConfig(
+            training_args=TrainingArgs(
+                num_train_epochs=1,
+                per_device_train_batch_size=1,
+                gradient_accumulation_steps=1,
+                optim="adamw_8bit",
+                learning_rate=2.0e-4,
+                logging_steps=100,
+            ),
+            sft_args=SftArgs(max_seq_length=512, neftune_noise_alpha=None),
+        ),
+        inference=InferenceConfig(
+            max_length=128,
+            do_sample=False,
+            num_beams=5,
+            temperature=1.0,
+            top_k=50,
+            top_p=1.0,
+            use_cache=True,
+        ),
+        data=DataConfig(
+            file_type="json",
+            path="path/to/dataset.json",
+            prompt="Your prompt here {column_name}",
+            prompt_stub="Stub for prompt {column_name}",
+            train_size=0.9,
+            test_size=0.1,
+            train_test_split_seed=42,
+        ),
+    )

test_utils/test_config.py

@@ -0,0 +1 @@
+# TODO

tests/data/test_dataset_generator.py

@@ -0,0 +1,136 @@
+import pytest
+from unittest.mock import patch, MagicMock, mock_open
+
+from llmtune.data.ingestor import (
+    CsvIngestor,
+    HuggingfaceIngestor,
+    JsonIngestor,
+    JsonlIngestor,
+    get_ingestor,
+)
+
+from datasets import Dataset
+
+
+def test_get_ingestor():
+    assert isinstance(get_ingestor("json")(""), JsonIngestor)
+    assert isinstance(get_ingestor("jsonl")(""), JsonlIngestor)
+    assert isinstance(get_ingestor("csv")(""), CsvIngestor)
+    assert isinstance(get_ingestor("huggingface")(""), HuggingfaceIngestor)
+
+    with pytest.raises(ValueError):
+        get_ingestor("unsupported_type")
+
+
+def test_json_ingestor_to_dataset(mocker):
+    mock_generator = mocker.patch("llmtune.data.ingestor.JsonIngestor._json_generator")
+    mock_dataset = mocker.patch("llmtune.data.ingestor.Dataset")
+    JsonIngestor("").to_dataset()
+
+    mock_dataset.from_generator.assert_called_once_with(mock_generator)
+
+
+def test_jsonl_ingestor_to_dataset(mocker):
+    mock_generator = mocker.patch(
+        "llmtune.data.ingestor.JsonlIngestor._jsonl_generator"
+    )
+    mock_dataset = mocker.patch("llmtune.data.ingestor.Dataset")
+    JsonlIngestor("").to_dataset()
+
+    mock_dataset.from_generator.assert_called_once_with(mock_generator)
+
+
+def test_csv_ingestor_to_dataset(mocker):
+    mock_generator = mocker.patch("llmtune.data.ingestor.CsvIngestor._csv_generator")
+    mock_dataset = mocker.patch("llmtune.data.ingestor.Dataset")
+    CsvIngestor("").to_dataset()
+
+    mock_dataset.from_generator.assert_called_once_with(mock_generator)
+
+
+def test_huggingface_to_dataset(mocker):
+    # Setup
+    path = "some_path"
+    ingestor = HuggingfaceIngestor(path)
+    mock_concatenate_datasets = mocker.patch(
+        "llmtune.data.ingestor.concatenate_datasets"
+    )
+    mock_load_dataset = mocker.patch("llmtune.data.ingestor.load_dataset")
+    mock_dataset = mocker.patch("llmtune.data.ingestor.Dataset")
+
+    # Configure the mock objects
+    mock_dataset = MagicMock(spec=Dataset)
+    mock_load_dataset.return_value = {"train": mock_dataset, "test": mock_dataset}
+    mock_concatenate_datasets.return_value = mock_dataset
+
+    # Execute
+    result = ingestor.to_dataset()
+
+    # Assert
+    assert isinstance(result, Dataset)
+    mock_load_dataset.assert_called_once_with(path)
+    mock_concatenate_datasets.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "file_content,expected_output",
+    [
+        (
+            '[{"column1": "value1", "column2": "value2"}, {"column1": "value3", "column2": "value4"}]',
+            [
+                {"column1": "value1", "column2": "value2"},
+                {"column1": "value3", "column2": "value4"},
+            ],
+        )
+    ],
+)
+def test_json_ingestor_generator(file_content, expected_output, mocker):
+    mocker.patch("builtins.open", mock_open(read_data=file_content))
+    mocker.patch("ijson.items", side_effect=lambda f, prefix: iter(expected_output))
+    ingestor = JsonIngestor("dummy_path.json")
+
+    assert list(ingestor._json_generator()) == expected_output
+
+
+@pytest.mark.parametrize(
+    "file_content,expected_output",
+    [
+        (
+            '{"column1": "value1", "column2": "value2"}\n{"column1": "value3", "column2": "value4"}',
+            [
+                {"column1": "value1", "column2": "value2"},
+                {"column1": "value3", "column2": "value4"},
+            ],
+        )
+    ],
+)
+def test_jsonl_ingestor_generator(file_content, expected_output, mocker):
+    mocker.patch("builtins.open", mock_open(read_data=file_content))
+    mocker.patch(
+        "ijson.items",
+        side_effect=lambda f, prefix, multiple_values: (
+            iter(expected_output) if multiple_values else iter([])
+        ),
+    )
+    ingestor = JsonlIngestor("dummy_path.jsonl")
+
+    assert list(ingestor._jsonl_generator()) == expected_output
+
+
+@pytest.mark.parametrize(
+    "file_content,expected_output",
+    [
+        (
+            "column1,column2\nvalue1,value2\nvalue3,value4",
+            [
+                {"column1": "value1", "column2": "value2"},
+                {"column1": "value3", "column2": "value4"},
+            ],
+        )
+    ],
+)
+def test_csv_ingestor_generator(file_content, expected_output, mocker):
+    mocker.patch("builtins.open", mock_open(read_data=file_content))
+    ingestor = CsvIngestor("dummy_path.csv")
+
+    assert list(ingestor._csv_generator()) == expected_output

tests/data/test_ingestor.py

@@ -0,0 +1,28 @@
+import pytest
+
+from llmtune.finetune.generics import Finetune
+
+
+class MockFinetune(Finetune):
+    def finetune(self):
+        return "finetuning complete"
+
+    def save_model(self):
+        return "model saved"
+
+
+def test_finetune_method():
+    mock_finetuner = MockFinetune()
+    result = mock_finetuner.finetune()
+    assert result == "finetuning complete"
+
+
+def test_save_model_method():
+    mock_finetuner = MockFinetune()
+    result = mock_finetuner.save_model()
+    assert result == "model saved"
+
+
+def test_finetune_abstract_class_instantiation():
+    with pytest.raises(TypeError):
+        _ = Finetune()