web.py

import re
from datasets import load_dataset

OUTPUT_FILE = "pretrain_small.txt"

MAX_CHARS = 450_000_000  # ~300MB raw text cap (safe under 400MB)

HTML_RE = re.compile(r"<.*?>")
MULTISPACE_RE = re.compile(r"\s+")

def clean(text):
    text = text.lower()
    text = re.sub(HTML_RE, " ", text)
    text = re.sub(MULTISPACE_RE, " ", text)
    return text.strip()

def main():
    print("Loading streaming datasets...")

    wiki = load_dataset(
        "wikimedia/wikipedia",
        "20231101.en",
        split="train",
        streaming=True
    )

    web = load_dataset(
        "openwebtext",
        split="train",
        streaming=True
    )

    total_chars = 0

    with open(OUTPUT_FILE, "w", encoding="utf-8") as f:

        # Wikipedia first
        for item in wiki:
            text = clean(item["text"])
            if len(text) < 10:
                continue

            f.write(text + "\n\n")
            total_chars += len(text)

            if total_chars > MAX_CHARS:
                break

        # OpenWebText second
        for item in web:
            text = clean(item["text"])
            if len(text) < 10:
                continue

            f.write(text + "\n\n")
            total_chars += len(text)

            if total_chars > MAX_CHARS:
                break

    print(f"Done. Total chars: {total_chars}")
    print("Saved to", OUTPUT_FILE)

if __name__ == "__main__":
    main()