data_loader.py

from datasets import load_dataset

def load_documents(dataset_name="ag_news"):
    """
    Повертає список документів (рядків)
    """
    ds = load_dataset(dataset_name, split="train")
    docs = []
    for item in ds:
        text = item["text"]
        docs.append(text)
    return docs