# Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
from datasets import load_dataset
import json
import unicodedata

def remove_control_characters(s):
    return "".join(ch for ch in s if unicodedata.category(ch)[0]!="C")

from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 1536, # chars, not llm tokens
    chunk_overlap  = 0,
    length_function = len,
    is_separator_regex = False,
)

documents = load_dataset('Shitao/MLDR', "corpus-en", split='corpus')    
feed_file = "/tmp/vespa_feed_file_en.json"
with open(feed_file, "w") as f:
    for doc in documents:
        id = doc["docid"]
        text = doc['text']
        chunks = text_splitter.create_documents([text])
        text_chunks = [chunk.page_content for chunk in chunks]
        text_chunks = [remove_control_characters(chunk) for chunk in text_chunks]
        vespa_feed_doc = {
            "put": "id:%s:doc::%s" %  ("en", id),
                "fields": {
                    "text": text_chunks
                }
            }
        f.write(json.dumps(vespa_feed_doc))
        f.write("\n")

langchain.text_splitter.RecursiveCharacterTextSplitter
[((362, 477), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(1536)', 'chunk_overlap': '(0)', 'length_function': 'len', 'is_separator_regex': '(False)'}), '(chunk_size=1536, chunk_overlap=0,\n    length_function=len, is_separator_regex=False)\n', (392, 477), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((539, 595), 'datasets.load_dataset', 'load_dataset', (['"""Shitao/MLDR"""', '"""corpus-en"""'], {'split': '"""corpus"""'}), "('Shitao/MLDR', 'corpus-en', split='corpus')\n", (551, 595), False, 'from datasets import load_dataset\n'), ((1148, 1174), 'json.dumps', 'json.dumps', (['vespa_feed_doc'], {}), '(vespa_feed_doc)\n', (1158, 1174), False, 'import json\n'), ((244, 268), 'unicodedata.category', 'unicodedata.category', (['ch'], {}), '(ch)\n', (264, 268), False, 'import unicodedata\n')]