from zephon import Pipeline
from zephon.io import Dataset
from zephon.work import StaticMixtureWorkSource, TokenEstimation
if __name__ == "__main__":
    wiki = Dataset.from_path(
        "wiki", "hf://Salesforce/wikitext/wikitext-2-raw-v1/train"
    )
    news = Dataset.from_path("news", "hf://fancyzhx/ag_news/default/train")
    source = StaticMixtureWorkSource(
        datasets=[wiki, news],
        mixture={"wiki": 0.5, "news": 0.5}, 
        token_estimation=TokenEstimation(), # Enable token-aware mixing since sources are strings
        seed=42,
    )
    pipeline = (
        Pipeline(source)
        .options(
            io_options={"cache": {"enabled": True}},
        )
        .tokenize(tokenizer_id="gpt2", field="text", parallelism=4)
        .pack_flat(max_length=2048, algorithm="wrap")
        .batch(microbatch_size=8)
    )
    batch = next(iter(pipeline))
    print(batch.to_training())