Chunking & tokenization with Data Prep Kit

WIKI_ACCESS_TOKEN='eyxxx'
HF_READ_ACCESS_TOKEN='hf_xxx'

%%capture
%pip install "data-prep-toolkit-transforms[docling2parquet,doc_chunk,tokenization]"
%pip install pandas
%pip install "numpy<2.0"
from dotenv import load_dotenv

load_dotenv(".env", override=True)

def load_corpus(articles: list, folder: str) -> int:
    import os
    import re

    import requests

    headers = {"Authorization": f"Bearer {os.getenv('WIKI_ACCESS_TOKEN')}"}
    count = 0
    for article in articles:
        try:
            endpoint = f"https://api.enterprise.wikimedia.com/v2/articles/{article}"
            response = requests.get(endpoint, headers=headers)
            response.raise_for_status()
            doc = response.json()
            for article in doc:
                filename = re.sub(r"[^a-zA-Z0-9_]", "_", article["name"])
                with open(f"{folder}/{filename}.html", "w") as f:
                    f.write(article["article_body"]["html"])
                    count = count + 1
        except Exception as e:
            print(f"Failed to retrieve content: {e}")
    return count

import os
import tempfile

datafolder = tempfile.mkdtemp(dir=os.getcwd())
articles = ["Science,_technology,_engineering,_and_mathematics"]
assert load_corpus(articles, datafolder) > 0, "Faild to download any documents"

%%capture
from dpk_docling2parquet import Docling2Parquet, docling2parquet_contents_types

result = Docling2Parquet(
    input_folder=datafolder,
    output_folder=f"{datafolder}/docling2parquet",
    data_files_to_use=[".html"],
    docling2parquet_contents_type=docling2parquet_contents_types.MARKDOWN,  # markdown
).transform()

%%capture
from dpk_doc_chunk import DocChunk

result = DocChunk(
    input_folder=f"{datafolder}/docling2parquet",
    output_folder=f"{datafolder}/doc_chunk",
    doc_chunk_chunking_type="li_markdown",
    doc_chunk_chunk_size_tokens=128,  # default 128
    doc_chunk_chunk_overlap_tokens=30,  # default 30
).transform()

%%capture
from dpk_tokenization import Tokenization

Tokenization(
    input_folder=f"{datafolder}/doc_chunk",
    output_folder=f"{datafolder}/tkn",
    tkn_tokenizer="hf-internal-testing/llama-tokenizer",
    tkn_chunk_size=20_000,
).transform()

from pathlib import Path

import pandas as pd

parquet_files = list(Path(f"{datafolder}/tkn/").glob("*.parquet"))
pd.concat(pd.read_parquet(file) for file in parquet_files)