Skip to content

resolve_base_data

Resolve input to a list of text strings.

Supports HuggingFace dataset ID, pandas DataFrame, CSV path, or list of strings. Data is shuffled with seed before applying max_size to avoid bias from ordered sources (e.g. chronological, single-author).

Parameters:

Name Type Description Default
source Union[str, DataFrame, List[str]]

HF dataset ID (str), CSV path (str ending in .csv or path exists), pandas DataFrame, or List[str].

 required
text_column str

Column name for text (default "text"). Ignored for List[str].

 'text'
max_size Optional[int]

Cap on number of items (applied after shuffle). None = no cap.

 None
split str

Dataset split for HF (default "train").

 'train'
seed int

Random seed for shuffle.

 42
hf_config Optional[str]

HuggingFace dataset config/subset (e.g. "20220301.en" for wikipedia). Only used when source is an HF dataset ID.

 None
trust_remote_code Optional[bool]

Optional value passed to load_dataset when loading from HF. None means do not pass the keyword.

 None

Returns:

Type Description
List[str]

List of strings (texts).
Source code in gradiend/data/core/base_loader.py 
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def resolve_base_data(
    source: Union[str, pd.DataFrame, List[str]],
    text_column: str = "text",
    max_size: Optional[int] = None,
    split: str = "train",
    seed: int = 42,
    hf_config: Optional[str] = None,
    trust_remote_code: Optional[bool] = None,
) -> List[str]:
    """Resolve input to a list of text strings.

    Supports HuggingFace dataset ID, pandas DataFrame, CSV path, or list of strings.
    Data is shuffled with `seed` before applying `max_size` to avoid bias from
    ordered sources (e.g. chronological, single-author).

    Args:
        source: HF dataset ID (str), CSV path (str ending in .csv or path exists),
            pandas DataFrame, or List[str].
        text_column: Column name for text (default "text"). Ignored for List[str].
        max_size: Cap on number of items (applied after shuffle). None = no cap.
        split: Dataset split for HF (default "train").
        seed: Random seed for shuffle.
        hf_config: HuggingFace dataset config/subset (e.g. "20220301.en" for wikipedia).
            Only used when source is an HF dataset ID.
        trust_remote_code: Optional value passed to load_dataset when loading from HF.
            None means do not pass the keyword.

    Returns:
        List of strings (texts).
    """
    if not isinstance(text_column, str):
        raise TypeError(f"text_column must be str, got {type(text_column).__name__}")
    if max_size is not None and not isinstance(max_size, int):
        raise TypeError(f"max_size must be int or None, got {type(max_size).__name__}")
    if not isinstance(split, str):
        raise TypeError(f"split must be str, got {type(split).__name__}")
    if not isinstance(seed, int):
        raise TypeError(f"seed must be int, got {type(seed).__name__}")
    if hf_config is not None and not isinstance(hf_config, str):
        raise TypeError(f"hf_config must be str or None, got {type(hf_config).__name__}")
    if trust_remote_code is not None and not isinstance(trust_remote_code, bool):
        raise TypeError(f"trust_remote_code must be bool or None, got {type(trust_remote_code).__name__}")

    texts: List[str]
    if isinstance(source, list):
        if not all(isinstance(x, str) for x in source):
            raise TypeError("source as list must contain only strings")
        texts = [str(x).strip() for x in source if str(x).strip()]
    elif isinstance(source, pd.DataFrame):
        if text_column not in source.columns:
            raise ValueError(f"DataFrame missing column '{text_column}'. Columns: {list(source.columns)}")
        texts = source[text_column].dropna().astype(str).str.strip().tolist()
        texts = [x for x in texts if x]
    elif isinstance(source, str):
        texts = _load_from_string_source(
            source, text_column, split, hf_config, trust_remote_code, max_size
        )
    else:
        raise TypeError(f"source must be str, DataFrame, or List[str]; got {type(source)}")

    rng = random.Random(seed)
    rng.shuffle(texts)
    if max_size is not None and len(texts) > max_size:
        texts = texts[:max_size]
    logger.debug(f"resolve_base_data: {len(texts)} texts")
    return texts