1
$\begingroup$

A normal and stratified split option is provided by sklearn method that can be used for ML problems like multi-class classification. This is relatively easier to do as (1) one sample has one class, and (2) you can split samples per class-wise to have the equal distribution of classes in train-val-test splits.

Now there seems to be a problem with NER (Named Entity Recognition) problem, as (1) there could be multiple entities, and also (2) each sample may have a different distribution of entities. So for example, say we have the following sample set,

Sample 1: contains DATE, PER, ORG
Sample 2: contains DATE, PER
Sample 3: contains DATE, ORG
Sample 4: contains PER, ORG
Sample 5: contains ORG

Now the unique entities and their overall count are DATE=3, PER=3, and ORG=4. If you want to do an 80-20 train-test split (for simplicity's sake), the best option seems to be keeping Sample 1 in the test and rest in train - as only then you will have a somewhat desired distribution of entities in the splits. On the other hand, if you select Sample 5 as a test, for example, we won't have any DATE and PER instances in the test at all.

So this is my question -- what is the best practice to split the dataset at an entity level for the NER task? Do we even split at the entity level for stratification or randomly split at sample level a couple of times and pick the one with the best split at entity level?

$\endgroup$

1 Answer 1

0
$\begingroup$

I know it's a bit late, but I had the same question and developed a method which is available here:

!pip install deep_utils
from deep_utils import stratify_train_test_split_multi_label

The code is as follows:

from typing import Union
from deep_utils.utils.algorithm_utils.main import subset_sum
import numpy as np


def stratify_train_test_split_multi_label(x: Union[list, tuple, np.ndarray], y: np.ndarray, test_size=0.2,
                                          closest_ratio=False):
    """
        A handy function for splitting multi-label samples based on their number of classes. This is mainly useful for
    object detection and ner-like tasks that each sample may contain several objects/tags from different classes! The
    process of splitting starts from classes with the smallest number of samples to make sure their ratio is saved
    because they have small numbers of samples, retaining the ratio for them is challenging compared to those classes
    with more samples
    :param x: A list, Tuple or ndarray that contains the samples
    :param y: A 2D array that represents the number of labels in each class. Each column is representative of a class.
    As an example: y = np.array([[2, 3], [1, 1]]) says that sample one has
    two objects/tags for class 0 and 3 objects/tags for class 1 and so on
    :param test_size: size of the test set
    :param closest_ratio: For huge arrays extracting the closest ratio requires an intensive recursive function to work
     which could result in maximum recursion error. Being set to True will choose samples from the those with the smallest difference to the target number to ensure the best ratio. Set this variable to True if you are sure. by default is set to False.
    :return:
    >>> y = np.array([[1, 2, 0], [1, 0, 0], [1, 2, 0]])
    >>> x = np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]])
    >>> stratify_train_test_split_multi_label(x, y, test_size=0.3)
    (array([[2, 2, 2],
           [3, 3, 3]]), array([[1, 1, 1]]), array([[1, 0, 0],
           [1, 2, 0]]), array([[1, 2, 0]]))
    >>> x = np.array([[1, 2], [3, 4], [1, 2], [3, 4], [1, 2], [3, 4], [1, 2], [3, 4]])
    >>> y = np.array([[0, 0], [0, 0], [0, 1], [0, 1], [1, 1], [1, 1], [1, 0], [1, 0]])
    >>> x_train, x_test, y_train, y_test = stratify_train_test_split_multi_label(x, y, test_size=0.5, closest_ratio=False)
    >>> x_train
    array([[1, 2],
           [3, 4],
           [1, 2],
           [3, 4]])
    >>> x_test
    array([[1, 2],
           [3, 4],
           [1, 2],
           [3, 4]])
    >>> y_train
    array([[0, 1],
           [0, 1],
           [1, 0],
           [1, 0]])
    >>> y_test
    array([[1, 1],
           [1, 1],
           [0, 0],
           [0, 0]])
    >>> print("class ratio:", tuple(y_test.sum(0) / y.sum(0)))
    class ratio: (0.5, 0.5)
    >>> print("sample ratio:", y_test.shape[0] / y.shape[0])
    sample ratio: 0.5
    """
    assert len(y.shape) == 2, "y should be 2D"
    assert test_size > 0.0, "test_size cannot be a zero or negative value!"
    x = np.array(x, dtype=np.object_) if not isinstance(x, np.ndarray) else x

    # excluding samples with no objects/tags
    non_objects = np.any(y.sum(1) == 0)
    if non_objects:
        y_no_objects = y[y.sum(1) == 0]
        x_no_objects = x[y.sum(1) == 0]
        x = x[y.sum(1) > 0]
        y = y[y.sum(1) > 0]

    available_samples = np.ones((y.shape[0]), dtype=np.bool8)
    test_samples = np.zeros((y.shape[0]), dtype=np.bool8)
    train_samples = np.zeros((y.shape[0]), dtype=np.bool8)
    class_sample_counts = y.sum(axis=0)
    ideal_train_size = np.floor(sum(class_sample_counts) * (1 - test_size))

    # stratify starts from a class with the lowest number of samples
    class_indices = np.argsort(class_sample_counts)
    for class_index in class_indices:
        test_number_samples = y[:, class_index][test_samples].sum()
        n_test = np.ceil(class_sample_counts[class_index] * test_size)
        n_test = max(0, n_test - test_number_samples)
        input_labels = y[:, class_index].copy()
        input_labels[np.invert(available_samples)] = 0
        if n_test == 0 or len(input_labels) == 0:
            continue
        if closest_ratio:
            chosen_indices, *_ = subset_sum(input_numbers=input_labels, target_number=n_test)
        else:
            sorted_indices = np.argsort(input_labels)
            cum_sum_values = np.cumsum(input_labels[sorted_indices])
            chosen_indices = sorted_indices[cum_sum_values < n_test].tolist()
            if len(chosen_indices) < len(sorted_indices):
                chosen_indices.append(sorted_indices[len(chosen_indices)])
        # Update available_samples, train_samples, test_samples
        for update_index, n_label in enumerate(input_labels):
            if n_label == 0:
                # samples that have no elements are ignored ...
                continue
            if update_index in chosen_indices:
                test_samples[update_index] = True
                train_samples[update_index] = False
            else:
                test_samples[update_index] = False
                train_samples[update_index] = True
            available_samples[update_index] = False
    # Allocating all the remaining samples to train because the code structure ensures the ratio of test
    # samples to the whole dataset.
    train_samples = np.bitwise_or(train_samples, np.bitwise_not(test_samples))

    if non_objects:
        # splitting samples with no objects trying to save the balance between train and test numbers

        train_left = int(ideal_train_size - sum(train_samples))
        indices = np.arange(len(y_no_objects))
        np.random.shuffle(indices)

        x_no_objects_train, y_no_objects_train = x_no_objects[:train_left], y_no_objects[:train_left]
        x_no_objects_test, y_no_objects_test = x_no_objects[train_left:], y_no_objects[train_left:]

        return np.concatenate([x[train_samples], x_no_objects_train]), \
               np.concatenate([x[test_samples], x_no_objects_test]), \
               np.concatenate([y[train_samples], y_no_objects_train]), \
               np.concatenate([y[test_samples], y_no_objects_test])
    else:
        return x[train_samples], x[test_samples], y[train_samples], y[test_samples]

Two examples are provided in the code's description.

Link to stratify_train_test_split_multi_label code

Link to deep_utils library

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.