This repository is dedicated to providing cutting-edge tools and methodologies to evaluate and curate datasets specifically designed for Large Language Models (LLMs). Leveraging the capabilities of LLMs themselves, combined with programmatic best practices, our toolkit ensures a robust evaluation and refinement process for your datasets.
from typing import List, Tuple
from datasets import DatasetDict, Dataset
from transformers import PreTrainedTokenizer
from datawarden.utils import analyze_token_counts

class DataWardenDataset:
    """
    A custom dataset class for processing various types of input data.
    """

    def __init__(self, dataset: DatasetDict) -> None:
        """
        Initialize the dataset.

        Args:
            dataset (DatasetDict): A dataset containing 'train' and 'test' subsets.
        """
        self.dataset = dataset
        self.train_questions, self.train_answers = [], []
        self.test_questions, self.test_answers = [], []

        self._process_dataset(self.dataset['train'])
        self._process_dataset(self.dataset['test'])

        self.train_data = [(self.train_questions[i], self.train_answers[i]) for i in range(len(self.train_questions))]
        self.test_data = [(self.test_questions[i], self.test_answers[i]) for i in range(len(self.test_questions))]

    def _process_dataset(self, dataset_subset: Dataset) -> None:
        """
        Process a dataset subset and extract questions and answers.

        Args:
            dataset_subset: A subset of the dataset to be processed.
        """
        questions, answers = [], []

        if 'instruction' in dataset_subset.features:
            # Extract 'instruction' and 'input' as questions, 'output' as answers
            for row in dataset_subset:
                question = f"{row['instruction']} {row['input']}"
                answer = row['output']
                questions.append([question])
                answers.append([answer])
        elif 'conversations' in dataset_subset.features:
            # Extract 'value' where 'from' is 'human' as questions, 'gpt' as answers
            for row in dataset_subset:
                conversation = row['conversations']
                human_messages = [msg['value'] for msg in conversation if msg['from'] == 'human']
                gpt_messages = [msg['value'] for msg in conversation if msg['from'] == 'gpt']
                questions.append(human_messages)
                answers.append(gpt_messages)
        elif 'text' in dataset_subset.features:
            # Extract text after 'Human:' as questions, after 'Assistant:' as answers
            for row in dataset_subset:
                text_segments = row['text'].split('### ')
                human_segments = [seg.split('Human: ')[-1].strip() for seg in text_segments if 'Human:' in seg]
                assistant_segments = [seg.split('Assistant: ')[-1] for seg in text_segments if 'Assistant:' in seg]
                questions.append(human_segments)
                answers.append(assistant_segments)

        if dataset_subset == self.dataset['train']:
            self.train_questions = questions
            self.train_answers = answers
        elif dataset_subset == self.dataset['test']:
            self.test_questions = questions
            self.test_answers = answers

    def get_token_counts(self, tokenizer: PreTrainedTokenizer, min_tokens_question: int = 256, min_tokens_answer: int = 256) -> Tuple[
            List[Tuple[List[str], List[str]]], List[Tuple[List[str], List[str]]], List[int], List[int]]:
        """
        Analyze token counts in the train and test subsets of the dataset.

        Args:
            tokenizer (PreTrainedTokenizer): The tokenizer to use for encoding text.
            min_tokens_question (int, optional): The minimum number of tokens required for questions. Defaults to 256.
            min_tokens_answer (int, optional): The minimum number of tokens required for answers. Defaults to 256.

        Returns:
            Tuple[List[Tuple[List[str], List[str]]], List[Tuple[List[str], List[str]]], List[int], List[int]]: A tuple containing four lists -
                - problematic_rows_train: Pairs in the train set with either the question or answer having fewer than the specified min_tokens.
                - clean_rows_train: Pairs in the train set with both the question and answer having at least the specified min_tokens.
                - problematic_indexes_train: Indexes of problematic rows in the train set.
                - clean_indexes_train: Indexes of clean rows in the train set.
                - problematic_rows_test: Pairs in the test set with either the question or answer having fewer than the specified min_tokens.
                - clean_rows_test: Pairs in the test set with both the question and answer having at least the specified min_tokens.
                - problematic_indexes_test: Indexes of problematic rows in the test set.
                - clean_indexes_test: Indexes of clean rows in the test set.
        """
        problematic_rows_train, clean_rows_train, problematic_indexes_train, clean_indexes_train = analyze_token_counts(
            self.train_data, tokenizer, min_tokens_question, min_tokens_answer)

        problematic_rows_test, clean_rows_test, problematic_indexes_test, clean_indexes_test = analyze_token_counts(
            self.test_data, tokenizer, min_tokens_question, min_tokens_answer)

        return (
            problematic_rows_train, clean_rows_train, problematic_indexes_train, clean_indexes_train,
            problematic_rows_test, clean_rows_test, problematic_indexes_test, clean_indexes_test
        )
datawarden/dataset.py
e-xperiments/datawarden
from typing import List, Tuple
from transformers import PreTrainedTokenizer

def analyze_token_counts(dataset: List[Tuple[List[str], List[str]]], tokenizer: PreTrainedTokenizer, min_tokens_question: int = 256, min_tokens_answer: int = 256) -> Tuple[List[Tuple[List[str], List[str]]], List[Tuple[List[str], List[str]]], List[int], List[int]]:
    """
    Analyze token counts in the dataset.

    Args:
        dataset (List[Tuple[List[str], List[str]]]): A list of pairs where the first element is a list of tokens for the question,
            and the second element is a list of tokens for the answer.
        tokenizer (PreTrainedTokenizer): The tokenizer to use for encoding text.
        min_tokens_question (int, optional): The minimum number of tokens required for questions. Defaults to 256.
        min_tokens_answer (int, optional): The minimum number of tokens required for answers. Defaults to 256.

    Returns:
        Tuple[List[Tuple[List[str], List[str]]], List[Tuple[List[str], List[str]]], List[int], List[int]]: A tuple containing four lists -
            - problematic_rows: Pairs with either the question or answer having fewer than the specified min_tokens.
            - clean_rows: Pairs with both the question and answer having at least the specified min_tokens.
            - problematic_indexes: Indexes of problematic rows in the input dataset.
            - clean_indexes: Indexes of clean rows in the input dataset.
    """
    problematic_rows = [(question, answer) for question, answer in dataset if sum(len(tokenizer.encode(tokens)) for tokens in question) < min_tokens_question or sum(len(tokenizer.encode(tokens)) for tokens in answer) < min_tokens_answer]
    clean_rows = [(question, answer) for question, answer in dataset if sum(len(tokenizer.encode(tokens)) for tokens in question) >= min_tokens_question and sum(len(tokenizer.encode(tokens)) for tokens in answer) >= min_tokens_answer]
    
    problematic_indexes = [i for i, (question, answer) in enumerate(dataset) if (question, answer) in problematic_rows]
    clean_indexes = [i for i, (question, answer) in enumerate(dataset) if (question, answer) in clean_rows]
    
    return problematic_rows, clean_rows, problematic_indexes, clean_indexes

datawarden/utils.py
e-xperiments/datawarden
from setuptools import setup, find_packages

setup(
    name='datawarden',
    version='0.0.1',
    description='This package is dedicated to providing cutting-edge tools and methodologies to evaluate and curate datasets specifically designed for Large Language Models (LLMs). Leveraging the capabilities of LLMs themselves, combined with programmatic best practices, our toolkit ensures a robust evaluation and refinement process for your datasets.',
    author='e-xperiments',
    author_email='ahm.rimer@gmail.com',
    packages=find_packages(),  # Automatically discover and include all packages
    install_requires=[
        # List your dependencies here
        'transformers',
    ],
)

setup.py
e-xperiments/datawarden
import pytest
from datasets import Dataset, DatasetDict
from datawarden.dataset import DataWardenDataset
from transformers import AutoTokenizer

# Define sample data for testing
sample_data_instruction = {
    'train': Dataset.from_dict({
        'instruction': ['instr1', 'instr2'],
        'input': ['input1', 'input2'],
        'output': ['output1', 'output2']
    }),
    'test': Dataset.from_dict({
        'instruction': ['test_instr1', 'test_instr2'],
        'input': ['test_input1', 'test_input2'],
        'output': ['test_output1', 'test_output2']
    })
}

sample_data_conversations = {
    'train': Dataset.from_dict({
        'conversations': [
            [{'from': 'human', 'value': 'Hello'}, {'from': 'gpt', 'value': 'Hi'}],
            [{'from': 'human', 'value': 'How are you?'}, {'from': 'gpt', 'value': 'I am fine.'}],
        ]
    }),
    'test': Dataset.from_dict({
        'conversations': [
            [{'from': 'human', 'value': 'Good morning'}, {'from': 'gpt', 'value': 'Good morning!'}],
            [{'from': 'human', 'value': 'What is your name?'}, {'from': 'gpt', 'value': 'I am a chatbot.'}],
        ]
    })
}

sample_data_text = {
    'train': Dataset.from_dict({
        'text': [
            '### Human: How are you? ### Assistant: I am fine.',
            '### Human: What is your name? ### Assistant: I am a chatbot.',
        ]
    }),
    'test': Dataset.from_dict({
        'text': [
            '### Human: Good morning ### Assistant: Good morning!',
            '### Human: How can I help you? ### Assistant: You can ask me anything.',
        ]
    })
}

@pytest.fixture(params=[sample_data_instruction, sample_data_conversations, sample_data_text])
def sample_dataset(request):
    return DatasetDict(request.param)

@pytest.fixture(params=[sample_data_instruction, sample_data_conversations, sample_data_text])
def sample_dataset_with_tokens(request):
    dataset = DataWardenDataset(request.param)
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")  # Replace with your desired tokenizer
    return dataset, tokenizer

def test_init(sample_dataset):
    dataset = DataWardenDataset(sample_dataset)
    if 'instruction' in sample_dataset['train'].features:
        assert dataset.train_questions == [['instr1 input1'], ['instr2 input2']]
        assert dataset.train_answers == [['output1'], ['output2']]
        assert dataset.test_questions == [['test_instr1 test_input1'], ['test_instr2 test_input2']]
        assert dataset.test_answers == [['test_output1'], ['test_output2']]
    elif 'conversations' in sample_dataset['train'].features:
        assert dataset.train_questions == [['Hello'], ['How are you?']] 
        assert dataset.train_answers == [['Hi'], ['I am fine.']]
        assert dataset.test_questions == [['Good morning'], ['What is your name?']]
        assert dataset.test_answers == [['Good morning!'], ['I am a chatbot.']]
    elif 'text' in sample_dataset['train'].features:
        assert dataset.train_questions == [['How are you?'], ['What is your name?']]
        assert dataset.train_answers == [['I am fine.'], ['I am a chatbot.']]
        assert dataset.test_questions == [['Good morning'], ['How can I help you?']]
        assert dataset.test_answers == [['Good morning!'], ['You can ask me anything.']]

def test_get_token_counts(sample_dataset_with_tokens):
    dataset, tokenizer = sample_dataset_with_tokens
    min_tokens_question = 10
    min_tokens_answer = 10

    (
        problematic_rows_train,
        clean_rows_train,
        problematic_indexes_train,
        clean_indexes_train,
        problematic_rows_test,
        clean_rows_test,
        problematic_indexes_test,
        clean_indexes_test,
    ) = dataset.get_token_counts(tokenizer, min_tokens_question, min_tokens_answer)
    
    assert isinstance(problematic_rows_train, list)
    assert isinstance(clean_rows_train, list)
    assert isinstance(problematic_indexes_train, list)
    assert isinstance(clean_indexes_train, list)
    assert isinstance(problematic_rows_test, list)
    assert isinstance(clean_rows_test, list)
    assert isinstance(problematic_indexes_test, list)
    assert isinstance(clean_indexes_test, list)

if __name__ == "__main__":
    pytest.main()

tests/test_dataset.py
e-xperiments/datawarden
import pytest
from transformers import AutoTokenizer
from datawarden.utils import analyze_token_counts

# Sample dataset for testing
TEST_DATASET = [
    (["How", "are", "you?"], ["I'm", "fine.", "Thanks!"]),
    (["What", "is", "your", "name?"], ["My", "name", "is", "ChatGPT."]),
    (["Hello"], ["Hi"]),
    (
        ["This", "is", "a", "long", "question", "with", "a", "long", "answer."],
        ["Yes,", "indeed.", "This", "is", "a", "very", "very", "long", "answer."],
    ),
]


@pytest.fixture
def tokenizer():
    """Fixture to initialize the tokenizer."""
    return AutoTokenizer.from_pretrained("bert-base-uncased")

def test_analyze_token_counts(tokenizer):
    # Test with default min_tokens_question and min_tokens_answer
    problematic_rows, clean_rows, problematic_indexes, clean_indexes = analyze_token_counts(TEST_DATASET, tokenizer)
    
    assert len(problematic_rows) == 4
    assert len(clean_rows) == 0
    assert problematic_indexes == [0, 1, 2, 3]
    assert clean_indexes == []

    # Test with custom min_tokens_question and min_tokens_answer
    problematic_rows, clean_rows, problematic_indexes, clean_indexes = analyze_token_counts(TEST_DATASET, tokenizer, min_tokens_question=5, min_tokens_answer=5)
    
    assert len(problematic_rows) == 1
    assert len(clean_rows) == 3
    assert problematic_indexes == [2]
    assert clean_indexes == [0, 1, 3]

    # Test with an empty dataset
    empty_dataset = []
    problematic_rows, clean_rows, problematic_indexes, clean_indexes = analyze_token_counts(empty_dataset, tokenizer)
    
    assert len(problematic_rows) == 0
    assert len(clean_rows) == 0
    assert problematic_indexes == []
    assert clean_indexes == []

if __name__ == "__main__":
    pytest.main()

tests/test_utils.py
e-xperiments/datawarden