from collections import Counter
import json
import logging
import os
from enum import Enum
from urllib.error import HTTPError
from typing import Iterable, List, Optional, TextIO, Tuple, Union

import wget

import pyarrow as pa

from smart_open import open

from pathlib import Path
from lm_datasets.io.parquet import save_texts_to_parquet_chunks

from lm_datasets.systems import get_path_by_system
from lm_datasets.utils import get_parquet_compression

MILLION = 1_000_000
BILLION = 1_000_000_000

KB = 1024
MB = 1024 * 1024
GB = 1024 * 1024 * 1024
TB = 1024 * 1024 * 1024


TOKENS_PER_BYTE = 0.29335  # The Pile GPT2 tokenizeer (from the paper)


class Availability(Enum):
    DIRECT_DOWNLOAD = "direct_download"  # Yes - it has a direct download link or links
    SIGNIN_DOWNLOAD = "signin_download"  # Yes - after signing a user agreement"
    ON_REQUEST = "on_request"  # No - but the current owners/custodians have contact information for data queries
    PRIVATE = "private"  # No - we would need to spontaneously reach out to the current owners/custodians

    def __str__(self):
        return str(self.value)


class QualityWarning(Enum):
    SHORT_TEXT = "short_text"  # Dataset contains mostly short text (below paragraphs)
    BAD_ENCODING = "bad_encoding"  # Text might bit wrong encoded (utf-8 etc.)
    BAD_WHITESPACES = (  # Text might contain bad whitespaces, e.g., before punctuation (converting tokens to text)
        "bad_whitespaces"
    )
    BAD_LINEBREAKS = "bad_linebreaks"  # Text has missing or too many line breaks
    BAD_PUNCTUATION = "bad_punctuation"

    def __str__(self):
        return str(self.value)


class Genre(Enum):
    LEGAL = "legal"
    SCIENCE = "science"
    MATH = "math"
    NEWS = "news"
    DIALOGUE = "dialogue"
    GOVERNMENT = "government"
    LITERATURE = "literature"
    BIOMEDICAL = "biomedical"

    def __str__(self):
        return str(self.value)


# deprecated
AVAILIBILITY_OPTIONS = [
    "Yes - it has a direct download link or links",
    "Yes - after signing a user agreement",
    "No - but the current owners/custodians have contact information for data queries",
    "No - we would need to spontaneously reach out to the current owners/custodians",
]

logger = logging.getLogger(__name__)


class BaseDataset(object):
    """
    Base class for all datasets. It implements all generic loading, processing, and writing methods.
    """

    DATASET_ID = None
    SOURCE_ID = None

    TITLE = None
    DESCRIPTION = None
    HOMEPAGE = None
    AVAILIBILITY = None
    DOWNLOAD_URLS = []
    LOCAL_DIRS = []
    VERSION = None
    DOI = None

    LICENSE = None
    PII = None

    LANGUAGES = []

    WEB_CRAWLED = False
    QUALITY_WARNINGS = []
    GENRES = []
    USED_BY = None
    DUMMY = False
    SINGLE_OUTPUT_FILE = True

    # Statistics
    TOKENS = None
    BYTES = None

    counter = Counter()

    def __init__(
        self,
        output_dir,
        raw_datasets_dir: Optional[str] = None,
        workers: int = 1,
        output_text_field: str = "text",
        override_output: bool = False,
        limit: int = 0,
        skip_items: int = 0,
        hf_auth_token: str = None,
        print_write_progress: int = 10_000,
        min_length: int = 250,
        json_ensure_ascii: bool = False,
        title_delimiter: str = ":\n\n",
        paragraph_delimiter: str = "\n\n",
        sentence_delimiter: str = " ",
        output_format: str = "jsonl",  # options: jsonl, parquet
        output_compression: Optional[
            str
        ] = None,  # jsonl: gzip, parquet: ‘NONE’, ‘SNAPPY’, ‘GZIP’, ‘BROTLI’, ‘LZ4’, ‘ZSTD’
        output_batch_size: int = 1000,
        shuffled_output_dir: Optional[str] = None,
        max_output_chunk_uncompressed_bytes: int = 0,
    ) -> None:
        self.output_dir = output_dir
        self.raw_datasets_dir = raw_datasets_dir
        self.workers = workers
        self.output_text_field = output_text_field
        self.override_output = override_output
        self.limit = limit
        self.skip_items = skip_items
        self.hf_auth_token = hf_auth_token
        self.print_write_progress = print_write_progress
        self.min_length = min_length
        self.json_ensure_ascii = json_ensure_ascii
        self.title_delimiter = title_delimiter
        self.paragraph_delimiter = paragraph_delimiter
        self.sentence_delimiter = sentence_delimiter
        self.output_format = output_format
        self.output_compression = output_compression
        self.output_batch_size = output_batch_size
        self.shuffled_output_dir = shuffled_output_dir
        self.max_output_chunk_uncompressed_bytes = max_output_chunk_uncompressed_bytes

    def get_source_id(self):
        if self.SOURCE_ID:
            return self.SOURCE_ID
        else:
            return self.DATASET_ID

    def get_language_code(self, unknown: str = "unknown", mixed: str = "mixed"):
        if len(self.LANGUAGES) == 1:
            lang = self.LANGUAGES[0]
        elif len(self.LANGUAGES) == 0:
            lang = unknown
        else:
            lang = mixed

        return lang

    def has_output_files(self, min_file_size: int = 1, shuffled=False) -> bool:
        return self.has_single_output_file(
            min_file_size=min_file_size, shuffled=shuffled
        ) or self.has_chunked_output_files(min_file_size=min_file_size, shuffled=shuffled)

    def has_single_output_file(self, min_file_size: int = 1, shuffled=False) -> bool:
        fp = self.get_single_output_file_path(shuffled=shuffled)

        return os.path.exists(fp) and os.stat(fp).st_size >= min_file_size

    def has_chunked_output_files(self, min_file_size: int = 1, shuffled=False) -> bool:
        for fp in self.get_chunked_output_file_paths(shuffled=shuffled):
            if os.path.exists(fp) and os.stat(fp).st_size >= min_file_size:
                return True
            break

        return False

    def get_output_file_paths(self, single=False, chunked=False, shuffled=False) -> List[str]:
        if single:
            return [self.get_single_output_file_path(shuffled=shuffled)]
        elif chunked:
            return self.get_chunked_output_file_paths(shuffled=shuffled)
        else:
            # auto determine based on existing files
            if self.has_chunked_output_files(shuffled=shuffled):
                return self.get_chunked_output_file_paths(shuffled=shuffled)
            else:
                return [self.get_single_output_file_path(shuffled=shuffled)]

    def get_output_file_path(self):
        raise NotImplementedError("Use `get_output_file_paths` instead!")

    def get_output_extension(self, with_dot: bool = True, shuffled: bool = False) -> str:
        extension = "." if with_dot else ""

        if shuffled:
            extension += "shuffled."

        extension += self.output_format

        if self.output_format == "jsonl" and self.output_compression == "gzip":
            # Simply add ".gz" as extension as smart_open will take about the compression
            extension += ".gz"

        return extension

    def get_output_dir(self, shuffled=False):
        if shuffled:
            if self.shuffled_output_dir:
                return self.shuffled_output_dir
            raise ValueError("shuffled_output_dir is not set")
        else:
            return self.output_dir

    def get_single_output_file_path(self, shuffled=False) -> str:
        return os.path.join(
            self.get_output_dir(shuffled=shuffled), self.DATASET_ID + self.get_output_extension(shuffled=shuffled)
        )

    def get_chunked_output_file_paths(self, shuffled=False) -> List[str]:
        output_dir_path = Path(self.get_output_dir(shuffled=shuffled))

        return list(
            output_dir_path.glob(f"{self.DATASET_ID}.part-*-of-*{self.get_output_extension(shuffled=shuffled)}")
        )

    def get_chunked_output_file_path(self, part: int, total_parts: Optional[int] = None, shuffled=False) -> str:
        if total_parts is None:
            fn = f"{self.DATASET_ID}.part-{part:04d}{self.get_output_extension(shuffled=shuffled)}"
        else:
            fn = f"{self.DATASET_ID}.part-{part:04d}-of-{total_parts:04d}{self.get_output_extension(shuffled=shuffled)}"

        return os.path.join(self.get_output_dir(shuffled=shuffled), fn)

    def get_single_or_chunked_output_file_path(
        self, part: Optional[int] = None, total_parts: Optional[int] = None, shuffled=False
    ) -> str:
        if part is None:
            return self.get_single_output_file_path(shuffled=shuffled)
        else:
            return self.get_chunked_output_file_path(part, total_parts, shuffled=shuffled)

    # def has_output_file(self, min_file_size: int = 1):
    #     if self.SINGLE_OUTPUT_FILE:
    #         fps = [self.get_output_file_path()]
    #     else:
    #         fps = self.get_output_file_paths()

    #     for fp in fps:
    #         if os.path.exists(fp) and os.stat(fp).st_size >= min_file_size:
    #             pass
    #         else:
    #             return False

    #     return True

    def filter_texts(self, texts: Iterable[str]):
        """
        Applies basic filtering on the texts before saving
        """
        for text in texts:
            if self.min_length > 0 and len(text) < self.min_length:
                # skip because of short text length
                self.counter.update({"filtered_short_text": 1})
                continue

            yield text

    def remove_texts(self):
        for fp in self.get_output_file_paths():
            logger.warning(f"Removing {fp}")
            os.remove(fp)

    def save_texts(self, texts: Iterable[str], append: bool = False):
        """
        Save texts in different formats
        """
        if self.has_output_files() and not self.override_output:
            raise FileExistsError(f"Output exists already (override not enabled): {self.get_output_file_paths()}")

        if self.output_format == "jsonl":
            docs_count = self.save_texts_to_jsonl(texts, append=append)

        elif self.output_format == "parquet":
            if append:
                raise NotImplementedError("Appending is not supported by parquet output format")

            docs_count = self.save_texts_to_parquet(texts)

        else:
            raise ValueError(f"Unsupported output format: {self.output_format}")

        logger.info(f"Documents saved: {docs_count:,}")

        if docs_count == 0:
            logger.warning("No documents have been saved!")

            # delete empty output file
            if self.has_output_files():
                self.remove_texts()

        return docs_count

    def save_texts_to_parquet(self, texts: Iterable[str], file_path: Optional[str] = None, apply_filter: bool = True):
        """
        Save text in parquet (single column schema, in batches)
        """
        if file_path is None:
            file_path = self.get_output_file_paths(single=True)[0]

        if apply_filter:
            texts = self.filter_texts(texts)

        schema = pa.schema(
            [
                (self.output_text_field, pa.string()),
            ]
        )
        # Max. chunk size is multiplied with this factor
        # (to account for inaccurate chunk sizes due to batching)
        safety_factor = 0.975

        # Save as Parquet file
        logger.info(f"Writing parquet output ({self.output_batch_size=}; {self.limit=}; {self.output_compression=})")

        saved_docs, saved_chunks = save_texts_to_parquet_chunks(
            texts=texts,
            schema=schema,
            max_chunk_uncompressed_bytes=self.max_output_chunk_uncompressed_bytes * safety_factor,
            output_path_func=self.get_single_or_chunked_output_file_path,
            compression=get_parquet_compression(self.output_compression),
            batch_size=self.output_batch_size,
            print_write_progress=self.print_write_progress,
            limit=self.limit,
        )

        if hasattr(texts, "terminate"):
            logger.info("Killing all remaining workers, if any (iterator end)")
            texts.terminate()

        return saved_docs

    def save_texts_to_jsonl(self, texts: Iterable[str], append: bool = False):
        """
        Write JSONL files to <output_dir>/<DATASET_ID>.jsonl
        (each line is a JSON object with "doc" field and text as plain text)
        """
        mode = "a" if append else "w"
        fp = self.get_output_file_paths(single=True)[0]

        # Save as JSONL
        logger.info(f"Writing JSONL output to {fp} ({mode=})")

        docs_count = 0

        with open(fp, mode) as f:
            for docs_count, text in enumerate(self.filter_texts(texts), 1):
                f.write(json.dumps({self.output_text_field: text}, ensure_ascii=self.json_ensure_ascii) + "\n")

                if docs_count > 0 and (docs_count % self.print_write_progress) == 0:
                    logger.info(f"Writen {docs_count:,} docs ...")

                if self.limit > 0 and docs_count >= self.limit:
                    logger.warning(f"Limit reached ({docs_count:,} docs)")

                    if hasattr(texts, "terminate"):
                        logger.info("Killing all remaining workers, if any")
                        texts.terminate()
                    break

        if hasattr(texts, "terminate"):
            logger.info("Killing all remaining workers, if any (iterator end)")
            texts.terminate()

        return docs_count

    def get_hf_auth_token(self):
        if self.hf_auth_token:
            return self.hf_auth_token
        else:
            env_token = os.environ.get("HF_PASSWORD")

            if env_token:
                logger.info("Using HF auth token from env var")
                return env_token

        return None

    def get_local_dataset_dir(self):
        if self.LOCAL_DIRS:
            # manually defined dataset directory
            return get_path_by_system(self.LOCAL_DIRS)
        elif self.raw_datasets_dir:
            # automatically based on language + dataset_id
            return os.path.join(self.raw_datasets_dir, self.get_language_code(), self.DATASET_ID)
        else:
            raise ValueError("Either `LOCAL_DIRS` or `raw_datasets_dir` must be defined.")

    def get_dataset_file_paths(
        self,
        dataset_dir: Optional[str] = None,
        single_file: bool = False,
        subdirectories: bool = False,
        needed_suffix: Optional[Union[str, Tuple[str]]] = None,
        return_none_if_not_dir_exists: bool = False,
    ):
        if dataset_dir is None:
            dataset_dir = self.get_local_dataset_dir()

        if not os.path.exists(dataset_dir):
            logger.warning(f"Download directory does not exist: {dataset_dir}")

            if return_none_if_not_dir_exists:
                return None
            else:
                return []

        if subdirectories:
            # find files in all subdirectories
            logger.info(f"Finding dataset files in all subdirectories: {dataset_dir}")
            fps = [os.path.join(path, name) for path, subdirs, files in os.walk(dataset_dir) for name in files]

        else:
            # root-level files
            fps = [os.path.join(dataset_dir, f) for f in os.listdir(dataset_dir)]

        # filter by suffix
        fps = [f for f in fps if needed_suffix is None or f.endswith(needed_suffix)]

        # filter by file type
        fps = [fp for fp in fps if os.path.isfile(fp)]

        if single_file:
            if len(fps) != 1:
                raise ValueError(f"Multiple files in download directory but only a single one was expected: {fps}")

            return fps[0]

        return fps

    def decompress(self):
        raise NotImplementedError

    def is_dummy(self):
        return self.DUMMY

    def is_downloaded(self):
        return False

    def download(self):
        # Download all DOWNLOAD_URLS into local dataset dir
        output_dir = self.get_local_dataset_dir()

        logger.info(f"Downloading {len(self.DOWNLOAD_URLS)} files to {output_dir}")

        if not os.path.exists(output_dir):
            logger.info(f"Creating download dir: {output_dir}")
            os.makedirs(output_dir)

        for source_url in self.DOWNLOAD_URLS:
            if isinstance(source_url, tuple):
                source_url, target_filename = source_url
                output_filepath = os.path.join(output_dir, target_filename)

                if os.path.exists(output_filepath):
                    logger.warning(f"Output exists already: {output_filepath}")
                    continue
            else:
                output_filepath = output_dir  # auto file name

            try:
                logger.info(f"Download URL: {source_url}")
                logger.info(f"Output path: {output_filepath}")

                out_filename = wget.download(source_url, out=output_filepath)
                logger.info(f"Completed {out_filename}")
            except HTTPError as e:
                logger.error(f"Error {e}")

    def get_tokens(self):
        if self.TOKENS:
            return self.TOKENS
        elif self.get_bytes():
            return int(self.get_bytes() * TOKENS_PER_BYTE)
        else:
            return None

    def get_bytes(self):
        return self.BYTES

    def get_texts_from_conllu_file(self, file_handler: TextIO):
        import conllu

        text = None

        # try:
        for sentence in conllu.parse_incr(file_handler):
            if "newdoc id" in sentence.metadata:
                if text is not None:
                    # doc completed
                    yield text
                text = ""  # init empty document

            # append text to doc
            if "text" in sentence.metadata:
                if not text:
                    text = ""  # some conllu are not using doc ids -> force init
                else:
                    text += " "  # whitespace betweeen sentences

                text += sentence.metadata["text"]

            if "title" in sentence.metadata:
                text += self.title_delimiter

        # yield last document
        if text:
            yield text

        # except ParseException as e:
        #     # TODO
        #     logger.error(e)

    def get_texts(self) -> Iterable[str]:
        raise NotImplementedError

    def extract_plaintext(self):
        self.save_texts(self.get_texts())

        if self.counter:
            logger.info(f"Statistics {self.counter}")

src/lm_datasets/datasets/base.py
malteos/lm-datasets
import logging

from lm_datasets.datasets.base import BaseDataset, Availability, MB
from tqdm.auto import tqdm

logger = logging.getLogger(__name__)


class BGNCAdminEURDataset(BaseDataset):
    """
    Part of Bulgarian National Corpus

    TODO overlap with eurlex_bg?
    """

    DATASET_ID = "bgnc_admin_eur"
    TITLE = "ADMIN_EUR Corpus of EU legislation (bg)"
    HOMEPAGE = "https://eur-lex.europa.eu/homepage.html"

    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["bg"]

    DOWNLOAD_URLS = ["https://dcl.bas.bg/BulNC-registration/dl.php?dl=feeds/ADMIN_EUR.BG.zip"]

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/bg/BgNC/admin_eur/ADMIN_EUR.BG"]

    BYTES = 257 * MB

    def download(self):
        """
        DOWNLOAD
        -----------

        Instruction

        - Downloaded locally by clicking on the download link in the browser:

        https://dcl.bas.bg/BulNC-registration/dl.php?dl=feeds/ADMIN_EUR.BG.zip

        - Copy local file to server:

        scp /Local/Path/to/ADMIN_EUR.BG.zip username@clustername:/data/datasets/ele/bg/BgNC/admin_eur

        - Extract files:

        unzip ADMIN_EUR.BG.zip

        """
        pass

    def decompress(self):
        # unzip ADMIN_EUR.BG.zip
        pass

    def get_texts(self):
        files_path = self.get_dataset_file_paths(subdirectories=True, needed_suffix=".txt")

        logger.info(f"Found {len(files_path):,} files")

        for input_file in tqdm(files_path, desc="Reading files"):
            # skip if is metadata
            if "METADATA" in input_file:
                logger.warning(f"Skip {input_file}")
                continue

            # each file is one documentt
            with open(input_file, "r") as inp:
                text = inp.read()
                yield text.strip()

src/lm_datasets/datasets/bg/bgnc_admin_eur.py
malteos/lm-datasets
import logging

from lm_datasets.datasets.base import BaseDataset, Availability, MB
from tqdm.auto import tqdm

logger = logging.getLogger(__name__)


class BGNCNewsCorpusDataset(BaseDataset):
    """
    Part of Bulgarian National Corpus
    """

    DATASET_ID = "bgnc_news_corpus"
    TITLE = "News Corpus (bg)"
    HOMEPAGE = "https://dcl.bas.bg/BulNC-registration/?lang=EN#feeds/page/2"

    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["bg"]

    DOWNLOAD_URLS = ["https://dcl.bas.bg/BulNC-registration/dl.php?dl=feeds/JOURNALISM.BG.zip"]

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/bg/BgNC/news_corpus"]

    BYTES = 57 * MB

    def download(self):
        """
        DOWNLOAD
        -----------

        Instruction

        - Downloaded locally by clicking on the download link in the browser:

        https://dcl.bas.bg/BulNC-registration/dl.php?dl=feeds/JOURNALISM.BG.zip

        - Copy local file to server:

        scp /Local/Path/to/JOURNALISM.BG.zip username@clustername:/data/datasets/ele/bg/BgNC/news_corpus

        - Extract files:

        unzip JOURNALISM.BG.zip

        """
        pass

    def decompress(self):
        # unzip JOURNALISM.BG.zip
        pass

    def get_texts(self):
        files_path = self.get_dataset_file_paths(subdirectories=True, needed_suffix=".txt")

        logger.info(f"Found {len(files_path):,} files")

        for input_file in tqdm(files_path, desc="Reading files"):
            # skip if is metadata
            if "METADATA" in input_file:
                logger.warning(f"Skip {input_file}")
                continue

            # each file is one documentt
            with open(input_file, "r", encoding="utf-8") as inp:
                text = inp.read()
                yield text.strip()

src/lm_datasets/datasets/bg/bgnc_news_corpus.py
malteos/lm-datasets
from lm_datasets.datasets.base import MB, BaseDataset

import logging

logger = logging.getLogger(__name__)


class BulgarianNewsDataset(BaseDataset):
    DATASET_ID = "bulgarian_news"
    DOWNLOAD_URLS = ["http://old.dcl.bas.bg/dataset/Bulgarian_news.7z"]
    DESCRIPTION = (
        "The collection was collected by crawling Bulgarian websites in Bulgarian. Text samples are in json format. We"
        " can provide raw tests."
    )

    LANGUAGES = ["bg"]
    BYTES = 919 * MB

    def decompress(self):
        # 7z x Bulgarian_news.7z
        pass

    def get_texts(self):
        # from py7zr import SevenZipFile

        # read from extracted JSON files
        # fps = self.get_dataset_file_paths(subdirectories=True, needed_suffix=".json")

        # read from 7z
        # fp = self.get_dataset_file_paths(single_file=True, needed_suffix=".7z")
        # logger.info(f"Reading from {fp}")

        # with SevenZipFile(fp, "r") as zip:
        #     fn = "Bulgarian_news/n/nbox.bg/bg_00000634-24dd-4c46-b846-af4e5c3db559.json"
        #     print(fn)

        #     content = zip.read(fn)[fn].read()

        #     # for fname, bio in zip.read(targets=[fn]).items():
        #     #     print(f"{fname}: {bio.read(10)}...")
        #     #     break

        #     # allfiles = zip.getnames()

        # with SevenZipFile(fp, "r") as zip:
        #     for fn in allfiles:
        #         if fn.endswith(".json"):
        #             for fname, bio in zip.read(fn).items():
        #                 print(f"{fname}: {bio.read(10)}...")
        #         break

        # with lzma.open(fp) as f:
        #     file_content = f.read()

        print("x")

src/lm_datasets/datasets/bg/bulgarian_news.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, GB, Availability
import logging

logger = logging.getLogger(__name__)


class BulNCDataset(BaseDataset):
    DATASET_ID = "bulnc"
    TITLE = "Bulgarian National Corpus"
    AVAILIBILITY = Availability.ON_REQUEST
    DOWNLOAD_URLS = ["http://old.dcl.bas.bg/dataset/BulNC.7z"]
    DESCRIPTION = (
        "The Bulgarian National Corpus contains a wide range of texts in various sizes, media types (written and "
        "spoken), styles, periods (synchronic and diachronic), and licenses. Each text in the collection is supplied "
        "with metadata. The Bulgarian National Corpus  was first compiled using the Bulgarian Lexicographic Archive "
        "and the Text Archive of Written Bulgarian, which account for 55.95% of the corpus. Later, the EMEA corpus "
        "(medical administrative texts) and the OpenSubtitles corpus (film subtitles) were added, accounting for "
        "1.27% and 8.61% of the BulNC, respectively. The remaining texts were crawled automatically and include a "
        "large number of administrative texts, news from monolingual and multilingual sources, scientific texts, and "
        "popular science. The BulNC is not fully downloadable due to the inclusion of copyrighted material. We've "
        "provided a link to a password-protected archive for evaluation."
    )

    LANGUAGES = ["bg"]
    BYTES = 1.8 * GB

    def decompress(self):
        # 7z x BulNC.7z
        pass

    def get_texts(self):
        # # read from extracted JSON files
        # fps = self.get_dataset_file_paths(subdirectories=True, needed_suffix=".json")
        raise NotImplementedError()

src/lm_datasets/datasets/bg/bulnc.py
malteos/lm-datasets
from pathlib import Path
from lm_datasets.datasets.base import BaseDataset, Availability
from lm_datasets.datasets.hf_dataset import HFDataset


# files already downloaded => use parquet DS instead!
class StarcoderHFDataset(HFDataset):
    DATASET_ID = "starcoder"

    TITLE = "StarCoder"
    HOMEPAGE = "https://huggingface.co/datasets/bigcode/starcoderdata"
    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    LANGUAGES = ["code"]

    TOKENS = 250_000_000_000

    HF_DATASET_ID = "bigcode/starcoderdata"
    HF_DATASET_SPLIT = "train"

    text_column_name = "text"
    # remove_columns = ["doc_id", "LICENSE", "uri", "date_built"]


class ParquetDataset(BaseDataset):
    SINGLE_OUTPUT_FILE = False

    def extract_plaintext(self):
        if self.output_format == "parquet":
            raise ValueError("Dataset is already in parquet format; no text extraction needed!")
        else:
            super().extract_plaintext()


class StarcoderDataset(ParquetDataset):
    """
    Output files:
    /netscratch/ortiz/corpora/starcoder/starcoderdata/<code_lang>/train-<i>-of-<n>.parquet
    """

    DATASET_ID = "starcoder"

    TITLE = "StarCoder"
    HOMEPAGE = "https://huggingface.co/datasets/bigcode/starcoderdata"
    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    LANGUAGES = ["code"]

    TOKENS = 250_000_000_000

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/starcoder/starcoderdata"]

    HF_DATASET_ID = "bigcode/starcoderdata"
    HF_DATASET_SPLIT = "train"

    text_column_name = "text"
    # remove_columns = ["doc_id", "LICENSE", "uri", "date_built"]

    def get_output_file_paths(self):
        return Path(self.get_local_dataset_dir()).glob("*.parquet")

src/lm_datasets/datasets/code/starcoder.py
malteos/lm-datasets
import json
import logging
import argparse
import os
from typing import Iterable
from smart_open import open
import re
import gzip

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger(__name__)

# dataset metadata
DATASET_ID = "cs_en_parallel"
TITLE = "Czech-English Parallel Corpus 1.0 (CzEng 1.0)"
HOMEPAGE = "http://hdl.handle.net/11234/1-1458"

# language codes of the dataset (see ./resources/languages.json)
LANGUAGES = ["cs"]


"""
DOWNLOAD
-----------

Instruction

- Run the command on the server:

curl --remote-name-all \
    https://lindat.mff.cuni.cz/repository/xmlui/bitstream/handle/11234/1-1458/data-plaintext-format.tar

- Extract the archive:

tar -xvf data-plaintext-format.tar

"""


def file_iterator(data_path: str) -> Iterable[str]:
    """
    Iterates through all files in a given directory and outputs path names
    """

    for file in os.listdir(data_path):
        filename = os.path.join(data_path, file)
        yield filename


def text_extractor(files_path: Iterable[str]) -> Iterable[str]:
    """
    Extracts all texts from each file. Each file contains several texts, and one line per sentence.
    This function concatenates all sentences belonging to one text and then outputs the text.
    """
    for input_file in file_iterator(files_path):
        with gzip.open(input_file, "rt") as inp:
            lines = inp.readlines()
            cs_texts = []
            current_filename = ""
            for line in lines:
                rows = line.split("\t")
                name = rows[0]
                filename = re.sub("-s[0-9]+", "", name)
                if current_filename == "":
                    current_filename = filename
                elif current_filename != filename:
                    yield " ".join(cs_texts)
                    cs_texts = []
                    current_filename = filename
                # score = rows[1]
                cs_text = rows[2]
                # en_text = rows[3]
                cs_texts.append(cs_text.strip("\n"))
            yield " ".join(cs_texts)


def save_texts_to_jsonl(
    texts: Iterable[str],
    output_file_path: str,
    append: bool = False,
    limit: int = 0,
    min_length: int = 0,
    print_write_progress: int = 1_000,
):
    """
    Saves a list or generator of texts into a JSON-line file (each line is a valid JSON object)
    """
    mode = "a" if append else "w"
    output_text_field = "text"

    # Save as JSONL
    logger.info(f"Writing output to {output_file_path} ({mode=})")

    docs_count = 0

    with open(output_file_path, mode, encoding="utf-8") as f:
        for i, text in enumerate(texts):
            if min_length > 0 and len(text) < min_length:
                # skip because of short text length
                continue

            f.write(json.dumps({output_text_field: text}, ensure_ascii=False) + "\n")
            docs_count += 1

            if i > 0 and (i % print_write_progress) == 0:
                logger.info(f"Writen {i:,} lines ...")

            if limit > 0 and docs_count >= limit:
                logger.warning(f"Limit reached ({docs_count:,} docs)")
                break


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("input_path", help="Path to input file or directory")
    parser.add_argument(
        "output_path",
        help="Output is saved in file path (JSONL file format)",
    )
    parser.add_argument("--override", action="store_true", help="Override existing output files")
    parser.add_argument("--download", action="store_true", help="Download dataset")
    parser.add_argument("--limit", default=0, type=int, help="Limit dataset size (for debugging)")
    parser.add_argument("--min_text_length", default=0, type=int, help="Min. text length (shorter texts are discarded)")
    args = parser.parse_args()

    if os.path.exists(args.output_path) and not args.override:
        raise FileExistsError(f"Output exists already ({args.output_path}). Fix this error with --override")

    if not os.path.exists(args.input_path):
        raise FileNotFoundError(f"Input does not exist: {args.input_path}")

    # if args.download:
    #     download()

    texts = text_extractor(args.input_path)

    save_texts_to_jsonl(texts, args.output_path, limit=args.limit, min_length=args.min_text_length)

src/lm_datasets/datasets/cs/cs_en_parallel.py
malteos/lm-datasets
import logging

import lzma as xz

from lm_datasets.datasets.base import BaseDataset, Availability


logger = logging.getLogger(__name__)


class SynV9Dataset(BaseDataset):
    """ """

    DATASET_ID = "syn_v9"
    TITLE = "SYN v9: large corpus of written Czech"
    HOMEPAGE = "https://lindat.mff.cuni.cz/repository/xmlui/handle/11234/1-4635"
    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    LANGUAGES = ["cs"]

    LICENSE = "Academic Use"

    USED_BY = ["https://arxiv.org/abs/2103.13031"]

    TOKENS = 4_700_000_000

    def get_texts(self):
        filepath = self.get_dataset_file_paths(single_file=True, needed_suffix=".xz")

        logger.info(f"Reading from {filepath}")
        doc_counter = 0
        with xz.open(filepath, "rt", encoding="utf-8") as f_in:
            doc_lines = []

            # Read file line by line
            for i, line in enumerate(f_in):
                doc_lines.append(line)

                # Extract <doc> ... </doc>
                if line.startswith("<doc "):
                    doc_lines = [line]
                elif line.startswith("</doc>"):
                    # End of doc -> process
                    doc_counter += 1
                    doc_text = ""

                    # Parse a full <doc> item
                    sent = ""
                    for doc_line in doc_lines:
                        if doc_line.startswith("</block"):
                            # Each block is a paragraph -> paragraph completed
                            yield doc_text.rstrip()  # remove last white space
                            doc_text = ""
                        elif doc_line.startswith("<s "):
                            # beginning of sentence
                            sent = ""
                        elif doc_line.startswith("</s>"):
                            # parse sentence
                            s = ""
                            last_line = "<g/>"
                            for sent_line in sent.splitlines():
                                sent_line = sent_line.strip()

                                if sent_line.startswith("<"):
                                    pass
                                else:
                                    if last_line != "<g/>":
                                        s += " "

                                    cols = sent_line.split("\t", 2)

                                    if len(cols) > 1:  # is valid entry (CONLLU-like format)
                                        original = cols[0]  # first column contains original string
                                        s += original

                                last_line = sent_line

                            # sentence parsed
                            # WARNING: do not concatenate sentences to a document since
                            # they have been shuffled - only paragraph level.
                            doc_text += s + " "

                        else:
                            sent += doc_line

        logger.info(f"Done. Processed {doc_counter:,} shuffled documents (and more paragraphs)")

src/lm_datasets/datasets/cs/syn_v9.py
malteos/lm-datasets
import json
import logging
from typing import Iterable
from smart_open import open

from lm_datasets.datasets.base import GB, BaseDataset, Availability

logger = logging.getLogger(__name__)


class DANewsroomDataset(BaseDataset):
    DATASET_ID = "danewsroom"
    TITLE = "DaNewsroom: A Large-scale Danish Summarisation Dataset"
    HOMEPAGE = "https://github.com/danielvarab/da-newsroom"  # "https://aclanthology.org/2020.lrec-1.831/"

    LANGUAGES = ["da"]
    # LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/danewsroom"]

    AVAILIBILITY = Availability.ON_REQUEST
    BYTES = 1.5 * GB

    def download(self):
        """
        gdown "https://drive.google.com/u/0/uc?id=1u22Hcs__CUu_GAzY4HQ275gfHvLGOckv&export=download"
        """
        pass

    def get_texts(self) -> Iterable[str]:
        """
        Extracts the text from each JSONL file.
        """

        with open(self.get_dataset_file_paths(single_file=True, needed_suffix=".jsonl.gz")) as f:
            for line in f:
                doc = json.loads(line)
                text = doc["title"] + self.title_delimiter + doc["text"]

                yield text

src/lm_datasets/datasets/da/danewsroom.py
malteos/lm-datasets
import json
import os
from lm_datasets.datasets.base import BILLION, Availability, BaseDataset
from lm_datasets.datasets.hf_dataset import HFDataset


class DanishGigawordDataset(HFDataset):
    DATASET_ID = "danish_gigaword"

    TITLE = "Danish GiagaWord"
    HOMEPAGE = "https://sprogteknologi.dk/dataset/danish-gigaword"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["da"]

    DESCRIPTION = (
        "A billion-word corpus of Danish text. Split into many sections, and covering many dimensions ",
        "of variation (spoken/written, formal/informal, modern/old, rigsdansk/dialect, and so on).",
        "",
        "The license is CC-BY 4.0, Creative Commons with Attribution. Owners: ITU; Leon Derczynski, Manuel R. Ciosici",
    )
    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "public domain"

    # Size: 1 billion words
    TOKENS = 1 * BILLION

    HF_DATASET_ID = "DDSC/partial-danish-gigaword-no-twitter"
    HF_DATASET_SPLIT = "train"

    excluded_sources = {
        # exclude all Wikimedia sources to avoid duplicated content with the original Wikimedia dataset
        "wiki",
        "wikibooks",
        "wikisource",
    }

    text_column_name = "text"
    remove_columns = ["doc_id", "LICENSE", "uri", "date_built"]

    def get_filter_func(self):
        def filter_func(example):
            return example["source"] not in self.excluded_sources

        return filter_func


# deprecated
class DanishGigawordOriginalDataset(BaseDataset):
    """
    Dataset based on orignal source (not HF version)
    """

    DATASET_ID = "danish_gigaword"
    TITLE = "Danish GiagaWord"
    HOMEPAGE = "https://sprogteknologi.dk/dataset/danish-gigaword"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["da"]

    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "public domain"

    # Size: 1 billion words
    TOKENS = 1 * BILLION

    DOWNLOAD_URLS = ["https://itu.dk/research/dagw/dagw_v1.0-release.zip"]
    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/danish_gigaword"]

    def get_dataset_file_path(self):
        # TODO handle datasets with multiple filles
        dataset_dir = self.get_local_dataset_dir()
        files = os.listdir(dataset_dir)

        if len(files) == 1:
            return os.path.join(dataset_dir, files[0])
        else:
            raise ValueError(f"Cannot determine file path. Either none or multiple files in dataset dir: {files=}")

    def get_texts(self):
        import zipfile

        dataset_file_path = self.get_dataset_file_path()
        print(dataset_file_path)

        with zipfile.ZipFile(dataset_file_path, "r") as zf:
            # zf.printdir()
            zfns = zf.namelist()

            for fn in zfns:
                if "sektioner" in fn:
                    print(fn)

                    with zf.open(fn) as content_file:
                        for json_str in content_file:
                            doc = json.loads(json_str)
                            print(doc)

                            yield doc["text"]

                        # content = content_file.read()

                        # for l in myfile:
                        #     print(l)
                    # break

src/lm_datasets/datasets/da/danish_gigaword.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, Availability


# @DeprecationWarning("Folketinget is already part of Danish GigaWord")
class DanishParliamentCorpusDataset(BaseDataset):
    DATASET_ID = "danish_parliament_corpus"
    TITLE = "The Danish Parliament Corpus 2009 - 2017, v1"
    HOMEPAGE = " http://hdl.handle.net/20.500.12115/8"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["da"]

    DESCRIPTION = (
        "The Danish Parliament Corpus 2009 - 2017 contains Hansards (transcripts of parliamentary speeches) from the "
        "sittings in the Chamber of the Danish Parliament, Folketinget. The corpus consists of xml files, one for "
        "each parliamentary year, running from October to the following June. The files are marked for meetings, "
        "item title and number, speeches, name and party of speakers, date, time etc. The Danish Parliament Corpus "
        "2009-2017 follows the license for Open Data stating the following: `The Danish Parliament grants a "
        "world-wide, free, non-exclusive and otherwise unrestricted right of use of the data in the Danish "
        "Parliament's open data catalogue. The data can be freely: § copied, distributed and published, § "
        "adapted and combined with other material, § exploited commercially and non-commercially. ` Following "
        "the copyright act, the speeches can be distributed without the consent of the speaker but only in a way "
        "where the author/speaker of each text/speech is clearly stated. Furthermore, the Danish Parliament must "
        "be acknowledged as the source. Version 1 of the corpus includes meetings until May 4th, 2017, and the "
        "reports for the latest parliamentary year have not been published as the final edition. The reports of "
        "all other meetings are the final editions. "
        "Project Leader: Costanza Navarretta, UCPH"
    )
    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "public domain"

    # Size:  40,841,226 words
    TOKENS = 40_841_226

    DOWNLOAD_URLS = []
    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/danish_parliament_corpus"]

src/lm_datasets/datasets/da/danish_parliament_corpus.py
malteos/lm-datasets
import logging
from typing import Iterable

import zipfile


from lm_datasets.datasets.base import BaseDataset, Availability, GB

logger = logging.getLogger(__name__)


class DKClarinDataset(BaseDataset):
    DATASET_ID = "dk_clarin"
    TITLE = "DK-CLARIN Reference Corpus of General Danish"
    HOMEPAGE = "https://korpus.dsl.dk/clarin/"

    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD
    LICENSE = "Academic Use; CLARIN-ACA-NC"

    LANGUAGES = ["da"]

    BYTES = 1.4 * GB

    def get_texts(self) -> Iterable[str]:
        """
        Extracts the text from the zip files
        """
        from bs4 import BeautifulSoup

        zip_fps = self.get_dataset_file_paths(needed_suffix=".zip")

        for zip_fp in zip_fps:
            logger.info(f"Extracting from {zip_fp}")

            with zipfile.ZipFile(zip_fp) as zf:
                needed_suffix = ".zip"
                members_fns = [fn for fn in zf.namelist() if fn.endswith(needed_suffix)]

                for member_fn in members_fns:
                    logger.info(f"---> Extracting from {member_fn}")

                    if "wikipedia" in member_fn:
                        logger.info(f"---> Skipping {member_fn} (exclude Wikipedia)")
                        continue

                    with zf.open(member_fn) as member_f:
                        with zipfile.ZipFile(member_f) as zf_2:
                            needed_suffix_2 = ".xml"
                            members_fns_2 = [fn for fn in zf_2.namelist() if fn.endswith(needed_suffix_2)]

                            for member_fn_2 in members_fns_2:
                                if "__MACOSX" in member_fn_2:
                                    logger.info(f"Skipping {member_fn_2}")
                                    continue

                                logger.info(f"Parsing from {member_fn_2}")

                                # TEIP5DKCLARIN-format
                                # p => paragraph
                                # w => word
                                # p => punctutation
                                # s => space
                                # type="s" => s

                                with zf_2.open(member_fn_2) as member_f_2:
                                    soup = BeautifulSoup(member_f_2, features="lxml")

                                paragraphs_texts = []
                                paragraphs = soup.select("p")

                                for i, p in enumerate(paragraphs):
                                    paragraph_text = ""
                                    for ele in p.find_all():
                                        paragraph_text += ele.get_text()
                                        attrs = ele.attrs

                                        if "type" in attrs and attrs["type"] == "s":
                                            paragraph_text += " "

                                    paragraphs_texts.append(paragraph_text)

                                title = soup.title.get_text().replace("CTB version: ", "")  # TODO strip "CTB version: "

                                if title == "nil":
                                    text = ""
                                else:
                                    text = title + self.title_delimiter

                                text += self.paragraph_delimiter.join(paragraphs_texts)

                                yield text

        logger.info("All ZIP files extracted")

src/lm_datasets/datasets/da/dk_clarin.py
malteos/lm-datasets
from typing import List, Optional, Union
from .multilingual.wikimedia import get_wikimedia_auto_classes
from .multilingual.oscar_dummy import get_oscar_dummy_classes
from .multilingual.colossal_oscar import get_colossal_oscar_auto_classes
from .multilingual.eurlex import get_eurlex_auto_classes
from .multilingual.legal_mc4 import get_legal_mc4_auto_classes
from .nl.sonar import get_sonar_classes

import importlib
import logging


logger = logging.getLogger(__name__)


ALL_DATASET_IMPORTS = [
    # multilingual
    # curlicat
    ".multilingual.curlicat.CurlicatBGDataset",
    ".multilingual.curlicat.CurlicatHRDataset",
    ".multilingual.curlicat.CurlicatHUDataset",
    ".multilingual.curlicat.CurlicatPLDataset",
    ".multilingual.curlicat.CurlicatRODataset",
    ".multilingual.curlicat.CurlicatSKDataset",
    ".multilingual.curlicat.CurlicatSLDataset",
    # macocu
    ".multilingual.macocu.MacocuBGDataset",
    ".multilingual.macocu.MacocuHRDataset",
    ".multilingual.macocu.MacocuELDataset",
    # ".multilingual.macocu.MacocuSQDataset",
    # ".multilingual.macocu.MacocuBSDataset",
    ".multilingual.macocu.MacocuCADataset",
    # ".multilingual.macocu.MacocuISDataset",
    # ".multilingual.macocu.MacocuMKDataset",
    ".multilingual.macocu.MacocuMTDataset",
    # ".multilingual.macocu.MacocuCNRDataset",
    ".multilingual.macocu.MacocuSRDataset",
    ".multilingual.macocu.MacocuSLDataset",
    # ".multilingual.macocu.MacocuTRDataset",
    ".multilingual.macocu.MacocuUKDataset",
    # redpajama
    ".multilingual.redpajama.RedPajamaBookDataset",
    # ".multilingual.redpajama.RedPajamaArxivDataset",
    ".multilingual.redpajama.RedPajamaStackexchangeDataset",
    # en
    ".en.wikihow.WikihowDataset",
    ".en.pes2o.PeS2oDataset",
    ".en.proof_pile.ProofPileDataset",
    ".en.dialogstudio.DialogstudioDataset",
    # bg
    ".bg.bgnc_admin_eur.BGNCAdminEURDataset",
    ".bg.bgnc_news_corpus.BGNCNewsCorpusDataset",
    ".bg.bulgarian_news.BulgarianNewsDataset",
    ".bg.bulnc.BulNCDataset",
    # de
    ".de.openlegaldata.OpenLegalDataDataset",
    ".de.dewac.DEWacDataset",
    # ga
    ".ga.ga_bilingual_legistation.GABilingualLegislationDataset",
    ".ga.ga_universal_dependencies.GAUniversalDependenciesDataset",
    # hr
    ".hr.hrwac.HRWACDataset",
    ".hr.styria_news.StyriaNewsDataset",
    ".hr.croatian_news_engri.CroatianNewsENGRIDataset",
    # it
    ".it.itwac.ITWacDataset",
    # mt
    ".mt.korpus_malti.KorpusMaltiDataset",
    # nl
    *get_sonar_classes(),
    ".nl.sonar_new_media.SonarNewMediaDataset",
    # sl
    ".sl.cc_gigafida.CCGigaFidaDataset",
    ".sl.academic_slovene_kas.AcademicSloveneKASDataset",
    ".sl.slwac_web.SLWaCWebDataset",
    # sk
    ".sk.sk_court_decisions.SKCourtDecisionsDataset",
    ".sk.sk_laws.SKLawsDataset",
    # cs
    ".cs.syn_v9.SynV9Dataset",
    # da
    ".da.danish_gigaword.DanishGigawordDataset",
    # danish_parliament_corpus.DanishParliamentCorpusDataset,
    ".da.danewsroom.DANewsroomDataset",
    ".da.dk_clarin.DKClarinDataset",
    # fr
    ".fr.cabernet.CaBeRnetDataset",
    # no
    ".no.norwegian_cc.NorwegianCCNNDataset",
    ".no.norwegian_cc.NorwegianCCNODataset",
    # nak.NAKDataset,
    # nbdigital.NBDigitalDataset,
    # maalfrid_2021.Maalfrid2021Dataset,
    # parlamint.ParlaMintDataset,
    # parliamentary_proceedings.ParliamentaryProceedingsDataset,
    # sakspapir_nno.SakspapirNNODataset,
    # pl
    # ".pl.luna_pl.LunaPL",  ### BAD data
    ".pl.pl_nkjp.NKJPPodkorpusMilionowyDataset",
    ".pl.pl_parliamentary_corpus.PLParliamentaryCorpusDataset",
    # pt
    ".pt.parlamento_pt.ParlamentoPtDataset",
    ".pt.brwac.BrWacDataset",
    # lt
    ".lt.seimas_lt_en.SeismasLTENDataset",
    # lv
    ".lv.state_related_latvian_web.StateRelatedLatvianWebDataset",
    # el
    ".el.greek_legal_code.GreekLegalCodeDataset",
    ".el.greek_web_corpus.GreekWebCorpus",
    # et
    ".et.estonian_reference_corpus.EstonianReferenceCorpusDataset",
    ".et.enc.ENC2021Dataset",
    # eu
    ".eu.euscrawl.EUSCrawlDataset",
    # es
    ".es.spanish_legal.SpanishLegalDataset",
    # ".es.escorpius.ESCorpiusDataset",  # based on CC => use OSCAR instead
    # fi
    ".fi.ylenews.YLENewsDataset",
    # sv
    ".sv.sv_gigaword.SVGigawordDataset",
    # sr
    ".sr.srpkor.SrpKorDataset",
    # ro
    ".ro.marcell_legislative_subcorpus_v2.MarcellLegislativeSubcorpusV2Dataset",
    # uk
    ".uk.uk_laws.UKLawsDataset",
    # code
    ".code.starcoder.StarcoderDataset",
]


def get_class_by_import_string(
    import_string_or_cls: Union[str, object], relative_base_package: str = "lm_datasets.datasets"
):
    """
    Import dataset class based on import string

    Allowed formats:
    - ".lang.DatasetClass"   # relative
    - "lm_datasets.lang.DatasetClass"  # absolute
    """

    if isinstance(import_string_or_cls, str):
        # import from string
        if import_string_or_cls.startswith("."):  # relative import
            package = relative_base_package
        else:
            package = None  # absolute

        import_split = import_string_or_cls.split(".")
        cls_name = import_split[-1]
        mod_str = ".".join(import_split[:-1])

        mod = importlib.import_module(mod_str, package=package)

        return getattr(mod, cls_name)
    else:
        return import_string_or_cls  # already object, no need to import from string



def get_registered_dataset_classes(extra_dataset_registries: Optional[Union[str, List[str]]] = None):
    dataset_classes = (
        [get_class_by_import_string(clss) for clss in ALL_DATASET_IMPORTS]
        + get_oscar_dummy_classes()
        + get_eurlex_auto_classes()
        + get_legal_mc4_auto_classes()
        + get_wikimedia_auto_classes()
        + get_colossal_oscar_auto_classes()
    )

    if extra_dataset_registries:
        if isinstance(extra_dataset_registries, str):
            extra_dataset_registries = extra_dataset_registries.split(",")

        # Iterate over registeries
        for extra_dataset_registry_str in extra_dataset_registries:
            logger.info(f"Loading datasets from registry: {extra_dataset_registry_str}")
            extra_dataset_registry_package = importlib.import_module(extra_dataset_registry_str)

            extra_dataset_registry_getter = getattr(extra_dataset_registry_package, "get_registered_dataset_classes")
            extra_dataset_classes = extra_dataset_registry_getter()

            logger.info(f"Extra datasets: {extra_dataset_classes}")

            dataset_classes += extra_dataset_classes

    return dataset_classes


src/lm_datasets/datasets/dataset_registry.py
malteos/lm-datasets
import logging
from typing import Iterable
from lm_datasets.datasets.base import BILLION, Availability, BaseDataset
from smart_open import open

logger = logging.getLogger(__name__)


class DEWacDataset(BaseDataset):
    """
    See also: ITWacDataset
    """

    DATASET_ID = "dewac"
    TITLE = "DeWaC"
    DESCRIPTION = (
        "DeWaC is a 1.7 billion word corpus constructed from the Web limiting the crawl to the .de domain and using"
        " medium-frequency words from the SudDeutsche Zeitung corpus and basic German vocabulary lists as seeds."
    )
    LANGUAGES = ["de"]
    AVAILIBILITY = Availability.ON_REQUEST
    HOMEPAGE = "https://docs.sslmit.unibo.it/doku.php?id=corpora:dewac"

    TOKENS = 1.7 * BILLION

    def decompress(self):
        #  7z x dewac_preproc.7z
        pass

    def get_texts(self) -> Iterable[str]:
        file_path = self.get_dataset_file_paths(single_file=True, needed_suffix=".txt")

        # chardet.detect(line)

        # New documents are separated by "CURRENT URL <url>"
        doc_text = ""
        with open(file_path, "r", encoding="ISO-8859-1") as f:
            for i, line in enumerate(f):
                if line.startswith("CURRENT URL "):
                    if doc_text:
                        # print(doc_text)
                        yield doc_text

                    doc_text = ""
                else:
                    doc_text += line.strip()

                pass

            if doc_text:
                # last document
                yield doc_text

src/lm_datasets/datasets/de/dewac.py
malteos/lm-datasets
from lm_datasets.datasets.base import Availability, Genre
from lm_datasets.datasets.jsonl_dataset import JSONLDataset


class OpenLegalDataDataset(JSONLDataset):
    DATASET_ID = "openlegaldata"
    TITLE = "Open Legal Data - German court decisions and laws"
    HOMEPAGE = "https://openlegaldata.io/"
    VERSION = "oldp_cases_v1_v2_fix"  # "oldp-cases_merged_v1_v2"  # dump date: 2022-10-18

    LANGUAGES = ["de"]
    GENRES = [Genre.LEGAL]
    AVAILIBILITY = Availability.ON_REQUEST

    TOKENS = 9_700_000_000

    raw_jsonl_paths = [
        "v1.jsonl.gz",
        "v2.jsonl.gz",
        # "merged_v1_v2.jsonl.gz",
        # "merged_v1_v2.jsonl",
    ]

    doc_ids = set()

    def get_text_from_item(self, item):
        # filter by doc ID
        if item["doc_id"] in self.doc_ids:
            return None
        else:
            self.doc_ids.add(item["doc_id"])
            return item[self.raw_jsonl_text_field]

src/lm_datasets/datasets/de/openlegaldata.py
malteos/lm-datasets
from lm_datasets.datasets.base import Availability, MB
from lm_datasets.datasets.hf_dataset import HFDataset


_CITATION = """
@inproceedings{papaloukas-etal-2021-glc,
    title = "Multi-granular Legal Topic Classification on Greek Legislation",
    author = "Papaloukas, Christos and Chalkidis, Ilias and Athinaios, Konstantinos and Pantazi, Despina-Athanasia and Koubarakis, Manolis",
    booktitle = "Proceedings of the Natural Legal Language Processing Workshop 2021",
    year = "2021",
    address = "Punta Cana, Dominican Republic",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/2109.15298",
    doi = "10.48550/arXiv.2109.15298",
    pages = "63--75"
}
"""  # noqa


class GreekLegalCodeDataset(HFDataset):
    DATASET_ID = "greek_legal_code"

    TITLE = "Greek Legal Code"
    DESCRIPTION = (
        "Greek_Legal_Code (GLC) is a dataset consisting of approx. 47k legal resources from Greek legislation. The"
        " origin of GLC is “Permanent Greek Legislation Code - Raptarchis”, a collection of Greek legislative "
        " documents classified into multi-level (from broader to more specialized) categories."
    )

    HOMEPAGE = "https://huggingface.co/datasets/greek_legal_code"

    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["el"]

    HF_DATASET_ID = "greek_legal_code"
    HF_DATASET_SPLIT = "train"
    HF_DATASET_CONFIGS = [  # all the same text but different labels
        "chapter",
        # "subject",
        # "volume",
    ]

    BYTES = (438 / 3) * MB

    text_column_name = "text"
    remove_columns = ["label"]

src/lm_datasets/datasets/el/greek_legal_code.py
malteos/lm-datasets
from typing import Iterable
from lm_datasets.datasets.base import GB, BaseDataset, Availability, QualityWarning
from smart_open import open
import logging

logger = logging.getLogger(__name__)


class GreekWebCorpus(BaseDataset):
    DATASET_ID = "greek_web_corpus"
    TITLE = "Greek Web Corpus"
    AVAILIBILITY = Availability.ON_REQUEST
    HOMEPAGE = "http://nlp.polytechnique.fr/resources-greek"
    CITATION = """@article{Outsios2018,
        title = {Word Embeddings from Large-Scale Greek Web content},
        author = {Outsios, Stamatis and Skianis, Konstantinos and Meladianos, Polykarpos and Xypolopoulos, Christos and Vazirgiannis, Michalis},
         year={2018},
         journal = {arXiv preprint arXiv:1810.06694},
    }"""  # noqa
    LANGUAGES = ["el"]
    USED_BY = [
        "https://arxiv.org/abs/2304.00869",  # GreekBART
    ]

    BYTES = 10 * GB

    QUALITY_WARNINGS = [QualityWarning.SHORT_TEXT, QualityWarning.BAD_PUNCTUATION]

    def get_texts(self) -> Iterable[str]:
        # txt.gz
        fp = self.get_dataset_file_paths(needed_suffix=".txt.gz", single_file=True)
        logger.info(f"Reading from {fp}")

        with open(fp) as f:
            for i, line in enumerate(f):
                text = line.strip()  # TODO each line a new doc?

                if text:
                    yield text

src/lm_datasets/datasets/el/greek_web_corpus.py
malteos/lm-datasets
from typing import Iterable
from lm_datasets.datasets.base import Genre
from lm_datasets.datasets.hf_dataset import HFDataset


DIALOGSTUDIO_DATASETS = {
    "natural_language_understanding": [
        "ATIS",
        "ATIS-NER",
        "BANKING77",
        "BANKING77-OOS",
        "CLINC-Single-Domain-OOS-banking",
        "CLINC-Single-Domain-OOS-credit_cards",
        "CLINC150",
        "DSTC8-SGD",
        "HWU64",
        "MIT-Movie",
        "MIT-Restaurant",
        "RESTAURANTS8K",
        "SNIPS",
        "SNIPS-NER",
        "TOP",
        "TOP-NER",
    ],
    "task_oriented": [
        "ABCD",
        "AirDialogue",
        "BiTOD",
        "CaSiNo",
        "CraigslistBargains",
        "Disambiguation",
        "DSTC2-Clean",
        "FRAMES",
        "GECOR",
        "HDSA-Dialog",
        "KETOD",
        "KVRET",
        "MetaLWOZ",
        "MS-DC",
        "MuDoCo",
        "MulDoGO",
        "MultiWOZ_2.1",
        "MULTIWOZ2_2",
        "SGD",
        "SimJointGEN",
        "SimJointMovie",
        "SimJointRestaurant",
        "STAR",
        "Taskmaster1",
        "Taskmaster2",
        "Taskmaster3",
        "WOZ2_0",
    ],
    "dialogue_summarization": [
        "AMI",
        "CRD3",
        "DialogSum",
        "ECTSum",
        "ICSI",
        "MediaSum",
        "QMSum",
        "SAMSum",
        "TweetSumm",
        "ConvoSumm",
        "SummScreen_ForeverDreaming",
        "SummScreen_TVMegaSite",
    ],
    "conversational_recommendation": [
        "Redial",
        "DuRecDial-2.0",
        "OpenDialKG",
        "SalesBot",
    ],
    "open_domain": [
        "chitchat-dataset",
        "ConvAI2",
        "AntiScam",
        "Empathetic",
        "HH-RLHF",
        "PLACES3.5",
        "Prosocial",
        "SODA",
    ],
    "knowledge_grounded": [
        "CompWebQ",
        "CoQA",
        "CoSQL",
        "DART",
        "FeTaQA",
        "GrailQA",
        "HybridQA",
        "MTOP",
        "MultiModalQA",
        "SParC",
        "Spider",
        "SQA",
        "ToTTo",
        "WebQSP",
        "WikiSQL",
        "WikiTQ",
        "wizard_of_internet",
        "wizard_of_wikipedia",
    ],
}


class DialogstudioDataset(HFDataset):
    """
    Rather a fine-tuning dataset
    """

    DATASET_ID = "dialogstudio"
    TITLE = "DialogStudio: Unified Dialog Datasets and Instruction-Aware Models for Conversational AI"
    DESCRIPTION = (
        "The proof-pile is a 13GB pre-training dataset of mathematical text that comprises 8.3 billion tokens (using"
        " the gpt-neox tokenizer). Models trained on this dataset are coming soon :) The dataset is composed of diverse"
        " sources of both informal and formal mathematics, namely"
    )
    HOMEPAGE = "https://huggingface.co/datasets/Salesforce/dialogstudio"
    DOI = "10.48550/arXiv.2307.10172"

    LANGUAGES = ["en"]
    GENRES = [Genre.DIALOGUE]
    LICENSE = "mixed"

    HF_DATASET_ID = "Salesforce/dialogstudio"
    HF_DATASET_SPLIT = "train"

    text_column = "text"

    @property
    def HF_DATASET_CONFIGS(self):
        return [ds for datasets in DIALOGSTUDIO_DATASETS.values() for ds in datasets]

    def get_texts(self) -> Iterable[str]:
        # See https://github.com/salesforce/DialogStudio/tree/main
        raise NotImplementedError("Structured data needs to be converted into plaintext")

src/lm_datasets/datasets/en/dialogstudio.py
malteos/lm-datasets
from lm_datasets.datasets.base import BILLION, Genre
from lm_datasets.datasets.hf_dataset import HFDataset


class PeS2oDataset(HFDataset):
    """
    peS2o V2

    Knowledge cutoff: 2023-01-03
    Number of documents: 38.97M
    Number of whitespace-separated tokens*: 42.01B

    """

    DATASET_ID = "pes2o"
    TITLE = "peS2o"
    DESCRIPTION = (
        "The peS2o dataset is a collection of ~40M creative open-access academic papers, cleaned, filtered, and"
        " formatted for pre-training of language models. It is derived from the Semantic Scholar Open Research"
        " Corpus(Lo et al, 2020), or S2ORC."
    )
    HOMEPAGE = "https://huggingface.co/datasets/allenai/peS2o"
    VERSION = "V2"

    LANGUAGES = ["en"]
    GENRES = [Genre.SCIENCE]

    TOKENS = 42.01 * BILLION

    HF_DATASET_ID = "allenai/peS2o"
    HF_DATASET_SPLIT = "train"
    HF_DATASET_CONFIGS = ["v2"]

    text_column = "text"
    streaming = True

src/lm_datasets/datasets/en/pes2o.py
malteos/lm-datasets
from lm_datasets.datasets.base import BILLION, Genre
from lm_datasets.datasets.hf_dataset import HFDataset


class ProofPileDataset(HFDataset):
    """
    Duplicated content / overlap with arxiv: PeS2oDataset
    """

    DATASET_ID = "proof_pile"
    TITLE = "proof-pile"
    DESCRIPTION = (
        "The proof-pile is a 13GB pre-training dataset of mathematical text that comprises 8.3 billion tokens (using"
        " the gpt-neox tokenizer). Models trained on this dataset are coming soon :) The dataset is composed of diverse"
        " sources of both informal and formal mathematics, namely"
    )
    HOMEPAGE = "https://huggingface.co/datasets/hoskinson-center/proof-pile"

    LANGUAGES = ["en"]
    GENRES = [Genre.SCIENCE, Genre.MATH]

    TOKENS = 8.3 * BILLION

    HF_DATASET_ID = "hoskinson-center/proof-pile"
    HF_DATASET_SPLIT = "train"

    text_column_name = "text"

src/lm_datasets/datasets/en/proof_pile.py
malteos/lm-datasets
from lm_datasets.datasets.base import Availability, MB
from lm_datasets.datasets.hf_dataset import HFDataset


class WikihowDataset(HFDataset):
    DATASET_ID = "wikihow"

    TITLE = "WikiHow"
    DESCRIPTION = (
        "WikiHow is a new large-scale dataset using the online WikiHow (http://www.wikihow.com/) knowledge base."
    )

    HOMEPAGE = "https://github.com/mahnazkoupaee/WikiHow-Dataset"

    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["en"]

    DESCRIPTION = """WikiHow is a new large-scale dataset using the online WikiHow"""

    HF_DATASET_ID = "wikihow"
    HF_DATASET_CONFIGS = ["all"]
    HF_DATASET_SPLIT = "train"

    BYTES = 5.21 * MB

    @property
    def HF_DATA_DIR(self):
        return self.get_local_dataset_dir()

    text_column_name = "text"
    # remove_columns = ["doc_id", "LICENSE", "uri", "date_built"]

src/lm_datasets/datasets/en/wikihow.py
malteos/lm-datasets
from lm_datasets.datasets.hf_dataset import HFDataset


class ESCorpiusDataset(HFDataset):
    """
    OVERLAP with CC/OSCAR
    """

    DATASET_ID = "escorpius"
    TITLE = "esCorpius: A Massive Spanish Crawling Corpus "

    HOMEPAGE = "https://huggingface.co/datasets/LHF/escorpius"
    CITATION = """@inproceedings{gutierrezfandino22_iberspeech,
        author={Asier Gutiérrez-Fandiño and David Pérez-Fernández and Jordi Armengol-Estapé and David Griol and Zoraida Callejas},
        title={{esCorpius: A Massive Spanish Crawling Corpus}},
        year=2022,
        booktitle={Proc. IberSPEECH 2022},
        pages={126--130},
        doi={10.21437/IberSPEECH.2022-26}
        }"""  # noqa

    LANGUAGES = ["es"]
    WEB_CRAWLED = True

    HF_DATASET_ID = "LHF/escorpius"
    HF_DATASET_SPLIT = "train"

    text_column = "text"

    def get_texts(self):
        raise ValueError("Do not use -> overlap with OSCAR")

src/lm_datasets/datasets/es/escorpius.py
malteos/lm-datasets
import logging
import zipfile
from lm_datasets.datasets.base import Availability, BaseDataset, Genre, GB, QualityWarning

logger = logging.getLogger(__name__)


class SpanishLegalDataset(BaseDataset):
    """
    Warning: Document splitting only done by "\n\n" => short texts and bad line breaks.
    """

    DATASET_ID = "spanish_legal"
    TITLE = "Spanish Legal Domain Corpora"
    HOMEPAGE = "https://github.com/PlanTL-GOB-ES/lm-legal-es"
    CITATION = """@misc{gutierrezfandino2021legal,
            title={Spanish Legalese Language Model and Corpora},
            author={Asier Gutiérrez-Fandiño and Jordi Armengol-Estapé and Aitor Gonzalez-Agirre and Marta Villegas},
            year={2021},
            eprint={2110.12201},
            archivePrefix={arXiv},
            primaryClass={cs.CL}
        }
        """
    LANGUAGES = ["es"]
    GENRES = [Genre.LEGAL]
    AVAILIBITY = Availability.DIRECT_DOWNLOAD
    QUALITY_WARNINGS = [QualityWarning.BAD_LINEBREAKS, QualityWarning.SHORT_TEXT]
    BYTES = 8.9 * GB

    USED_BY = ["https://huggingface.co/PlanTL-GOB-ES/RoBERTalex"]

    DOWNLOAD_URLS = [("https://zenodo.org/record/5495529/files/corpus.zip?download=1", "corpus.zip")]

    def is_downloaded(self):
        return bool(self.get_dataset_file_paths(needed_suffix=".zip", single_file=True))

    def get_texts(self):
        if not self.is_downloaded():
            self.download()

        zip_fp = self.get_dataset_file_paths(needed_suffix=".zip", single_file=True)

        with zipfile.ZipFile(zip_fp) as zf:
            member_fns = zf.namelist()
            for fn in member_fns:
                # Exclude: EURLEX (part of other dataset)
                # TODO "opus", "europarl"
                if fn.endswith(".txt") and "eurlex" not in fn:
                    with zf.open(fn) as member_f:
                        txt = member_f.read().decode()

                        for doc_text in txt.split("\n\n"):  # TODO
                            doc_text = doc_text.strip()

                            yield doc_text

src/lm_datasets/datasets/es/spanish_legal.py
malteos/lm-datasets
import json
import logging
import argparse
import os
from typing import Iterable
from smart_open import open

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger(__name__)

# dataset metadata
DATASET_ID = "et_ekspress"
TITLE = "Ekspress news article archive (only Estonian) 1.0"
HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1408"

# language codes of the dataset (see ./resources/languages.json)
LANGUAGES = ["et"]


"""
DOWNLOAD
-----------

Instruction

- Run the command on the server:

wget "https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1408/ee_articles_2009-2019.zip"

- Extract the archive:

unzip ee_articles_2009-2019.zip

- Run python script:

python ekspress.py /data/datasets/ele/et/ekspress/ee_articles_2009-2019/ \
    /data/datasets/ele/et/ekspress/ekspress_processed.jsonl --limit 1000
"""


def file_iterator(data_path: str) -> Iterable[str]:
    """
    Iterates through all files in a given directory and outputs path names
    """

    for file in os.listdir(data_path):
        if file.endswith(".json"):
            filename = os.path.join(data_path, file)
        yield filename


def text_extractor(files_path: Iterable[str]) -> Iterable[str]:
    """
    Extracts the article texts from each json file.
    """
    for file in file_iterator(files_path):
        with open(file, "r", encoding="utf8") as inp:
            data = json.load(inp)
            for article in data:
                if article["channelLanguage"] == "nat":  # only include estonian
                    if len(article["bodyText"]) > 250:  # at least over 250 characters
                        yield article["bodyText"]
                else:
                    continue


def save_texts_to_jsonl(
    texts: Iterable[str],
    output_file_path: str,
    append: bool = False,
    limit: int = 0,
    min_length: int = 0,
    print_write_progress: int = 1_000,
):
    """
    Saves a list or generator of texts into a JSON-line file (each line is a valid JSON object)
    """
    mode = "a" if append else "w"
    output_text_field = "text"

    # Save as JSONL
    logger.info(f"Writing output to {output_file_path} ({mode=})")

    docs_count = 0

    with open(output_file_path, mode) as f:
        for i, text in enumerate(texts):
            if min_length > 0 and len(text) < min_length:
                # skip because of short text length
                continue

            f.write(json.dumps({output_text_field: text}, ensure_ascii=False) + "\n")
            docs_count += 1

            if i > 0 and (i % print_write_progress) == 0:
                logger.info(f"Writen {i:,} lines ...")

            if limit > 0 and docs_count >= limit:
                logger.warning(f"Limit reached ({docs_count:,} docs)")
                break


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("input_path", help="Path to input file or directory")
    parser.add_argument(
        "output_path",
        help="Output is saved in file path (JSONL file format)",
    )
    parser.add_argument("--override", action="store_true", help="Override existing output files")
    parser.add_argument("--download", action="store_true", help="Download dataset")
    parser.add_argument("--limit", default=0, type=int, help="Limit dataset size (for debugging)")
    parser.add_argument("--min_text_length", default=0, type=int, help="Min. text length (shorter texts are discarded)")
    args = parser.parse_args()

    if os.path.exists(args.output_path) and not args.override:
        raise FileExistsError(f"Output exists already ({args.output_path}). Fix this error with --override")

    # TODO this can may be optional
    if not os.path.exists(args.input_path):
        raise FileNotFoundError(f"Input does not exist: {args.input_path}")

    # if args.download:
    #     download()

    texts = text_extractor(args.input_path)
    save_texts_to_jsonl(texts, args.output_path, limit=args.limit, min_length=args.min_text_length)

src/lm_datasets/datasets/et/ekspress.py
malteos/lm-datasets
import logging
from typing import Iterable
from lm_datasets.datasets.base import BaseDataset

from smart_open import open


logger = logging.getLogger(__name__)


class ENC2021Dataset(BaseDataset):
    DATASET_ID = "enc2021"
    TITLE = "Estonian National Corpus 2021"
    HOMEPAGE = (
        "https://entu.keeleressursid.ee/shared/9939/EVKultjxSeFA2QhFkbE7fGGDGNT1zmJLOUGFK9hw53tq9Rx2YBTejI1IoKhy65zq"
    )
    LANGUAGES = ["et"]
    DOI = "10.5128/ERYa18.12"

    DOWNLOAD_URLS = [
        # (
        #     "https://entu.keeleressursid.ee/api2/file-24747?key=EVKultjxSeFA2QhFkbE7fGGDGNT1zmJLOUGFK9hw53tq9Rx2YBTejI1IoKhy65zq",
        #     "nc19_Web_2019.prevert.gz",
        # ),
        (
            "https://entu.keeleressursid.ee/api2/file-24753?key=EVKultjxSeFA2QhFkbE7fGGDGNT1zmJLOUGFK9hw53tq9Rx2YBTejI1IoKhy65zq",  # noqa
            "nc21_Web_2021.prevert.gz",
        ),
        (
            "https://entu.keeleressursid.ee/api2/file-24738?key=EVKultjxSeFA2QhFkbE7fGGDGNT1zmJLOUGFK9hw53tq9Rx2YBTejI1IoKhy65zq",  # noqa
            "nc21_Feeds.prevert.gz",
        ),
        (
            "https://entu.keeleressursid.ee/api2/file-24720?key=EVKultjxSeFA2QhFkbE7fGGDGNT1zmJLOUGFK9hw53tq9Rx2YBTejI1IoKhy65zq",  # noqa
            "nc21_DOAJ.prevert.gz",
        ),
        (
            "https://entu.keeleressursid.ee/api2/file-24729?key=EVKultjxSeFA2QhFkbE7fGGDGNT1zmJLOUGFK9hw53tq9Rx2YBTejI1IoKhy65zq",  # noqa
            "nc21_Fiction.prevert.gz",
        ),
    ]

    def is_downloaded(self):
        return len(self.get_dataset_file_paths(needed_suffix=".prevert.gz")) == len(self.DOWNLOAD_URLS)

    def get_texts(self) -> Iterable[str]:
        from lm_datasets.io.prevert_file import PrevertFile

        if not self.is_downloaded():
            self.download()

        file_paths = self.get_dataset_file_paths(needed_suffix=".prevert.gz")

        for fp in file_paths:
            logger.info(f"Reading from {fp}")
            with open(fp) as f:
                dset = PrevertFile(f)

                for doc in dset:  # iterating through documents of a dataset
                    text = str(doc).replace("\n", " ")  # TODO replace line breaks with white spaces

                    yield text

src/lm_datasets/datasets/et/enc.py
malteos/lm-datasets
from typing import Iterable
import zipfile
from lm_datasets.datasets.base import MILLION, BaseDataset, Availability, QualityWarning
from lm_datasets.utils import remove_whitespaces_before_punctuation

import logging

logger = logging.getLogger(__name__)


class EstonianReferenceCorpusDataset(BaseDataset):
    DATASET_ID = "estonian_reference_corpus"
    TITLE = "Estonian Reference Corpus"
    HOMEPAGE = "https://www.cl.ut.ee/korpused/segakorpus/"
    DESCRIPTION = (
        "This corpus includes Estonian texts (fiction, PhD theses, newspapers, magazines, parliamentary ",
        "transcriptions, computer-mediated communication) published between 1990 and 2007. ",
        "The corpus is encoded in TEI. ",
        "The corpus is available for online browsing through a dedicated concordancer and is available ",
        "for download from CELR.",
    )

    LANGUAGES = ["et"]

    DOWNLOAD_URLS = [
        "https://www.cl.ut.ee/korpused/segakorpus/eesti_ilukirjandus_1990/failid/Ilukirjandus.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/postimees/failid/xml/Postimees.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/ekspress/failid/Ekspress.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/epl/failid/xml/Paevaleht.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/maaleht/failid/Maaleht.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/slohtuleht/failid/xml/SLOleht.tar.gz",
        "https://www.cl.ut.ee/korpused/segakorpus/valgamaalane/failid/Valgamaalane.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/laane_elu/failid/LaaneElu.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/horisont/failid/Horisont.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/luup/failid/Luup.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/kroonika/failid/Kroonika.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/eestiarst/failid/EestiArst.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/arvutitehnika/failid/Arvutitehnika.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/agraarteadus/failid/Agraarteadus.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/teadusartiklid/failid/Teadusartiklid.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/seadused/failid/xml/Seadused.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/riigikogu/failid/Riigikogu.zip",
        "https://www.cl.ut.ee/korpused/segakorpus/doktoritood/failid/Doktoritood.zip",
    ]
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LICENSE = "free for non-commercial use"

    QUALITY_WARNINGS = [QualityWarning.BAD_WHITESPACES]

    TOKENS = 175 * MILLION

    def is_downloaded(self):
        return len(self.get_dataset_file_paths(needed_suffix=(".tar.gz", ".zip"))) == len(self.DOWNLOAD_URLS)

    def get_texts(self) -> Iterable[str]:
        """
        Adapted from https://github.com/estnltk/estnltk/blob/123fadf204fb99661da5320cf95172aa9b61c697/tutorials/corpus_processing/importing_text_objects_from_corpora.ipynb  # noqa
        """
        from estnltk.corpus_processing.parse_koondkorpus import get_div_target
        from estnltk.corpus_processing.parse_koondkorpus import parse_tei_corpus_file_content

        if not self.is_downloaded():
            self.download()

        # iterate over archives
        for archive_fp in self.get_dataset_file_paths(needed_suffix=(".tar.gz", ".zip")):
            if archive_fp.endswith(".zip"):
                with zipfile.ZipFile(archive_fp, "r") as zf:
                    # Folders 'bin' contain headers and corpus descriptions.
                    # The '.xml' files outside the 'bin' folders are the files with the actual textual content.
                    for zfn in zf.namelist():
                        if "/bin/" in zfn:
                            continue

                        if ".xml" not in zfn:
                            continue

                        with zf.open(zfn) as xml_file:
                            xml_str = xml_file.read()

                        # find out which subsection of the XML file forms a single document
                        fake_path = archive_fp + "/" + zfn
                        target = get_div_target(fake_path)

                        # import documents as Text objects
                        for text_obj in parse_tei_corpus_file_content(
                            xml_str.decode(),
                            file_path=fake_path,
                            target=[target],
                            sentence_separator=self.sentence_delimiter,
                            paragraph_separator=self.paragraph_delimiter,
                        ):
                            yield remove_whitespaces_before_punctuation(text_obj.text)

            elif archive_fp.endswith(".tar.gz"):
                # raise NotImplementedError
                # TODO
                logger.warning(f"TAR.GZ not implemented: {archive_fp}")

src/lm_datasets/datasets/et/estonian_reference_corpus.py
malteos/lm-datasets
from lm_datasets.datasets.base import MILLION
from lm_datasets.datasets.hf_dataset import HFDataset


class EUSCrawlDataset(HFDataset):
    DATASET_ID = "euscrawl"
    TITLE = "EusCrawl"
    LANGUAGES = ["eu"]
    DESCRIPTION = (
        "EusCrawl (http://www.ixa.eus/euscrawl/) is a high-quality corpus for Basque comprising 12.5 million ",
        "documents and 423 million tokens, totalling 2.1 GiB of uncompressed text. EusCrawl was built using ad-hoc ",
        "scrapers to extract text from 33 Basque websites with high-quality content, resulting in cleaner text ",
        "compared to general purpose approaches.",
    )

    TOKENS = 423 * MILLION
    WEB_CRAWLED = True

    HF_DATASET_ID = "HiTZ/euscrawl"
    HF_DATASET_SPLIT = "train"

    text_column_name = "plain_text"

src/lm_datasets/datasets/eu/euscrawl.py
malteos/lm-datasets
import json
from lm_datasets.datasets.base import BaseDataset, Availability
import zipfile
import logging
from tqdm.auto import tqdm

logger = logging.getLogger(__name__)


class YLENewsDataset(BaseDataset):
    DATASET_ID = "ylenews"
    TITLE = "Yle Finnish News Archive"
    HOMEPAGE = "http://urn.fi/urn:nbn:fi:lb-2021050401"

    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    LANGUAGES = ["fi"]

    def get_texts(self):
        """ """
        zip_fps = self.get_dataset_file_paths(needed_suffix=".zip")

        for zip_fp in zip_fps:
            logger.info(f"Extracting from {zip_fp}")
            with zipfile.ZipFile(zip_fp) as zf:
                member_fns = zf.namelist()
                for fn in tqdm(member_fns):
                    if fn.endswith(".json"):
                        with zf.open(fn) as member_f:
                            obj = json.load(member_f)
                            # data = obj["data"]

                            for doc in obj["data"]:
                                # doc_texts = []
                                doc_text = ""
                                for content in doc["content"]:
                                    if content["type"] == "text":
                                        doc_text += content["text"].rstrip()
                                        doc_text += " "

                                    elif content["type"] == "heading":
                                        if doc_text:
                                            doc_text += self.paragraph_delimiter

                                        doc_text += content["text"] + self.title_delimiter

                                # doc_text = "\n\n".join(doc_texts)
                                yield doc_text

                                # stop after doc
                                # break
                        # stop after file
                        # break
            # stop after archive
            # break

src/lm_datasets/datasets/fi/ylenews.py
malteos/lm-datasets
import logging
import tarfile
from lm_datasets.datasets.base import Availability, BaseDataset

logger = logging.getLogger(__name__)


class CaBeRnetDataset(BaseDataset):
    DATASET_ID = "cabernet"
    TITLE = "CaBeRnet: a New French Balanced Reference Corpus"
    LANGUAGES = ["fr"]
    HOMEPAGE = "https://aclanthology.org/2020.cmlc-1.3/"

    AVAILIBITY = Availability.ON_REQUEST
    TOKENS = 711_792_861

    def decompress(self):
        # gzip -d CaBeRnet.tar.gz
        pass

    def get_texts(self):
        # read from tar file
        archive_fp = self.get_dataset_file_paths(needed_suffix=".tar", single_file=True)

        logger.info(f"Extracting from {archive_fp}")

        with tarfile.open(archive_fp) as tar_f:
            member_fns = tar_f.getmembers()
            member_fns = [m for m in tar_f.getmembers() if m.name.endswith(".txt")]

            # Read from txt files
            for member_fn in member_fns:
                logger.info(f"Decompress: {member_fn}")
                decompressed_member = tar_f.extractfile(member_fn).read()
                txt = decompressed_member.decode()

                if txt:
                    pass

                # for doc_text in txt.split("\n\n"):
                #     yield doc_text

                break
                print("x")

src/lm_datasets/datasets/fr/cabernet.py
malteos/lm-datasets
import logging
from lm_datasets.datasets.base import BaseDataset, MB, Availability


logger = logging.getLogger(__name__)


class GABilingualLegislationDataset(BaseDataset):
    """

    TODO only sentences no documents

    """

    DATASET_ID = "ga_bilingual_legistation"
    TITLE = "Irish legislation"
    HOMEPAGE = "https://portulanclarin.net/repository/browse/the-gaois-bilingual-corpus-of-english-irish-legislation-processed/daeac17c9e3511ea9b7f02420a000407b83de243dc0b469aab41084386c5b80f/"  # noqa

    LANGUAGES = ["ga"]

    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    DOWNLOAD_URLS = []

    USED_BY = ["gaBERT"]

    BYTES = 0.5 * 25 * MB

    def get_texts(self):
        from translate.storage.tmx import tmxfile

        with open(self.get_dataset_file_paths(single_file=True, needed_suffix=".tmx"), "rb") as fin:
            tmx_file = tmxfile(fin, "ga", "en")

        for i, node in enumerate(tmx_file.unit_iter()):
            text = node.source  # ga
            # en => node.target
            yield text

src/lm_datasets/datasets/ga/ga_bilingual_legistation.py
malteos/lm-datasets
import logging
import os
from lm_datasets.datasets.base import BaseDataset, MB, Availability
from lm_datasets.systems import get_path_by_system

import wget

logger = logging.getLogger(__name__)


class GAUniversalDependenciesDataset(BaseDataset):
    """

    TODO only sentences no documents


    """

    DATASET_ID = "ga_universal_dependencies"
    TITLE = "Irish Universal Dependencies"
    HOMEPAGE = "https://universaldependencies.org/"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["ga"]

    DOWNLOAD_URLS = [
        "https://github.com/UniversalDependencies/UD_Irish-IDT/raw/r2.12/ga_idt-ud-train.conllu",
        "https://github.com/UniversalDependencies/UD_Irish-IDT/raw/r2.12/ga_idt-ud-dev.conllu",
        "https://github.com/UniversalDependencies/UD_Irish-IDT/raw/r2.12/ga_idt-ud-test.conllu",
        "https://github.com/UniversalDependencies/UD_Irish-Cadhan/raw/r2.12/ga_cadhan-ud-test.conllu",
        "https://github.com/UniversalDependencies/UD_Irish-TwittIrish/raw/r2.12/ga_twittirish-ud-train.conllu",
        "https://github.com/UniversalDependencies/UD_Irish-TwittIrish/raw/r2.12/ga_twittirish-ud-dev.conllu",
        "https://github.com/UniversalDependencies/UD_Irish-TwittIrish/raw/r2.12/ga_twittirish-ud-test.conllu",
    ]

    USED_BY = ["gaBERT"]

    BYTES = 9.7 * MB

    def download(self):
        output_dir = get_path_by_system(self.LOCAL_DIRS)

        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        file_paths = []

        for url in self.DOWNLOAD_URLS:
            fn = url.split("/")[-1]
            fp = os.path.join(output_dir, fn)
            # try:
            out_filename = wget.download(url, out=fp)
            logger.info(f"Completed {out_filename}")

            file_paths.append(fp)
            # except HTTPError as e:
            #     logger.error(f"Error {e}")

        return file_paths

    def get_dataset_file_paths(self):
        # TODO handle datasets with multiple filles
        dataset_dir = get_path_by_system(self.LOCAL_DIRS)

        if not os.path.exists(dataset_dir):
            return []

        fps = [os.path.join(dataset_dir, f) for f in os.listdir(dataset_dir)]
        fps = [fp for fp in fps if os.path.isfile(fp) and fp.endswith(".conllu")]

        return fps

    def get_texts(self):
        import conllu

        file_paths = self.get_dataset_file_paths()

        if len(file_paths) != len(self.DOWNLOAD_URLS):
            file_paths = self.download()

        # Parse CONLL files and extract sentences
        for fp in file_paths:
            logger.info(f"Reading {fp}")

            with open(fp) as f:
                data = f.read()

            sentences = conllu.parse(data)

            for sent in sentences:
                text = sent.metadata["text"]

                yield text

src/lm_datasets/datasets/ga/ga_universal_dependencies.py
malteos/lm-datasets
import logging
from datasets import load_dataset, DatasetDict

from typing import Dict, List, Optional

from lm_datasets.datasets.base import BaseDataset


logger = logging.getLogger(__name__)


class HFDataset(BaseDataset):
    HF_DATASET_ID: str = None
    HF_DATASET_SPLIT: Optional[str] = None
    HF_DATASET_CONFIGS: Optional[List[str]] = None
    HF_DATA_DIR = None

    config_to_dataset: Optional[Dict] = None
    text_column_name = "text"
    title_column_name = None
    remove_columns = None
    streaming = False

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        if self.HF_DATASET_ID is None:
            raise ValueError("HF_DATASET_ID is not set")

    def get_hf_configs(self):
        if self.HF_DATASET_CONFIGS:
            return self.HF_DATASET_CONFIGS
        else:
            # if no config is used
            return [None]

    def download(self):
        self.config_to_dataset = {}

        for hf_config in self.get_hf_configs():
            logger.info(f"Downloading for {hf_config=}")

            ds = load_dataset(
                self.HF_DATASET_ID,
                hf_config,
                split=self.HF_DATASET_SPLIT,
                data_dir=self.HF_DATA_DIR,
                streaming=self.streaming,
                use_auth_token=self.get_hf_auth_token(),
                keep_in_memory=False,
            )

            # check dataset split
            if isinstance(ds, DatasetDict) and not self.HF_DATASET_SPLIT:
                logger.warning(f"HF returned DatasetDict but split not set: {DatasetDict}")

            if self.limit > 0:
                if self.streaming:
                    logger.warning(f"Limit requested ({self.limit=}) but streaming is enabled!")
                else:
                    logger.warning(f"Limiting dataset to: {self.limit}")
                    ds = ds.select(range(self.limit))

            if self.remove_columns is not None:
                logger.info(f"Removing columns (at download): {self.remove_columns}")

                ds = ds.remove_columns(self.remove_columns)

            filter_func = self.get_filter_func()
            if filter_func:
                logger.info(f"Dataset size before filter: {len(ds):,}")

                ds = ds.filter(filter_func, num_proc=self.workers)

                logger.info(f"Dataset size after filter: {len(ds):,}")

            self.config_to_dataset[hf_config] = ds

    def get_filter_func(self):
        return None

    def get_text_from_item(self, item) -> str:
        return item[self.text_column_name]

    def get_texts(self):
        self.download()

        # drop all non-text columns
        for ds_idx, config in enumerate(self.config_to_dataset):
            # remove non-text and non-title columns
            columns_to_remove = set(self.config_to_dataset[config].column_names) - {self.text_column_name}

            if self.title_column_name:
                columns_to_remove = columns_to_remove - {self.title_column_name}

            logger.info(f"Removing columns (get texts): {columns_to_remove}")

            self.config_to_dataset[config] = self.config_to_dataset[config].remove_columns(columns_to_remove)

            if self.title_column_name:
                logger.info(f"Prepending title to text column ({self.title_column_name=})")

                def prepend_title(example):
                    example[self.text_column_name] = (
                        example[self.title_column_name] + self.title_delimiter + example[self.text_column_name]
                    )

                    return example

                self.config_to_dataset[config] = self.config_to_dataset[config].map(prepend_title)

                # remove title column
                self.config_to_dataset[config] = self.config_to_dataset[config].remove_columns([self.title_column_name])

            ds_iterator = iter(self.config_to_dataset[config])

            for item in ds_iterator:
                yield self.get_text_from_item(item)

            # yield from self.config_to_dataset[config][self.text_column_name]

            # if self.text_column_name != self.output_text_field:
            #     # rename text column to output
            #     self.config_to_dataset[config] = self.config_to_dataset[config].rename_column(
            #         self.text_column_name, self.output_text_field
            #     )

            # write_mode = "a" if ds_idx > 0 else "w"

            # write to JSON line files
            # logger.info(f"Writing output to {self.get_output_file_path()}; {write_mode=}")
            # self.config_to_dataset[config].to_json(
            #     self.get_output_file_path(), lines=True, mode=write_mode, force_ascii=self.json_ensure_ascii
            # )

src/lm_datasets/datasets/hf_dataset.py
malteos/lm-datasets
"""

https://repository.pfri.uniri.hr/islandora/object/pfri%3A2156/datastream/DATASET0/download

"""

import logging


from lm_datasets.datasets.base import BaseDataset, Availability


from smart_open import open


logger = logging.getLogger(__name__)


class CroatianNewsENGRIDataset(BaseDataset):
    """


    https://www.clarin.si/repository/xmlui/handle/11356/1416

    """

    DATASET_ID = "croatian_news_engri"
    TITLE = "Corpus of Croatian news portals ENGRI (2014-2018)"
    HOMEPAGE = "https://repository.pfri.uniri.hr/islandora/object/pfri%3A2156"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    DESCRIPTION = (
        "The corpus consists of texts collected from the most popular (based on the Reuters Institute Digital News ",
        "Report for 2018, retrieved from http://www.digitalnewsreport.org in April, 2019) news portals in Croatia ",
        "in the period from 2014 to 2018: Direktno, Dnevno, Net Hr, Hrt, Index_Hr, Jutarnji, Novilist, Rtl, ",
        "SlobodnaDalmacija, Večernji, Tportal, Dnevnik.",
    )

    LANGUAGES = ["hr"]

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/hr/croatian_news_engri"]

    TOKENS = 694799268

    @property
    def DOWNLOAD_URLS(self):
        fns = "engri.24sata.hr.conllu.gz,/engri.direktno.hr.conllu.gz,/engri.dnevno.hr.conllu.gz,/engri.hrt.hr.conllu.gz,/engri.index.hr.conllu.gz,/engri.jutarnji.hr.conllu.gz,/engri.net.hr.conllu.gz,/engri.novilist.hr.conllu.gz,/engri.rtl.hr.conllu.gz,/engri.slobodnadalmacija.hr.conllu.gz,/engri.telegram.hr.conllu.gz,/engri.vecernji.hr.conllu.gz".split(  # noqa
            ",/"
        )

        return ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1416" + fn for fn in fns]

    def get_texts(self):
        import conllu

        # Parse CONLLU files and extract document-level texts

        for fp in self.get_dataset_file_paths(needed_suffix=".conllu.gz"):
            logger.info(f"Reading from {fp}")

            with open(fp) as f:
                text = None

                for sentence in conllu.parse_incr(f):
                    if "newdoc id" in sentence.metadata:
                        if text is not None:
                            # doc completed
                            yield text
                        text = ""  # init

                    # append text to doc
                    text += sentence.metadata["text"]

                    if "title" in sentence.metadata:
                        text += ":\n"

                # yield last document
                if text is not None:
                    yield text

src/lm_datasets/datasets/hr/croatian_news_engri.py
malteos/lm-datasets
import logging
from lm_datasets.datasets.base import BaseDataset, Availability

import html

from smart_open import open


logger = logging.getLogger(__name__)


class HRWACDataset(BaseDataset):
    """
    TODO only paragraphs no documents

    curl --remote-name-all https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1064{/hrWaC2.1.01.xml.gz,/hrWaC2.1.02.xml.gz,/hrWaC2.1.03.xml.gz,/hrWaC2.1.04.xml.gz,/hrWaC2.1.05.xml.gz,/hrWaC2.1.05.xml.gz,/hrWaC2.1.06.xml.gz,/hrWaC2.1.07.xml.gz,/hrWaC2.1.08.xml.gz,/hrWaC2.1.09.xml.gz,/hrWaC2.1.10.xml.gz,/hrWaC2.1.11.xml.gz,/hrWaC2.1.12.xml.gz,/hrWaC2.1.13.xml.gz,/hrWaC2.1.14.xml.gz}  # noqa
    """

    DATASET_ID = "hrwac"
    TITLE = "Croatian web corpus hrWaC 2.1"
    HOMEPAGE = "http://nlp.ffzg.hr/resources/corpora/hrwac/"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["hr"]

    DOWNLOAD_URLS = [
        f"https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1064/hrWaC2.1.{i:02d}.xml.gz"
        for i in range(1, 14)
    ]
    # LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/hr/hrWaC2.1"]

    TOKENS = 1_397_757_548

    def get_paragraphs(self, file_path: str, needed_language: str, min_length: int = 250):
        lang = None
        paragraph_i = 0

        with open(file_path, encoding="utf-8") as f:
            paragraph = ""
            sentence = ""
            last_line = ""

            for line in f:
                line = line.strip()
                # print(line)

                if line.startswith("<p"):
                    # print(line)
                    if 'lang="sr"' in line:
                        lang = "sr"
                    elif 'lang="hr"' in line:  # TODO
                        lang = "hr"
                    else:
                        raise ValueError("Cannot determine language")

                    paragraph = ""
                elif line == "</p>":
                    # end of paragraph
                    lang = None
                    # print(f"======={paragraph}")
                    paragraph = paragraph.strip()

                    if len(paragraph) > min_length:
                        paragraph_i += 1
                        yield paragraph

                        if paragraph_i > self.limit and self.limit > 0:
                            logger.warning("Limit reached")
                            break

                elif needed_language == lang:
                    if line == "<s>":
                        # empty = True
                        sentence = ""
                    elif line == "</s>":
                        # end of sentence
                        # print(f"======={sentence}")

                        paragraph += sentence

                        # if not empty:
                        #     print("")
                    # elif line == "<g/>":
                    #     whitespace = False
                    elif line.startswith("<"):
                        pass
                    else:
                        # empty = False
                        decoded = html.unescape(line)  # .decode("utf-8")
                        original, diacritic, lemma, pos = decoded.split("\t")
                        # word = (
                        #     diacritic[0].upper() + diacritic[1:] + "\t" + pos.rstrip() + "\t" + lemma
                        #     if lemma[0].isupper()
                        #     else diacritic + "\t" + pos.rstrip() + "\t" + lemma
                        # )

                        if last_line != "<g/>":
                            sentence += " "

                        sentence += original

                last_line = line

    def is_downloaded(self):
        return len(self.get_dataset_file_paths(needed_suffix=".xml.gz")) == len(self.DOWNLOAD_URLS)

    def get_texts(self):
        needed_language = "hr"  # TODO make sr dataset

        if not self.is_downloaded():
            self.download()

        file_paths = self.get_dataset_file_paths(needed_suffix=".xml.gz")

        # Parse XML files and extract paragraphs
        for i, fp in enumerate(file_paths):
            logger.info(f"Reading {fp}")

            yield from self.get_paragraphs(file_path=fp, needed_language=needed_language)

src/lm_datasets/datasets/hr/hrwac.py
malteos/lm-datasets
import logging

import pandas as pd

from lm_datasets.datasets.base import BaseDataset, GB, Availability, Genre


logger = logging.getLogger(__name__)


class StyriaNewsDataset(BaseDataset):
    """
    https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1410/Styria-articles.csv
    https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1399/Styria-user-comments.zip?sequence=7&isAllowed=y

    """

    DATASET_ID = "styria_news"
    TITLE = "24sata news article archive 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1410"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["hr"]
    GENRES = [Genre.NEWS]

    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1410/Styria-articles.csv"]

    BYTES = 1.3 * GB

    def is_downloaded(self):
        return bool(self.get_dataset_file_paths(single_file=True))

    def get_texts(self):
        if not self.is_downloaded():
            self.download()

        # read CSV in chunks
        for df in pd.read_csv(self.get_dataset_file_paths(single_file=True), chunksize=10_000):
            for title, content in zip(df.title, df.content):
                text = f"{title}: {content}"

                yield text

src/lm_datasets/datasets/hr/styria_news.py
malteos/lm-datasets
import logging
from typing import Iterable
from lm_datasets.datasets.base import BILLION, Availability, BaseDataset
from smart_open import open

logger = logging.getLogger(__name__)


class ITWacDataset(BaseDataset):
    """
    See also: DeWacDataset
    """

    DATASET_ID = "itwac"
    TITLE = "ITWaC"
    DESCRIPTION = (
        "itWaC: a 2 billion word corpus constructed from the Web limiting the crawl to the .it domain and using"
        " medium-frequency words from the Repubblica corpus and basic Italian vocabulary lists as seeds. "
    )
    LANGUAGES = ["it"]
    AVAILIBILITY = Availability.ON_REQUEST
    HOMEPAGE = "https://docs.sslmit.unibo.it/doku.php?id=corpora:itwac"

    TOKENS = 2 * BILLION

    def decompress(self):
        #  7z x itwac_preproc.7z
        pass

    def get_texts(self) -> Iterable[str]:
        file_path = self.get_dataset_file_paths(single_file=True, needed_suffix=".filtered.pre.pos.corpus")

        # chardet.detect(line)

        # New documents are separated by "CURRENT URL <url>"
        doc_text = ""
        with open(file_path, "r", encoding="ISO-8859-1") as f:
            for i, line in enumerate(f):
                if line.startswith("CURRENT URL "):
                    if doc_text:
                        # print(doc_text)
                        yield doc_text

                    doc_text = ""
                else:
                    doc_text += line.strip()

                pass

            if doc_text:
                # last document
                yield doc_text

src/lm_datasets/datasets/it/itwac.py
malteos/lm-datasets
import logging
import json
import os

from smart_open import open

import zstandard as zstd
import io


from lm_datasets.datasets.base import BaseDataset


logger = logging.getLogger(__name__)


class JSONLDataset(BaseDataset):
    raw_jsonl_paths = None
    raw_jsonl_text_field = "text"

    def is_downloaded(self):
        for fp in self.get_raw_jsonl_paths():
            if not os.path.exists(fp):
                return False

        return True

    def get_raw_jsonl_paths(self):
        if self.raw_jsonl_paths is None:
            raise ValueError("Cannot load JSONL dataset since `raw_jsonl_paths` is not set")

        dataset_dir = self.get_local_dataset_dir()

        for fp in self.raw_jsonl_paths:
            if fp.startswith("*"):
                # find based on suffix
                yield from self.get_dataset_file_paths(
                    dataset_dir=dataset_dir, subdirectories=True, single_file=False, needed_suffix=fp.lstrip("*")
                )
            else:
                yield os.path.join(dataset_dir, fp)

    def get_text_from_item(self, item):
        """
        This simply returns the text field from item (but dataset classes can override this to implement filtering etc.)
        """
        return item[self.raw_jsonl_text_field]

    def get_texts_from_file_handler(self, file_handler):
        for line in file_handler:
            item = json.loads(line)
            text = self.get_text_from_item(item)

            if text:
                yield text

    def get_texts(self):
        """
        Iterate over all input files and read JSON from each line.
        """
        for fp in self.get_raw_jsonl_paths():
            logger.info(f"Reading from {fp}")

            if fp.endswith(".zst"):  # zstd compression
                with open(fp, "rb") as zf:
                    dctx = zstd.ZstdDecompressor()  # uncompress zstd
                    with dctx.stream_reader(zf) as reader:
                        f = io.BufferedReader(reader)
                        yield from self.get_texts_from_file_handler(f)
            else:
                with open(fp) as f:  # jsonl or jsonl.fz (via smart_open)
                    yield from self.get_texts_from_file_handler(f)

src/lm_datasets/datasets/jsonl_dataset.py
malteos/lm-datasets
import zipfile
from lm_datasets.datasets.base import BaseDataset, Genre


class SeismasLTENDataset(BaseDataset):
    DATASET_ID = "seimas_lt_en"
    TITLE = "Bilingual English-Lithuanian parallel corpus from Seimas of the Republic of Lithuania website"
    HOMEPAGE = "https://live.european-language-grid.eu/catalogue/corpus/3009/download/"

    DOWNLOAD_URLS = [
        "https://elrc-share.eu/repository/download/4486f8e4e72711e7b7d400155d0267060b3d0987d08b43fd9c065ce3f05f99f8"
    ]
    LANGUAGES = ["lt"]
    GENRES = [Genre.GOVERNMENT]

    BYTES = 160 * 1024

    def get_texts(self):
        from translate.storage.tmx import tmxfile

        zip_fp = self.get_dataset_file_paths(needed_suffix=".zip", single_file=True)

        with zipfile.ZipFile(zip_fp) as zf:
            for fn in zf.namelist():
                if fn.endswith(".tmx"):
                    with zf.open(fn) as member_f:
                        tmx_file = tmxfile(member_f, "lt", "en")

                for i, node in enumerate(tmx_file.unit_iter()):
                    text = node.target  # lt
                    # en => node.source

                    yield text

src/lm_datasets/datasets/lt/seimas_lt_en.py
malteos/lm-datasets
import logging
import zipfile

from lm_datasets.datasets.base import MB, Availability, BaseDataset

logger = logging.getLogger(__name__)


class StateRelatedLatvianWebDataset(BaseDataset):
    DATASET_ID = "state_related_latvian_web"
    TITLE = "Corpus of State-related content from the Latvian Web (Processed)"
    HOMEPAGE = "http://catalog.elra.info/en-us/repository/browse/ELRA-W0169/"
    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    BYTES = 3.4 * MB
    LANGUAGES = ["lv"]

    def get_texts(self):
        from translate.storage.tmx import tmxfile

        zip_fp = self.get_dataset_file_paths(needed_suffix=".zip", single_file=True)

        with zipfile.ZipFile(zip_fp) as zf:
            for fn in zf.namelist():
                if fn.endswith(".tmx"):
                    with zf.open(fn) as member_f:
                        tmx_file = tmxfile(member_f, "lv", "en")

                    for i, node in enumerate(tmx_file.unit_iter()):
                        text = node.target  # lv
                        # en => node.source

                        yield text

src/lm_datasets/datasets/lv/state_related_latvian_web.py
malteos/lm-datasets
import json
from pathlib import Path
from lm_datasets.datasets.base import BaseDataset, Availability
from lm_datasets.datasets.hf_dataset import HFDataset


from tqdm.auto import tqdm
import logging

logger = logging.getLogger(__name__)


class KorpusMaltiHFDataset(HFDataset):
    DATASET_ID = "korpus_malti"

    TITLE = "Korpus Malti"
    HOMEPAGE = "https://huggingface.co/datasets/MLRS/korpus_malti"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["mt"]

    DESCRIPTION = (
        "General Corpora for the Maltese Language. This dataset is composed of texts ",
        "from various genres/domains written in Maltese.",
    )

    PII = "I have not checked the data source for personally identifiable or sensitive information."

    LICENSE = "cc-by-nc-sa-4.0"  # DFKI has a permission for LLM training with commercial license

    # From paper https://aclanthology.org/2022.deeplo-1.10.pdf (Table 1)
    TOKENS = 466_601_373 - 98_582_031 - 1_885_661  # full dataset - law_eu - wiki

    HF_DATASET_ID = "MLRS/korpus_malti"
    HF_DATASET_SPLIT = "train"
    HF_DATASET_CONFIGS = [
        "belles_lettres",
        "blogs",
        "comics",
        "court",
        "eu_docs",
        "gov_docs",
        "government_gazzette",
        # "law_eu" # avoid duplication with EUR-LEx
        "law_mt",
        "legal",
        "nonfiction",
        "parliament",
        "press_eu",
        "press_mt",
        "speeches",
        "theses",
        "umlib_oar",
        "web_general",
        # "wiki",  # avoid duplication with Wiki
    ]
    # text_column_name = "text"
    # remove_columns = ["doc_id", "LICENSE", "uri", "date_built"]

    def get_text_from_item(self, item) -> str:
        sentences = item["text"]
        text = self.sentence_delimiter.join(sentences)

        return text

    # def get_texts(self):
    #     # load dataset for each config
    #     for config in self.dataset_configs:
    #         logger.info(f"Processing: {config}")
    #         self.datasets[config] = load_dataset(
    #             self.HF_DATASET_ID,
    #             config,
    #             streaming=False,
    #             use_auth_token=self.get_hf_auth_token(),
    #         )

    #         for sentences in self.datasets[config]["train"]["text"]:
    #             text = self.sentence_delimiter.join(sentences)

    #             yield text


class KorpusMaltiDataset(BaseDataset):
    DATASET_ID = "korpus_malti"

    TITLE = "Korpus Malti"
    HOMEPAGE = "https://huggingface.co/datasets/MLRS/korpus_malti"
    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    LANGUAGES = ["mt"]

    DESCRIPTION = (
        "General Corpora for the Maltese Language. This dataset is composed of texts ",
        "from various genres/domains written in Maltese.",
    )

    PII = "I have not checked the data source for personally identifiable or sensitive information."

    LICENSE = "cc-by-nc-sa-4.0"  # DFKI has a permission for LLM training with commercial license

    # From paper https://aclanthology.org/2022.deeplo-1.10.pdf (Table 1)
    TOKENS = 466_601_373 - 98_582_031 - 1_885_661  # full dataset - law_eu - wiki

    HF_DATASET_ID = "MLRS/korpus_malti"
    HF_DATASET_SPLIT = "train"
    HF_DATASET_CONFIGS = [
        "belles_lettres",
        "blogs",
        "comics",
        "court",
        "eu_docs",
        "gov_docs",
        "government_gazzette",
        # "law_eu" # avoid duplication with EUR-LEx
        "law_mt",
        "legal",
        "nonfiction",
        "parliament",
        "press_eu",
        "press_mt",
        "speeches",
        "theses",
        "umlib_oar",
        "web_general",
        # "wiki",  # avoid duplication with Wiki
    ]

    def download(self):
        """
        git clone --depth 1 https://$HF_LOGIN:$HF_PASSWORD@huggingface.co/datasets/MLRS/korpus_malti
        """
        pass

    def get_texts(self):
        dataset_dir = self.get_local_dataset_dir()

        for subset in self.HF_DATASET_CONFIGS:
            subset_path = Path(dataset_dir) / "data" / subset
            # logger.info(f"")
            fps = list(subset_path.glob("*.jsonl"))
            for fp in tqdm(fps, desc=f"Reading subset: {subset}"):
                with open(fp) as f:
                    # logger.info(f"Reading from {fp}")
                    for line in f:
                        item = json.loads(line)
                        text = self.sentence_delimiter.join(item["text"])

                        yield text

src/lm_datasets/datasets/mt/korpus_malti.py
malteos/lm-datasets
import logging
import os
from pathlib import Path

import numpy as np

from lm_datasets.datasets.base import Availability
from lm_datasets.datasets.jsonl_dataset import JSONLDataset


logger = logging.getLogger(__name__)


DEFAULT_OSCAR_MIN_HARMFUL_PP = 20.0

"""
'colossal_oscar_05-06-23_bg,colossal_oscar_05-06-23_cs,colossal_oscar_05-06-23_da,colossal_oscar_05-06-23_el,colossal_oscar_05-06-23_et,colossal_oscar_05-06-23_fi,colossal_oscar_05-06-23_fr,colossal_oscar_05-06-23_ga,colossal_oscar_05-06-23_hr,colossal_oscar_05-06-23_hu,colossal_oscar_05-06-23_lt,colossal_oscar_05-06-23_lv,colossal_oscar_05-06-23_mt,colossal_oscar_05-06-23_nl,colossal_oscar_05-06-23_pl,colossal_oscar_05-06-23_pt,colossal_oscar_05-06-23_ro,colossal_oscar_05-06-23_sk,colossal_oscar_05-06-23_sl,colossal_oscar_05-06-23_sv,colossal_oscar_05-06-23_uk,colossal_oscar_05-06-23_sr,colossal_oscar_05-06-23_sh,colossal_oscar_05-06-23_nn,colossal_oscar_05-06-23_no,colossal_oscar_05-06-23_eu,colossal_oscar_05-06-23_ca,colossal_oscar_05-06-23_gl,colossal_oscar_03-04-23_bg,colossal_oscar_03-04-23_cs,colossal_oscar_03-04-23_da,colossal_oscar_03-04-23_el,colossal_oscar_03-04-23_et,colossal_oscar_03-04-23_fi,colossal_oscar_03-04-23_fr,colossal_oscar_03-04-23_ga,colossal_oscar_03-04-23_hr,colossal_oscar_03-04-23_hu,colossal_oscar_03-04-23_lt,colossal_oscar_03-04-23_lv,colossal_oscar_03-04-23_mt,colossal_oscar_03-04-23_nl,colossal_oscar_03-04-23_pl,colossal_oscar_03-04-23_pt,colossal_oscar_03-04-23_ro,colossal_oscar_03-04-23_sk,colossal_oscar_03-04-23_sl,colossal_oscar_03-04-23_sv,colossal_oscar_03-04-23_uk,colossal_oscar_03-04-23_sr,colossal_oscar_03-04-23_sh,colossal_oscar_03-04-23_nn,colossal_oscar_03-04-23_no,colossal_oscar_03-04-23_eu,colossal_oscar_03-04-23_ca,colossal_oscar_03-04-23_gl'
"""  # noqa
OSCAR_DUMPS = ["05-06-23", "03-04-23"]
LANGUAGES = "bg cs da el et fi ga hr hu lt lv mt nl pl pt ro sk sl sv uk sr sh nn no eu ca gl".split(" ")  # removed fr

EXCLUDE_CATEGORIES = {
    # See http://dsi.ut-capitole.fr/blacklists/index_en.php
    "agressif",
    "adult",
    "cryptojacking",
    "dangerous_material",
    "phishing",
    "warez",
}


class ColossalOscarBaseDataset(JSONLDataset):
    """
    Read OSCAR output from jsonl.zst files (as provided on HF)
    """

    SOURCE_ID = "colossal_oscar"

    DESCRIPTION = "Colossal OSCAR 1"
    HOMEPAGE = "https://huggingface.co/datasets/oscar-corpus/colossal-oscar-1.0"
    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    LOCAL_DIRS = ["pegasus:/netscratch/mostendorff/experiments/eulm/data/colossal-oscar-1.0"]

    DATASET_ID = "colossal_oscar_05-06-23_da"
    LANGUAGES = ["da"]
    DUMP_VERSION = "05-06-23"

    min_harmful_pp = None

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        # set min_harmful_pp

        # TODO hard coded
        # resources_dir = os.path.join("/netscratch/mostendorff/experiments/eulm", "resources", "oscar")

        # with open(os.path.join(resources_dir, "adult_pp_dict.pkl"), "rb") as f:
        #     c_dict_means = pickle.load(f)

        # with open(os.path.join(resources_dir, "adult_pp_dict_full.pkl"), "rb") as file:
        #     # Mean perplexity for documents with the Adult quality warning in the Oscar Dataset
        #     dict_full = pickle.load(file)

        dict_full = {}

        self.min_harmful_pp = (
            np.mean(dict_full[self.get_language_code()])
            if self.get_language_code() in dict_full
            else DEFAULT_OSCAR_MIN_HARMFUL_PP
        )

        # logger.info(f"{self.min_harmful_pp=}")

    def get_text_from_item(self, doc):
        if doc["metadata"]["quality_warnings"]:
            self.counter.update({"filtered_quality_warnings": 1})
            return None
        elif (
            "harmful_pp" in doc["metadata"]
            and doc["metadata"]["harmful_pp"]
            and doc["metadata"]["harmful_pp"] < self.min_harmful_pp
        ):
            self.counter.update({"filtered_harmful_pp": 1})
            return None
        elif doc["metadata"]["categories"] and len(set(doc["metadata"]["categories"]) & EXCLUDE_CATEGORIES) > 0:
            self.counter.update({"filtered_categories": 1})
            return None
        else:
            return doc["content"]

    def get_raw_jsonl_paths(self):
        lang = self.get_language_code()
        dataset_path = Path(os.path.join(self.get_local_dataset_dir(), self.DUMP_VERSION, f"{lang}_meta"))

        return sorted([str(p) for p in dataset_path.glob("*.jsonl.zst")])

    def get_bytes(self):
        return sum(os.stat(fp).st_size for fp in self.get_raw_jsonl_paths())


def get_colossal_oscar_class(lang, dump_version):
    class ColossalOscarDataset(ColossalOscarBaseDataset):
        DATASET_ID = f"colossal_oscar_{dump_version}_{lang}"
        LANGUAGES = [lang]
        DUMP_VERSION = dump_version

    return ColossalOscarDataset


def get_colossal_oscar_auto_classes():
    """
    Auto generate dataset classes
    """

    return [get_colossal_oscar_class(lang, dump_version) for dump_version in OSCAR_DUMPS for lang in LANGUAGES]

src/lm_datasets/datasets/multilingual/colossal_oscar.py
malteos/lm-datasets
import logging
from typing import Iterable
from smart_open import open

from lm_datasets.datasets.base import BaseDataset, Availability
from tqdm.auto import tqdm

from lm_datasets.io.conllu_file import get_texts_from_conllu_file

logger = logging.getLogger(__name__)


class CurlicatBaseDataset(BaseDataset):
    SOURCE_ID = "curlicat"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    conllup_encoding = "utf-8"

    def download(self):
        """
        DOWNLOAD
        -----------

        Instruction

        - Run the command on the server:

        wget -O file.zip \
            "https://elrc-share.eu/repository/download/fed6af2a590311ed9c1a00155d0267062ed273d01d2343f1b78d08d4d481679d/"

        - Extract the archive:

        unzip file.zip

        """  # noqa
        pass

    def decompress(self):
        # unzip
        pass

    def get_curclicat_file_paths(self):
        return self.get_dataset_file_paths(needed_suffix=".conllup")

    def get_texts(self) -> Iterable[str]:
        """
        Extracts the text from each conllu file.
        """

        # TODO read directly from compressed files
        fps = self.get_curclicat_file_paths()
        for i, fp in tqdm(enumerate(fps), total=len(fps), desc="Reading files"):
            # doc = []
            if self.skip_items > 0 and i < self.skip_items:
                logger.warning(f"Skipping {fp}")
                continue

            with open(fp, "r", encoding=self.conllup_encoding) as f:
                yield from get_texts_from_conllu_file(
                    f,
                    title_delimiter=self.title_delimiter,
                    sentence_delimiter=self.sentence_delimiter,
                    skip_sentence_prefixes=["tf.Tensor("],
                )


class CurlicatBGDataset(CurlicatBaseDataset):
    DATASET_ID = "curlicat_bg"
    TITLE = "CURLICAT Bulgarian Corpus"
    DESCRIPTION = """The Bulgarian CURLICAT corpus consists of texts from different sources, provided with  appropriate licences for distribution. We used three general types of sources with regard to the metadata extraction: Bulgarian National Corpus (provided that they have redistributable licensing terms); some public repositories with open and copyright free data; blogs with redistributable licenses, open content websites, etc. The Bulgarian CURLICAT collection contains 113 087 documents, distributed in seven thematic domains: Culture, Education, European Union, Finance, Politics, Economics, and Science. For more information see the CURLICAT website (http:curlicat-project.eu/deliverables) """  # noqa

    HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-bulgarian-corpus/fed6af2a590311ed9c1a00155d0267062ed273d01d2343f1b78d08d4d481679d/"  # noqa

    LANGUAGES = ["bg"]

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/bg/curlicat_bg/Bulgarian_Curlicat_corpus"]

    TOKENS = 35_319_695

    # done


class CurlicatHRDataset(CurlicatBaseDataset):
    """
    unzip and move to raw dir
    find hr_curlicat_2023-01-13 -name '*.*' -exec mv {} /home/mostendorff/experiments/eulm/data/ele/hr/curlicat_hr;
    """

    DATASET_ID = "curlicat_hr"
    TITLE = "CURLICAT Croatian Corpus"

    HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-croatian-corpus/00815518592811ed9c1a00155d026706bc4c59740fce4f7986213e7eef133023/"  # noqa

    LANGUAGES = ["hr"]

    DOWNLOAD_URLS = ["https://zzl.ffzg.unizg.hr/files/5cd9d6151612489fafb2c1f116140519/hr_curlicat_2023-01-13.zip"]

    TOKENS = 49_007_508

    # done


class CurlicatHUDataset(CurlicatBaseDataset):
    DATASET_ID = "curlicat_hu"
    TITLE = "CURLICAT Hungarian Corpus"
    HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-hungarian-corpus/8b6c8dcb58ea11ed9c1a00155d02670679a453431c8147079e5a7d9b879a9729/"  # noqa
    LANGUAGES = ["hu"]
    TOKENS = 61_196_946

    def get_curclicat_file_paths(self):
        # Hugarian files are stored in *.txt files subdirectories
        return self.get_dataset_file_paths(subdirectories=True, needed_suffix=".txt")

    # extracted


class CurlicatPLDataset(CurlicatBaseDataset):
    """
    unzip and move to raw dir
    find pl-20221026-annotated -name '*.*' -exec mv {} /home/mostendorff/experiments/eulm/data/ele/pl/curlicat_pl;
    """

    DATASET_ID = "curlicat_pl"
    TITLE = "CURLICAT Polish Corpus"

    HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-polish-corpus/f63ae912553911ed9c1a00155d02670648c0a234e0314895b52169af2af57dd7/"  # noqa

    LANGUAGES = ["pl"]
    TOKENS = 59_301_782

    # done


class CurlicatRODataset(CurlicatBaseDataset):
    DATASET_ID = "curlicat_ro"
    TITLE = "CURLICAT Romanian Corpus"
    DESCRIPTION = """The corpus contains 26,477 files, which represent our contribution to the CURLICAT project. It contains texts from 7 domains: science, politics, culture, economy, health, education, nature. Each file has multiple levels of annotation: tokenized, lemmatized, morphologically annotated, dependency parsed, named entities, nominal phrases, IATE terms and automatic domain-specific terms were identified as well. All processing tools are available within the RELATE platform. The corpus was automatically anonymized. Alternate download location: https://relate.racai.ro/resources/curlicat/"""  # noqa

    HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-romanian-corpus/8b6c8dca58ea11ed9c1a00155d026706fb03ef8b4c1847cfbe9cea869a82731e/"  # noqa

    LANGUAGES = ["ro"]  # TODO Romanian; Moldavian; Moldovan (ro)

    TOKENS = 94_925_454

    DOWNLOAD_URLS = ["https://relate.racai.ro/resources/curlicat/ro_curlicat_20221031.zip"]

    # See https://stackoverflow.com/questions/17912307/u-ufeff-in-python-string
    conllup_encoding = "utf-8-sig"
    # done


class CurlicatSKDataset(CurlicatBaseDataset):
    # WARNING: First document is very large! => http://bur.sk/share/2015/RozumnostAZ.pdf (536246 tokens)
    # https://www.juls.savba.sk/curlicat_en.html

    DATASET_ID = "curlicat_sk"
    TITLE = "CURLICAT Slovak Corpus (3rd version)"

    HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-slovak-corpus-v10/b419d7086ef411ed9c1a00155d0267066a930aa487824c08ba48f1183e993aca/"  # noqa

    LANGUAGES = ["sk"]

    DOWNLOAD_URLS = ["https://www.juls.savba.sk/data/curlicat/curlicat-sk-20221025-v1.0.conllup.xz"]

    TOKENS = 67_000_000

    def decompress(self):
        # xz --decompress  curlicat-sk-20221025-v1.0.conllup.xz
        pass

    # done


class CurlicatSLDataset(CurlicatBaseDataset):
    DATASET_ID = "curlicat_sl"
    TITLE = "CURLICAT Slovenian Corpus"
    HOMEPAGE = "https://elrc-share.eu/repository/browse/curlicat-slovenian-corpus/e549f298590711ed9c1a00155d026706db0d61d46f294d9a821307cf9c5df245/"  # noqa
    LANGUAGES = ["sl"]
    TOKENS = 43_481_563

    # done

src/lm_datasets/datasets/multilingual/curlicat.py
malteos/lm-datasets
import os
from lm_datasets.datasets.base import Availability
from lm_datasets.datasets.jsonl_dataset import JSONLDataset

# From https://docs.google.com/spreadsheets/d/1_rfLKa_Kq09YI0BPnfmoSL6-U3SHrW8tRwmk3-Qchzo/edit#gid=1698256287
RAW_LANG_TO_TOKENS = """bg	398067053
cs	471961631
da	671484862
de	695512401
el	696216541
en	769465561
es	725125274
et	328068754
fi	404265224
fr	828959218
ga	65030095
hr	258816068
hu	375253894
it	768605772
lt	364361783
lv	363239195
mt	367834815
nl	770312808
pl	406648795
pt	675152149
ro	415038571
sk	392235510
sl	394814289
sv	500085970"""


class EURLexBaseDataset(JSONLDataset):
    """
    Read preprocessed JSONL files + token counts (see EULM repo: extract_text_eurlex.py)
    """

    SOURCE_ID = "eurlex"
    DESCRIPTION = "EurlexResources: A Corpus Covering the Largest EURLEX Resources."
    HOMEPAGE = "https://huggingface.co/datasets/joelito/eurlex_resources"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LOCAL_DIRS = ["pegasus:/netscratch/mostendorff/experiments/eulm/data/docs_by_language"]

    def get_raw_jsonl_paths(self):
        dataset_dir = self.get_local_dataset_dir()

        return [os.path.join(dataset_dir, lang, "eurlex.jsonl") for lang in self.LANGUAGES]


def get_eurlex_auto_cls_by_language(lang, tokens):
    class EURLexLanguageDataset(EURLexBaseDataset):
        TOKENS = tokens
        DATASET_ID = "eurlex_" + lang
        TITLE = f"EURLex ({lang} subset)"
        LANGUAGES = [lang]

    return EURLexLanguageDataset


def get_eurlex_auto_classes():
    """
    Generate automatically dataset classes with token count
    """
    pass

    lang_to_tokens = {row.split("\t")[0]: int(row.split("\t")[1]) for row in RAW_LANG_TO_TOKENS.splitlines()}

    return [
        get_eurlex_auto_cls_by_language(lang, tokens * 10)  # TODO token count like OSCAR based on ten CC dumps
        for lang, tokens in lang_to_tokens.items()
    ]

src/lm_datasets/datasets/multilingual/eurlex.py
malteos/lm-datasets
from lm_datasets.datasets.base import Availability
from lm_datasets.datasets.hf_dataset import HFDataset


# From https://docs.google.com/spreadsheets/d/1_rfLKa_Kq09YI0BPnfmoSL6-U3SHrW8tRwmk3-Qchzo/edit#gid=1206328236
RAW_LANG_TO_TOKENS = """bg	2390349
cs	1840827375
da	10466716
de	6184578784
el	1155977
en	966539309
es	9058939804
et	110198368
fi	62799074
fr	2117306229
ga	32772
hu	244911748
it	3053920779
lt	9142223
lv	58702
mt	3479869
nl	21962633
pl	2235839721
pt	1338147828
ro	551372510
sk	349265172
sl	107493024
sv	328471555"""


class LegalMC4BaseDataset(HFDataset):
    SOURCE_ID = "legal_mc4"
    DESCRIPTION = "MC4_Legal: A Corpus Covering the Legal Part of MC4 for European Languages"
    HOMEPAGE = "https://huggingface.co/datasets/joelito/legal-mc4"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
    WEB_CRAWLED = True
    # DUMMY = True

    HF_DATASET_ID = "joelito/legal-mc4"
    HF_DATASET_SPLIT = "train"
    HF_DATASET_CONFIGS = None  # is set by language version

    streaming = True


def get_legal_mc4_auto_cls_by_language(lang, tokens):
    class LegalMC4LanguageDataset(LegalMC4BaseDataset):
        TOKENS = tokens
        DATASET_ID = "legal_mc4_" + lang
        TITLE = "legal_mc4_" + lang
        LANGUAGES = [lang]

        @property
        def HF_DATASET_CONFIGS(self):
            return [lang]

    return LegalMC4LanguageDataset


def get_legal_mc4_auto_classes():
    """
    Auto generate dataset classes with token count
    """
    lang_to_tokens = {row.split("\t")[0]: int(row.split("\t")[1]) for row in RAW_LANG_TO_TOKENS.splitlines()}

    return [get_legal_mc4_auto_cls_by_language(lang, tokens) for lang, tokens in lang_to_tokens.items()]

src/lm_datasets/datasets/multilingual/legal_mc4.py
malteos/lm-datasets
from io import TextIOWrapper
import logging
from typing import Set
import zipfile
from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning


logger = logging.getLogger(__name__)


class MacocuBaseDataset(BaseDataset):
    SOURCE_ID = "macocu"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
    WEB_CRAWLED = True

    languages_needed: Set[str] = None  # hbs_lat,en,hbs_cyr

    def is_downloaded(self):
        return len(self.get_dataset_file_paths(needed_suffix=".zip")) == len(self.DOWNLOAD_URLS)

    def get_texts(self):
        from lm_datasets.io.prevert_file import PrevertFile

        if not self.is_downloaded():
            self.download()

        if self.languages_needed is None:
            raise ValueError("languages_needed is not set")

        archive_path = self.get_dataset_file_paths(single_file=True, needed_suffix=".zip")

        logger.info(f"Extracting from {archive_path} ...")

        with zipfile.ZipFile(archive_path) as zf:
            member_fns = zf.namelist()
            for fn in member_fns:
                if fn.endswith(".xml"):
                    with zf.open(fn) as member_f:
                        logger.info(f"Parsing {fn} ...")

                        dset = PrevertFile(TextIOWrapper(member_f))

                        for doc in dset:  # iterating through documents of a dataset
                            lang = eval(doc.meta["lang_distr"])[0][0]  # hbs_lat,en,hbs_cyr
                            self.counter.update({f"lang_{lang}": 1})
                            text = str(doc)

                            if self.languages_needed is not None and lang in self.languages_needed:
                                yield text


class MacocuBGDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_bg"
    TITLE = "Bulgarian web corpus MaCoCu-bg 2.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1800"

    LANGUAGES = ["bg"]

    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1800/MaCoCu-bg-2.0.xml.zip"]

    TOKENS = 3_506_223_084

    languages_needed = {"bg"}


class MacocuHRDataset(MacocuBaseDataset):
    """
    Partially bad quality

    TODO maybe encoding error
    """

    DATASET_ID = "macocu_hr"
    TITLE = "Croatian MaCoCu corporus"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1806"

    LANGUAGES = ["hr"]

    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1806/MaCoCu-hr-2.0.xml.zip"]

    USED_BY = []

    TOKENS = 2_363_710_130

    QUALITY_WARNINGS = [QualityWarning.BAD_ENCODING]

    languages_needed = {"hr", "hbs_lat", "hbs_cyr"}


class MacocuELDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_el"
    TITLE = "Greek web corpus MaCoCu-el 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1839"
    LANGUAGES = ["el"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1839/MaCoCu-el-1.0.xml.zip"]
    TOKENS = 4_384_614_674

    languages_needed = {"el"}


class MacocuSQDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_sq"
    TITLE = "Albanian web corpus MaCoCu-sq 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1804"
    LANGUAGES = ["sq"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1804/MaCoCu-sq-1.0.xml.zip"]
    TOKENS = 625_726_547

    languages_needed = {"sq"}


class MacocuBSDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_bs"
    TITLE = "Bosnian web corpus MaCoCu-bs 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1808"
    LANGUAGES = ["bs"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1808/MaCoCu-bs-1.0.xml.zip"]
    TOKENS = 730342880

    languages_needed = {"bs"}


class MacocuCADataset(MacocuBaseDataset):
    DATASET_ID = "macocu_ca"
    TITLE = "Catalan web corpus MaCoCu-ca 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1837"
    LANGUAGES = ["ca"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1837/MaCoCu-ca-1.0.xml.zip"]
    TOKENS = 1736578493

    languages_needed = {"ca"}


class MacocuISDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_is"
    TITLE = "Icelandic web corpus MaCoCu-is 2.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1805"
    LANGUAGES = ["is"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1805/MaCoCu-is-2.0.xml.zip"]
    TOKENS = 887198687

    languages_needed = {"is"}


class MacocuMKDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_mk"
    TITLE = "Macedonian web corpus MaCoCu-mk 2.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1801"
    LANGUAGES = ["mk"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1801/MaCoCu-mk-2.0.xml.zip"]
    TOKENS = 524069254

    languages_needed = {"mk"}


class MacocuMTDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_mt"
    TITLE = "Maltese web corpus MaCoCu-mt 2.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1803"
    LANGUAGES = ["mt"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1803/MaCoCu-mt-2.0.xml.zip"]
    TOKENS = 347_855_619

    languages_needed = {"mt"}


class MacocuCNRDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_cnr"
    TITLE = "Montenegrin web corpus MaCoCu-cnr 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1809"
    LANGUAGES = ["cnr"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1809/MaCoCu-cnr-1.0.xml.zip"]
    TOKENS = 161361667

    languages_needed = {"cnr"}


class MacocuSRDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_sr"
    TITLE = "Serbian web corpus MaCoCu-sr 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1807"
    LANGUAGES = ["sr"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1807/MaCoCu-sr-1.0.xml.zip"]
    TOKENS = 2491019884

    languages_needed = {"hbs_lat", "hbs_cyr", "sr"}


class MacocuSLDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_sl"
    TITLE = "Slovene web corpus MaCoCu-sl 2.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1795"
    LANGUAGES = ["sl"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1795/MaCoCu-sl-2.0.xml.zip"]
    TOKENS = 1_920_089_135

    languages_needed = {"sl"}


class MacocuTRDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_tr"
    TITLE = "Turkish web corpus MaCoCu-tr 2.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1802"
    LANGUAGES = ["tr"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1802/MaCoCu-tr-2.0.xml.zip"]
    TOKENS = 4344850253

    languages_needed = {"tr"}


class MacocuUKDataset(MacocuBaseDataset):
    DATASET_ID = "macocu_uk"
    TITLE = "Ukrainian web corpus MaCoCu-uk 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1838"
    LANGUAGES = ["uk"]
    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1838/MaCoCu-uk-1.0.xml.zip"]
    TOKENS = 6181945683

    languages_needed = {"uk"}

src/lm_datasets/datasets/multilingual/macocu.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, Availability

# From https://docs.google.com/spreadsheets/d/1_rfLKa_Kq09YI0BPnfmoSL6-U3SHrW8tRwmk3-Qchzo/edit#gid=430577929
OSCAR_RAW_LANG_TO_TOKENS = """af	6,217,024
am	40,262,809
an	264
ar	10,081,452,882
arz	385,511
as	24,395,215
ast	10,917
av	1,073
az	316,850,330
azb	2,029,729
ba	26,036,637
be	164,729,607
bg	3,635,273,738
bh	507
bn	1,092,983,765
bo	6,062,558
bpy	346,947
br	4,759,407
bs	395
bxr	698
ca	2,240,460,836
ce	1,051,752
ceb	6,263,404
ckb	61,334,746
cs	9,717,378,559
cv	3,039,925
cy	52,250,043
da	2,217,634,340
de	73,848,586,648
dsb	84
dv	10,655,359
el	7,691,622,692
en	523,869,288,690
eo	67,774,923
es	63,388,237,965
et	938,296,892
eu	136,672,615
fa	9,609,206,698
fi	4,198,143,883
fr	62,127,088,294
fy	9,885,788
ga	9,054,923
gd	18,458
gl	38,345,625
gn	260
gom	121,035
gsw	34,328
gu	417,001,705
he	1,697,158,891
hi	2,475,605,444
hr	3,542,961
hsb	15,827
ht	20,671
hu	16,013,364,289
hy	336,045,041
ia	9,384
id	3,228,020,221
ie	0
ilo	4,411
io	2,598
is	294,471,539
it	36,327,274,203
ja	4,401,059,165
jbo	312,315
jv	3,286
ka	373,935,158
kk	214,679,857
km	59,880,231
kn	124,924,350
ko	3,435,866,935
krc	8,385
ku	25,921,607
kv	5,909
kw	80
ky	32,062,783
la	307,865
lb	2,514,838
lez	60,634
li	169
lmo	3,269
lo	10,659,203
lt	1,674,362,574
lv	845,459,899
mai	467
mg	1,416,430
mhr	1,615,215
min	305,589
mk	389,344,425
ml	236,597,838
mn	454,350,415
mr	252,706,331
mrj	60,180
ms	238,477
mt	149,886
multi	1,251,676,406
mwl	54
my	82,433,836
mzn	16,115
nah	279
nds	1,612,175
ne	278,901,036
new	229,703
nl	19,564,553,306
nn	575,518
no	373,160,033
oc	34,701
or	31,963,340
os	3,935,964
pa	104,235,418
pl	18,073,705,588
pms	510,087
pnb	11,844,695
ps	30,196,179
pt	15,172,557,311
qu	13
ro	6,302,600,833
ru	78,032,029,344
sa	2,479,345
sah	4,288,051
sd	14,667,207
sh	166,517
si	172,755,385
sk	2,704,716,280
sl	192,816,743
so	51
sq	462,694,599
sr	632,781,822
su	258
sv	6,993,719,601
sw	164,459
ta	738,824,392
te	201,575,815
tg	76,987,285
th	2,224,483,018
tk	325,786
tl	110,560,444
tr	8,290,890,087
tt	59,253,765
ug	14,659,554
uk	3,183,842,018
ur	434,023,273
uz	1,665,960
vi	22,424,984,210
vo	49,968
wa	6,347
war	19,665
wuu	1,199
x-eml	329
xal	27
xmf	283,316
yi	14,287,370
yo	2,396
zh	44,378,380,161"""

RAW_FILTER_LANGUAGES = """bg
cs
da
de
el
en
es
et
fi
fr
ga
hr
hu
it
lt
lv
mt
nl
pl
pt
ro
sk
sl
sv
code
uk
sr
sh
nn
no"""


class OscarBaseDataset(BaseDataset):
    SOURCE_ID = "oscar"
    DESCRIPTION = (
        "The OSCAR project (Open Super-large Crawled Aggregated coRpus) is an Open Source project aiming to provide"
        " web-based multilingual resources and datasets for Machine Learning (ML) and Artificial Intelligence (AI)"
        " applications."
    )
    HOMEPAGE = "https://oscar-project.org/"
    AVAILIBILITY = Availability.ON_REQUEST
    WEB_CRAWLED = True
    DUMMY = True


def get_oscar_dummy_cls_by_language(lang, tokens):
    class OscarLanguageDataset(OscarBaseDataset):
        TOKENS = tokens
        DATASET_ID = "oscar_" + lang
        TITLE = "oscar_" + lang
        LANGUAGES = [lang]

    return OscarLanguageDataset


def get_oscar_dummy_classes():
    """
    Generate dummy dataset classes with token count based on OSCAR 23.01
    """
    pass

    lang_to_tokens = {
        row.split("\t")[0]: int(row.split("\t")[1].replace(",", "")) for row in OSCAR_RAW_LANG_TO_TOKENS.splitlines()
    }

    filter_langs = set(RAW_FILTER_LANGUAGES.splitlines())

    return [
        get_oscar_dummy_cls_by_language(lang, tokens * 10)  # TODO token count like OSCAR based on ten CC dumps
        for lang, tokens in lang_to_tokens.items()
        if lang in filter_langs
    ]

src/lm_datasets/datasets/multilingual/oscar_dummy.py
malteos/lm-datasets
import os
from lm_datasets.datasets.base import Availability, BILLION
from lm_datasets.datasets.hf_dataset import HFDataset


class RedPajamaBaseDataset(HFDataset):
    """
    # full dataset
    wget 'https://data.together.xyz/redpajama-data-1T/v1.0.0/urls.txt'
    while read line; do
        dload_loc=${line#https://data.together.xyz/redpajama-data-1T/v1.0.0/}
        mkdir -p $(dirname $dload_loc)
        wget "$line" -O "$dload_loc"
    done < urls.txt

    # filter urls (books, arxiv, stackexnchange)
    cat urls.txt | grep -E "book|arxiv|stackexchange" > filtered_urls.txt
    while read line; do
        dload_loc=${line#https://data.together.xyz/redpajama-data-1T/v1.0.0/}
        mkdir -p $(dirname $dload_loc)
        wget "$line" -O "$dload_loc"
    done < filtered_urls.txt

    # checksums
    wget https://data.together.xyz/redpajama-data-1T/v1.0.0/sha256/arxiv_SHA256SUMS.txt
    wget https://data.together.xyz/redpajama-data-1T/v1.0.0/sha256/book_SHA256SUMS.txt
    wget https://data.together.xyz/redpajama-data-1T/v1.0.0/sha256/c4_SHA256SUMS.txt
    wget https://data.together.xyz/redpajama-data-1T/v1.0.0/sha256/common_crawl_SHA256SUMS.txt
    wget https://data.together.xyz/redpajama-data-1T/v1.0.0/sha256/github_SHA256SUMS.txt
    wget https://data.together.xyz/redpajama-data-1T/v1.0.0/sha256/stackexchange_SHA256SUMS.txt
    wget https://data.together.xyz/redpajama-data-1T/v1.0.0/sha256/wikipedia_SHA256SUMS.txt

    sha256sum --check stackexchange_SHA256SUMS.txt
    sha256sum --check book_SHA256SUMS.txt
    sha256sum --check arxiv_SHA256SUMS.txt

    """

    SOURCE_ID = "redpajama"

    TITLE = "RedPajama T1"
    HOMEPAGE = "https://huggingface.co/datasets/togethercomputer/RedPajama-Data-1T"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/multilingual/redpajama"]

    HF_DATASET_ID = "togethercomputer/RedPajama-Data-1T"
    HF_DATASET_SPLIT = "train"

    text_column_name = "text"

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

        # set download diretory
        os.environ["RED_PAJAMA_DATA_DIR"] = self.get_local_dataset_dir()


class RedPajamaBookDataset(RedPajamaBaseDataset):
    DATASET_ID = "redpajama_book"
    HF_DATASET_CONFIGS = ["book"]
    TOKENS = 26 * BILLION

    LANGUAGES = ["en"]


class RedPajamaArxivDataset(RedPajamaBaseDataset):
    DATASET_ID = "redpajama_arxiv"
    HF_DATASET_CONFIGS = ["arxiv"]
    TOKENS = 28 * BILLION

    LANGUAGES = ["en"]


class RedPajamaStackexchangeDataset(RedPajamaBaseDataset):
    DATASET_ID = "redpajama_stackexchange"
    HF_DATASET_CONFIGS = ["stackexchange"]
    TOKENS = 20 * BILLION

    LANGUAGES = ["en"]


# not implemented: commoncrawl, c4, wikipedia, github

src/lm_datasets/datasets/multilingual/redpajama.py
malteos/lm-datasets
import os
import logging
from lm_datasets.datasets.base import Availability, BaseDataset, TOKENS_PER_BYTE
from lm_datasets.utils.wikimedia import parse_and_clean_wikicode


import xml.etree.cElementTree as etree

from smart_open import open

logger = logging.getLogger(__name__)


WIKI_TYPES = [
    "wiki",  # Wikipedia
    "wikibooks",  # Wikibooks
    "wikiquote",  # Wikiquote
    "wikinews",  # Wikinews
    "wikisource",  # Wikisource
    "wikivoyage",  # Wikivoyage
]


# stats_fp = "data/docs_by_language/stats.json"

# stats = json.load(open(stats_fp))

# bytes_per_language_and_source = {
#     lang: {s: b for s, b in source_bytes.items() if s in WIKI_TYPES}
#     for lang, source_bytes in stats["bytes_per_language_and_source"].items()
# }

bytes_per_language_and_source = {
    # EU24 languages
    "bg": {
        "wiki": 1214121628,
        "wikibooks": 10017340,
        "wikiquote": 37156441,
        "wikinews": 4227797,
        "wikisource": 86661269,
    },
    "cs": {
        "wiki": 930310918,
        "wikibooks": 13039351,
        "wikiquote": 8167752,
        "wikinews": 5203801,
        "wikisource": 259915843,
    },
    "da": {"wiki": 224403313, "wikibooks": 21476337, "wikiquote": 1034165, "wikisource": 20123952},
    "de": {
        "wiki": 5197735740,
        "wikibooks": 169962198,
        "wikiquote": 14929927,
        "wikinews": 29553128,
        "wikisource": 531482040,
        "wikivoyage": 100404450,
    },
    "el": {
        "wiki": 1992383542,
        "wikibooks": 65173299,
        "wikiquote": 15485796,
        "wikinews": 14364210,
        "wikisource": 558477104,
        "wikivoyage": 11892155,
    },
    "en": {
        "wiki": 12573466166,
        "wikibooks": 438671226,
        "wikiquote": 425528832,
        "wikinews": 50512180,
        "wikisource": 2491597933,
        "wikivoyage": 163322890,
    },
    "es": {
        "wiki": 3635029311,
        "wikibooks": 83041289,
        "wikiquote": 18004839,
        "wikinews": 23282513,
        "wikisource": 382827125,
        "wikivoyage": 46463569,
    },
    "et": {"wiki": 208343404, "wikibooks": 3589699, "wikiquote": 42857536, "wikisource": 4411631},
    "fi": {
        "wiki": 466350200,
        "wikibooks": 15630626,
        "wikiquote": 8245391,
        "wikinews": 2549156,
        "wikisource": 62883499,
        "wikivoyage": 4194479,
    },
    "fr": {
        "wiki": 4968876644,
        "wikibooks": 82544091,
        "wikiquote": 1613702,
        "wikinews": 25773373,
        "wikisource": 130101343,
        "wikivoyage": 24004716,
    },
    "ga": {"wiki": 21473185, "wikibooks": 0, "wikiquote": 797},
    "hr": {"wiki": 222907670, "wikibooks": 1835214, "wikiquote": 3393450, "wikisource": 68521110},
    "hu": {
        "wiki": 1050228505,
        "wikibooks": 64029424,
        "wikiquote": 8793400,
        "wikinews": 1453968,
        "wikisource": 121765396,
    },
    "it": {
        "wiki": 2798247221,
        "wikibooks": 99975368,
        "wikiquote": 188532632,
        "wikinews": 17511625,
        "wikisource": 225622077,
        "wikivoyage": 44689019,
    },
    "lt": {"wiki": 114077820, "wikibooks": 2025822, "wikiquote": 4709667, "wikisource": 9228644},
    "lv": {"wiki": 98460508, "wikibooks": 112635},
    "mt": {"wiki": 20129053},
    "nl": {
        "wiki": 895801299,
        "wikibooks": 28281089,
        "wikiquote": 255527,
        "wikinews": 8871342,
        "wikisource": 54423027,
        "wikivoyage": 10058074,
    },
    "pl": {
        "wiki": 1229781898,
        "wikibooks": 33877793,
        "wikiquote": 99998973,
        "wikinews": 29401114,
        "wikisource": 63420701,
        "wikivoyage": 30125414,
    },
    "pt": {
        "wiki": 1588856067,
        "wikibooks": 45431838,
        "wikiquote": 24565736,
        "wikinews": 53043049,
        "wikisource": 119019274,
        "wikivoyage": 10344423,
    },
    "ro": {
        "wiki": 518232680,
        "wikibooks": 3900601,
        "wikiquote": 1486942,
        "wikinews": 2536562,
        "wikisource": 166964076,
        "wikivoyage": 1727955,
    },
    "sk": {"wiki": 216967151, "wikibooks": 6584629, "wikiquote": 4781513, "wikisource": 4674463},
    "sl": {"wiki": 264134366, "wikibooks": 5769223, "wikiquote": 2279993, "wikisource": 401940397},
    "sv": {
        "wiki": 443696163,
        "wikibooks": 9014999,
        "wikiquote": 1298260,
        "wikinews": 1840363,
        "wikisource": 32096731,
        "wikivoyage": 4359374,
    },
    # Additional languages
    "uk": {
        "wiki": 1,
        "wikibooks": 1,
        "wikiquote": 1,
        "wikinews": 1,
        "wikisource": 1,
        "wikivoyage": 1,
    },
    "sr": {
        "wiki": 1,
        "wikibooks": 1,
        "wikiquote": 1,
        "wikinews": 1,
        "wikisource": 1,
    },
    "sh": {
        "wiki": 1,
    },
    "nn": {
        "wiki": 1,
        "wikiquote": 1,
    },
    "no": {
        "wiki": 1,
        "wikibooks": 1,
        "wikiquote": 1,
        "wikinews": 1,
        "wikisource": 1,
        "wikivoyage": 1,
    },
    "eu": {
        "wiki": 1,
        "wikibooks": 1,
        "wikiquote": 1,
        "wikinews": 1,
        "wikisource": 1,
    },
    "ca": {
        "wiki": 1,
        "wikibooks": 1,
        "wikiquote": 1,
        "wikinews": 1,
        "wikisource": 1,
        "wikivoyage": 1,
    },
    "gl": {
        "wiki": 1,
        "wikibooks": 1,
        "wikiquote": 1,
        "wikinews": 1,
        "wikisource": 1,
    },
}


class WikimediaBaseDataset(BaseDataset):
    DESCRIPTION = "Wikimedia dumps."
    HOMEPAGE = "https://wikimedia.org"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD
    # DUMMY = True
    SOURCE_ID = None  # wiki,wikibooks,wikiquote,wikinews,wikisource,wikivoyage

    download_base_url = "https://dumps.wikimedia.org/"
    dump_date = "20230801"

    @property
    def DOWNLOAD_URLS(self):
        dump_url = (
            f"{self.download_base_url}{self.LANGUAGES[0]}{self.SOURCE_ID}/{self.dump_date}/{self.get_raw_file_name()}"
        )

        return [dump_url]

    def get_raw_file_name(self) -> str:
        return f"{self.LANGUAGES[0]}{self.SOURCE_ID}-{self.dump_date}-pages-articles.xml.bz2"

    def get_raw_file_path(self):
        return os.path.join(self.get_local_dataset_dir(), self.get_raw_file_name())

    def is_downloaded(self):
        return os.path.exists(self.get_raw_file_path())

    def get_texts(self):
        import mwparserfromhell

        if not self.is_downloaded():
            self.download()

        # Parse wiki dump
        archive_path = self.get_raw_file_path()

        logger.info(f"Reading from {archive_path}")

        with open(archive_path, "rb") as xml_f:
            # Iterate over XML file
            context = etree.iterparse(xml_f, events=("end",))

            for unused_event, elem in context:
                if not elem.tag.endswith("page"):
                    continue

                namespace = elem.tag[:-4]
                title = elem.find(f"./{namespace}title").text
                ns = elem.find(f"./{namespace}ns").text
                # id_ = elem.find(f"./{namespace}id").text
                red_ = elem.find(f"./{namespace}redirect")

                # Filter pages that are not in the "main" namespace.
                if ns != "0":
                    elem.clear()
                    continue

                raw_content = elem.find(f"./{namespace}revision/{namespace}text").text
                elem.clear()

                if raw_content is None or red_ is not None:
                    self.counter.update({"wiki_redirect": 1})
                    continue

                if "#REDIRECT" in raw_content:
                    self.counter.update({"wiki_redirect": 1})
                    continue

                try:
                    text = parse_and_clean_wikicode(
                        raw_content, parser=mwparserfromhell, language=self.get_language_code()
                    )
                except mwparserfromhell.parser.ParserError as e:
                    logger.error(f"mwparserfromhell ParseError: {e}")

                    self.counter.update({"error_mwparserfromhell": 1})
                    continue

                except Exception as e:
                    logger.error(f"Error: {e}")

                    self.counter.update({"error_other": 1})
                    continue

                plain_text = title + self.title_delimiter + text

                yield plain_text


def get_wikimedia_auto_cls_by_language(lang, source, bytes):
    class WikimediaLanguageDataset(WikimediaBaseDataset):
        SOURCE_ID = source
        TOKENS = int(TOKENS_PER_BYTE * bytes)
        BYTES = bytes
        DATASET_ID = source + "_" + lang
        TITLE = source + "_" + lang
        LANGUAGES = [lang]

    return WikimediaLanguageDataset


def get_wikimedia_auto_classes():
    """
    Auto generate dataset classes with token count
    """

    return [
        get_wikimedia_auto_cls_by_language(lang, s, b)
        for lang, source_bytes in bytes_per_language_and_source.items()
        for s, b in source_bytes.items()
    ]

src/lm_datasets/datasets/multilingual/wikimedia.py
malteos/lm-datasets
import logging
import os

from multiprocessing.pool import Pool

from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning
from lm_datasets.utils import remove_whitespaces_before_punctuation


logger = logging.getLogger(__name__)

try:
    import folia.main as folia
except ImportError:
    logger.error("Cannot load imports")


class SonarBaseDataset(BaseDataset):
    SOURCE_ID = "sonar"

    TITLE = "SoNaR Corpus NC 1.2"
    HOMEPAGE = "https://taalmaterialen.ivdnt.org/download/tstc-sonar-corpus/"
    DESCRIPTION = (
        "The SoNaR Corpus contains more than 500 million words from texts in standard Dutch ",
        "later than 1954. All texts were tokenized, tagged for part of speech and lemmatized. ",
        "The named entities were also labelled. All annotations were produced automatically, ",
        "no manual verification took place.",
    )

    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    LANGUAGES = ["nl"]

    USED_BY = ["GroNLP/bert-base-dutch-cased"]

    TOKENS = 0
    # at least 1,030,000 docs

    QUALITY_WARNINGS = [QualityWarning.BAD_WHITESPACES]

    # LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/nl/sonar/SONAR500/FoLiA/WR-P-E-A_discussion_lists"]
    # subsets = ["WR-P-E-J_wikipedia"]

    # # Excluded
    # subsets = [
    #     # "WR-P-E-A_discussion_lists",  # very short text, too many line breaks and whitespaces
    #     # "WR-P-E-J_wikipedia",  # exclude wikipedia (already part of other dataset)
    # ]

    def get_local_dataset_dir(self):
        return os.path.join(self.raw_datasets_dir, self.get_language_code(), self.SOURCE_ID)

    def decompress(self):
        # tgz
        pass

    def get_text_from_item(self, item) -> str:
        """
        item: file path
        """

        doc = folia.Document(file=item)
        text = doc.text()
        text = remove_whitespaces_before_punctuation(text)

        # logger.info(f"Text extracted {item}")

        return text

    def get_file_paths(self):
        dataset_dir = self.get_local_dataset_dir()

        # Iterate over subsets
        for subset in self.subsets:
            logger.info(subset)

            subset_dir = os.path.join(dataset_dir, "SONAR500", "FoLiA", subset)

            fns = os.listdir(subset_dir)

            # Some subsets have files in subdirectories
            folia_fns = [fn for fn in fns if fn.endswith(".folia.xml")]

            if folia_fns:
                for fn in folia_fns:
                    yield os.path.join(subset_dir, fn)
            else:
                # subdirectories
                for fn in fns:
                    dir_path = os.path.join(subset_dir, fn)
                    if os.path.isdir(dir_path):
                        for fn in os.listdir(dir_path):
                            if fn.endswith(".folia.xml"):
                                yield os.path.join(dir_path, fn)

    def get_texts(self):
        pool = Pool(processes=self.workers)
        iterable = pool.imap_unordered(self.get_text_from_item, self.get_file_paths())

        def terminate():
            pool.terminate()
            pool.close()
            pool.join()

        iterable.terminate = terminate

        return iterable

        # with Pool(processes=self.workers) as pool:
        #     yield from pool.imap_unordered(self.get_text_from_item, self.get_file_paths())

        # for fp in self.get_file_paths():
        #     yield self.get_text_from_item(fp)

    # def get_texts(self):
    #     # SLOW => because gzipped files need to be decompressed for members!
    #     # with tarfile.open("/netscratch/mostendorff/experiments/eulm/data/ele/nl/sonar.tgz") as tf:
    #     #     tfps = tf.getmembers()

    #     dataset_dir = self.get_local_dataset_dir()

    #     # Iterate over subsets
    #     for subset in self.subsets:
    #         logger.info(subset)

    #         subset_dir = os.path.join(dataset_dir, "SONAR500", "FoLiA", subset)

    #         fns = os.listdir(subset_dir)

    #         # Some subsets have files in subdirectories
    #         folia_fns = [fn for fn in fns if fn.endswith(".folia.xml")]

    #         if folia_fns:
    #             yield from self.get_texts_from_folia(folia_fns, subset_dir)
    #         else:
    #             # subdirectories
    #             for fn in fns:
    #                 dir_path = os.path.join(subset_dir, fn)
    #                 if os.path.isdir(dir_path):
    #                     folia_fns = [fn for fn in os.listdir(dir_path) if fn.endswith(".folia.xml")]

    #                     yield from self.get_texts_from_folia(folia_fns, dir_path)


class SonarWebDataset(SonarBaseDataset):
    DATASET_ID = "sonar_web"
    TOKENS = 500_000_000  # TODO all other subsets have zero tokens

    # Web
    subsets = [
        "WR-P-E-K_blogs",
        "WR-P-E-I_web_sites",
        "WR-P-E-H_teletext_pages",
    ]


class SonarSubtitlesDataset(SonarBaseDataset):
    DATASET_ID = "sonar_subtitles"

    # Subtitles
    subsets = [
        "WR-P-E-G_subtitles",
    ]


class SonarGovDataset(SonarBaseDataset):
    DATASET_ID = "sonar_gov"

    # Gov
    subsets = [
        "WR-P-P-I_policy_documents",
        "WR-P-P-F_legal_texts",
        "WR-P-P-J_proceedings",
        "WR-P-P-K_reports",
    ]


class SonarBooksDataset(SonarBaseDataset):
    DATASET_ID = "sonar_books"

    # Books
    subsets = [
        "WR-P-P-B_books",
        "WR-P-P-H_periodicals_magazines",
        "WR-P-E-C_e-magazines",
        "WR-P-P-C_brochures",
        "WS-U-E-A_auto_cues",
    ]


class SonarNewsDataset(SonarBaseDataset):
    DATASET_ID = "sonar_news"

    # News
    subsets = [
        "WR-P-E-E_newsletters",
        "WR-P-P-G_newspapers",
        "WR-P-E-F_press_releases",
        "WR-P-P-D_newsletters",
    ]


class SonarEduDataset(SonarBaseDataset):
    DATASET_ID = "sonar_edu"

    # Education
    subsets = [
        "WS-U-T-B_texts_for_the_visually_impaired",
        "WR-U-E-E_written_assignments",  # extra dataset
        "WR-P-P-E_guides_manuals",
    ]


def get_sonar_classes():
    # sonar_web,sonar_subtitles,sonar_gov,sonar_books,sonar_news,sonar_edu
    return [
        SonarSubtitlesDataset,
        SonarWebDataset,
        SonarGovDataset,
        SonarBooksDataset,
        SonarNewsDataset,
        SonarEduDataset,
    ]

src/lm_datasets/datasets/nl/sonar.py
malteos/lm-datasets
import logging


from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning

import zipfile


logger = logging.getLogger(__name__)


class SonarNewMediaDataset(BaseDataset):
    """ """

    DATASET_ID = "sonar_new_media"
    TITLE = "SoNaR Nieuwe Media Corpus (Version 1.0)"
    HOMEPAGE = "https://taalmaterialen.ivdnt.org/download/tstc-sonar-nieuwe-media-corpus-1/"
    DESCRIPTION = (
        "The SoNaR New Media Corpus 1.0 contains texts from new media (sms, tweets and ",
        "chat messages) that were collected within the STEVIN-project SoNaR.",
    )

    AVAILIBILITY = Availability.SIGNIN_DOWNLOAD

    QUALITY_WARNINGS = [QualityWarning.SHORT_TEXT]

    LANGUAGES = ["nl"]

    TOKENS = 36_000_000

    def get_texts(self):
        from bs4 import BeautifulSoup

        # import folia.main as folia

        # read directly from zip file
        zip_fp = self.get_dataset_file_paths(needed_suffix=".zip", single_file=True)

        with zipfile.ZipFile(zip_fp) as zf:
            needed_suffix = ".folia.xml"
            members_fns = [fn for fn in zf.namelist() if fn.endswith(needed_suffix)]

            for member_fn in members_fns:
                # tweets,chat,sms
                # if "tweets" in member_fn:
                #     continue

                if "chat" in member_fn:
                    # only very short messages
                    continue

                if "sms" in member_fn:
                    # bad quality
                    continue

                logger.info(f"Extracting from {member_fn}")

                with zf.open(member_fn) as member_f:
                    soup = BeautifulSoup(member_f, features="lxml")

                for tweet in soup.find_all("event"):
                    t = tweet.find("t")

                    if t:
                        text = t.get_text()

                        # replace _ with white space
                        text = text.replace("_", " ")

                        yield text

        logger.info("ZIP file extracted")

src/lm_datasets/datasets/nl/sonar_new_media.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, BILLION


class Maalfrid2021Dataset(BaseDataset):
    DATASET_ID = "maalfrid_2021"
    TITLE = "Målfrid 2021 - Freely Available Documents from Norwegian State Institutions"
    HOMEPAGE = "https://hdl.handle.net/21.11146/69"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["nb", "nn", "en", "se", "other"]

    DESCRIPTION = (
        "This corpus consists of documents from 339 internet domains of Norwegian state institutions and comprises"
        " approximately 4.1 billion tokens in total, which makes it one of the largest freely available resources for"
        " Norwegian Bokmål and Nynorsk. In addition to Norwegian, the corpus contains texts in Northern Sami, Lule"
        " Sami, Southern Sami and English."
    )
    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "open license"

    # Size: 4.1 billion tokens
    TOKENS = 4.1 * BILLION

    DOWNLOAD_URLS = ["https://www.nb.no/sbfil/tekst/maalfrid_2021/maalfrid_2021.tar.gz"]
    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/maalfrid_2021"]

src/lm_datasets/datasets/no/maalfrid_2021.py
malteos/lm-datasets
import os
import tarfile
from lm_datasets.datasets.base import BaseDataset, BILLION, MILLION
from lm_datasets.systems import get_path_by_system


class NAKDataset(BaseDataset):
    DATASET_ID = "nak"
    TITLE = "Norwegian Newspaper Corpus"
    HOMEPAGE = "https://hdl.handle.net/21.11146/4"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["nb", "nn"]

    DESCRIPTION = (
        "This version of Norwegian Newspaper Corpus contains text from 1998 to 2019. The corpus contains approximately"
        " 1,68 billion words for Norwegian Bokmål, and about 68 million words for Norwegian Nynorsk."
    )
    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "non-commercial use"

    # Size: 1.68 billion words for Bokmål, 68 million words for Nynorsk
    TOKENS = 1.68 * BILLION + 68 * MILLION

    DOWNLOAD_URLS = [
        "https://www.nb.no/sbfil/tekst/nak_2019.tar",
        "https://www.nb.no/sbfil/tekst/nak_2018.tar",
        "https://www.nb.no/sbfil/tekst/nak_2017.tar",
        "https://www.nb.no/sbfil/tekst/nak_2016.tar",
        "https://www.nb.no/sbfil/tekst/nak_2015.tar",
        "https://www.nb.no/sbfil/tekst/nak_2014.tar",
        "https://www.nb.no/sbfil/tekst/nak_2013.tar",
        "https://www.nb.no/sbfil/tekst/nak_2012.tar",
    ]
    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/nak"]

    def get_dataset_file_paths(self):
        # TODO handle datasets with multiple filles
        dataset_dir = get_path_by_system(self.LOCAL_DIRS)
        fps = [os.path.join(dataset_dir, f) for f in os.listdir(dataset_dir)]
        fps = [fp for fp in fps if os.path.isfile(fp)]

        if len(fps) == 8:
            return fps
        else:
            raise ValueError(f"Cannot determine file path. Either none or wrong number of files in dataset dir: {fps=}")

    def decompress(self):
        dataset_dir = get_path_by_system(self.LOCAL_DIRS)
        dataset_file_paths = self.get_dataset_file_paths()

        decompressed_dir = os.path.join(dataset_dir, "decompressed")

        if not os.path.exists(decompressed_dir):
            os.makedirs(decompressed_dir)

        for fp in dataset_file_paths:
            with tarfile.open(fp) as file_handler:
                fdir, fn = os.path.split(fp)

                file_handler.extractall(path=os.path.join(decompressed_dir, fn))

    def extract_plaintext(self):
        self.decompress()

        return super().extract_plaintext()

src/lm_datasets/datasets/no/nak.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, GB


class NBDigitalDataset(BaseDataset):
    DATASET_ID = "nbdigital"
    TITLE = "Public Domain Texts from NBdigital"
    HOMEPAGE = "https://hdl.handle.net/21.11146/34"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["nb", "nn"]

    DESCRIPTION = (
        "This text collection consists of texts from the digital National Library which are in the public domain. The"
        " collection contains 26.344 books (and other written material) by 10756 different authors (including, e.g.,"
        " public institutions for publically available material)."
    )
    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "public domain"

    # Size: 26344 books
    BYTES = 2.7 * GB

    DOWNLOAD_URLS = ["https://www.nb.no/sbfil/tekst/20150526_nbdig_txt01.tar.gz"]
    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/multilingual"]

src/lm_datasets/datasets/no/nbdigital.py
malteos/lm-datasets
from lm_datasets.datasets.base import Availability
from lm_datasets.datasets.hf_dataset import HFDataset


class NorwegianCCBaseDataset(HFDataset):
    DATASET_ID = "norwegian_cc"
    SOURCE_ID = "norwegian_cc"

    TITLE = "Norwegian Colossal Corpus"
    HOMEPAGE = "https://huggingface.co/datasets/NbAiLab/NCC"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    DESCRIPTION = (
        "The Norwegian Colossal Corpus is a collection of multiple smaller Norwegian corpuses "
        "suitable for training large language models. We have done extensive cleaning on the datasets, "
        "and have made them available in a common format. The total size of the NCC is currently 45GB. "
        "Documents: 20,830,348; Words/document: 331"
    )

    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "mixed"

    HF_DATASET_ID = "NbAiLab/NCC"
    HF_DATASET_SPLIT = "train"

    streaming = False

    excluded_doc_types = {
        # exclude all Wikimedia sources to avoid duplicated content with the original Wikimedia dataset
        "wikipedia",
    }

    text_column_name = "text"
    remove_columns = ["id", "publish_year", "lang_fasttext_conf"]

    def get_filter_func(self):
        languages_set = set(self.LANGUAGES)

        def filter_func(example):
            return example["doc_type"] not in self.excluded_doc_types and example["lang_fasttext"] in languages_set

        return filter_func


class NorwegianCCNODataset(NorwegianCCBaseDataset):
    DATASET_ID = "norwegian_cc_no"
    LANGUAGES = ["no"]

    # total: BYTES = 45 * GB
    TOKENS = 5_050_752_505  # total: 6_905_570_165


class NorwegianCCNNDataset(NorwegianCCBaseDataset):
    DATASET_ID = "norwegian_cc_no"
    LANGUAGES = ["nn"]

    # total: BYTES = 45 * GB
    TOKENS = 299_753_996  # total: 6_905_570_165

src/lm_datasets/datasets/no/norwegian_cc.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, MILLION


class ParlaMintDataset(BaseDataset):
    DATASET_ID = "parlamint"
    TITLE = "Norwegian part of the Parlamint corpus"
    HOMEPAGE = "https://hdl.handle.net/21.11146/77"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["no", "nn", "nb"]

    DESCRIPTION = (
        "ParlaMint-NO contains the Norwegian part of the ParlaMint project, an EU-funded project supported by CLARIN"
        " ERIC. The project’s aim is to create comparable and similarly annotated corpora of parliamentary meeting"
        " minutes."
    )
    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "public domain"

    # Size: 975 million tokens
    TOKENS = 975 * MILLION

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/parlamint"]

    DUMMY = True

src/lm_datasets/datasets/no/parlamint.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, BILLION


class ParliamentaryProceedingsDataset(BaseDataset):
    DATASET_ID = "parliamentary_proceedings"
    TITLE = "Norwegian Parliamentary Proceedings 1814-2000"
    HOMEPAGE = "https://hdl.handle.net/21.11146/74"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["no", "nb", "nn"]

    DESCRIPTION = (
        "This corpus contains published historical proceedings from the Norwegian parliament 1814-2000. A total of 2136"
        " volumes were digitized, OCR-read and processed at the National Library of Norway, and made available online"
        " in 2014."
    )
    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "open license"

    # Size: 1.5 billion tokens
    TOKENS = 1.5 * BILLION

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/parliamentary_proceedings"]

    DUMMY = True

src/lm_datasets/datasets/no/parliamentary_proceedings.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, MILLION


class SakspapirNNODataset(BaseDataset):
    DATASET_ID = "sakspapir_nno"
    TITLE = "Legal Documents from Norwegian Nynorsk Municipialities"
    HOMEPAGE = "https://hdl.handle.net/21.11146/60"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["nn", "nb"]

    DESCRIPTION = (
        "The texts in this corpus have been collected with the web crawler Veidemann in collaboration with the National"
        " Library’s Web Archive, based on a revised list of municipalities from the National Association of Nynorsk"
        " Municipalities (see lnk.no)."
    )
    PII = "I have not checked the data source for personally identifiable or sensitive information."
    LICENSE = "open license"

    # Size: 127 million tokens
    TOKENS = 127 * MILLION

    DOWNLOAD_URLS = ["https://www.nb.no/sbfil/tekst/sakspapir_nno/sakspapir_nno_01.tar.gz"]
    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/no/sakspapir_nno"]

    DUMMY = True

src/lm_datasets/datasets/no/sakspapir_nno.py
malteos/lm-datasets
import logging
import tarfile
from typing import Iterable
from lm_datasets.datasets.base import BaseDataset


logger = logging.getLogger(__name__)


class LunaPL(BaseDataset):
    """
    Original text cannot be reconstructed from TEI XML.
    """

    DATASET_ID = "luna_pl"
    TITLE = "LUNA.PL corpus"
    DESCRIPTION = "Human-human dialogues in TEI P5 format annotated for the 'Spoken Language UNderstanding "
    +"in multilingual communication systems' project"
    DOWNLOAD_URLS = [
        ("http://zil.ipipan.waw.pl/LUNA?action=AttachFile&do=get&target=LUNA_TEI.tar.gz", "LUNA_TEI.tar.gz")
    ]
    LANGUAGES = ["pl"]

    DUMMY = True

    # BYTES = 932 * MB

    def decompress(self):
        # gzip -d LUNA_TEI.tar.gz
        pass

    def get_texts(self) -> Iterable[str]:
        raise ValueError("bad dataset")

        # tar + TEI
        # read from tar file
        archive_fp = self.get_dataset_file_paths(needed_suffix=".tar", single_file=True)

        logger.info(f"Extracting from {archive_fp}")

        with tarfile.open(archive_fp) as tar_f:
            member_fns = tar_f.getmembers()
            member_fns = [m for m in tar_f.getmembers() if m.name.endswith(".xml")]

            for member_fn in member_fns:
                member_content = tar_f.extractfile(member_fn).read()
                tei_xml = member_content.decode()
                # TODO extract plain text

                if tei_xml:
                    pass
                break

src/lm_datasets/datasets/pl/luna_pl.py
malteos/lm-datasets
import logging
import tarfile

import xml.etree.cElementTree as ET

from tqdm.auto import tqdm
from lm_datasets.datasets.base import MILLION, BaseDataset, QualityWarning

logger = logging.getLogger(__name__)


class NKJPPodkorpusMilionowyDataset(BaseDataset):
    """
    DOWNLOAD
    -----------

    Instruction

    - Downloaded locally by clicking on the download link in the browser:

    http://clip.ipipan.waw.pl/NationalCorpusOfPolish?action=AttachFile&do=get&target=NKJP-PodkorpusMilionowy-1.2.tar.gz

    - Copy local file to server:

    $ scp /Local/Path/to/NKJP-PodkorpusMilionowy-1.2.tar.gz \
        username@clustername:/data/datasets/ele/pl/pl_nkjp_1-2/NKJP-PodkorpusMilionowy-1.2.tar.gz

    - Extract files:

    tar -xvf NKJP-PodkorpusMilionowy-1.2.tar.gz

    - Run python script locally with small sample:

    $ python pl_nkjp_1-2.py \
        /Users/melinaplakidis/Downloads/NKJP-PodkorpusMilionowy-1.2 pl_nkjp_processed.jsonl \
        --limit 100

    """

    DATASET_ID = "pl_nkjp"
    TITLE = "NKJP-PodkorpusMilionowy-1.2 (National Corpus of Polish)"
    HOMEPAGE = "http://clip.ipipan.waw.pl/NationalCorpusOfPolish"

    LANGUAGES = ["pl"]
    DOWNLOAD_URLS = [
        (
            "http://clip.ipipan.waw.pl/NationalCorpusOfPolish?action=AttachFile&do=get&target=NKJP-PodkorpusMilionowy-1.2.tar.gz",  # noqa
            "NKJP-PodkorpusMilionowy-1.2.tar.gz",
        )
    ]
    TOKENS = 1 * MILLION
    LICENSE = "CC-BY"
    QUALITY_WARNINGS = [QualityWarning.BAD_LINEBREAKS]

    def decompress(self):
        # gzip -d NKJP-PodkorpusMilionowy-1.2.tar.gz
        pass

    def get_texts(self):
        # read from tar file
        archive_fp = self.get_dataset_file_paths(needed_suffix=".tar", single_file=True)

        logger.info(f"Extracting from {archive_fp}")

        with tarfile.open(archive_fp) as tar_f:
            member_fns = tar_f.getmembers()
            member_fns = [m for m in member_fns if m.name.endswith("text.xml")]

            for member_fn in tqdm(member_fns):
                member_content = tar_f.extractfile(member_fn).read()
                xml_str = member_content.decode()

                # extract plain text
                tree = ET.ElementTree(ET.fromstring(xml_str))

                root = tree.getroot()
                # Print out some information ("name" and "desc" tags)
                utterances = [elem.text for elem in root.iter() if elem.tag.endswith("{http://www.tei-c.org/ns/1.0}ab")]
                text = " ".join(utterances)

                yield text

src/lm_datasets/datasets/pl/pl_nkjp.py
malteos/lm-datasets
import logging
import tarfile
from tqdm.auto import tqdm

import xml.etree.cElementTree as ET

from lm_datasets.datasets.base import BaseDataset, Genre


logger = logging.getLogger(__name__)


class PLParliamentaryCorpusDataset(BaseDataset):
    """
    DOWNLOAD
    -----------

    Instruction

    - Run the command on the server:

    wget http://manage.legis.nlp.ipipan.waw.pl/download/ppc-nanno.tar.gz

    - Extract files:

    tar -xvf ppc-nanno.tar.gz

    - Run python script (locally with small sample)

    $ python pl_parliamentary_corpus.py \
        /Users/melinaplakidis/Documents/DFKI/eulm/lm_datasets/pl/ppc-sample \
        pl_parl_processed.jsonl

    """

    DATASET_ID = "pl_parliamentary_corpus"
    TITLE = "Polish Parliamentary Corpus / Korpus Dyskursu Parlamentarnego"
    DESCRIPTION = (
        "The Polish Parliamentary Corpus (PPC) is a Polish corpus made up of documents from the proceedings of the"
        " Polish Parliament, Sejm, and Senate. The corpus includes data of the Polish Sejm corpus and consists of"
        " stenographic records of plenary sittings and committee sittings, segments of interpellations and questions."
        " Texts in the PPC corpus cover the period of a hundred years from 1919 to 2019."
    )
    HOMEPAGE = "http://clip.ipipan.waw.pl/PPC"

    LANGUAGES = ["pl"]
    LICENSE = "CC-BY"

    GENRES = [Genre.GOVERNMENT]

    DOWNLOAD_URLS = ["http://manage.legis.nlp.ipipan.waw.pl/download/ppc-nanno.tar.gz"]

    TOKENS = 671_292_351

    def decompress(self):
        # gzip -d ppc-nanno.tar.gz
        pass

    def get_texts(self):
        # read from tar file
        archive_fp = self.get_dataset_file_paths(needed_suffix=".tar", single_file=True)

        logger.info(f"Extracting from {archive_fp}")

        with tarfile.open(archive_fp) as tar_f:
            member_fns = tar_f.getmembers()
            member_fns = [m for m in member_fns if m.name.endswith("text_structure.xml")]

            for member_fn in tqdm(member_fns):
                member_content = tar_f.extractfile(member_fn).read()
                xml_str = member_content.decode()

                # extract plain text
                tree = ET.ElementTree(ET.fromstring(xml_str))
                root = tree.getroot()
                utterances = [elem.text for elem in root.iter() if elem.tag.endswith("{http://www.tei-c.org/ns/1.0}u")]
                text = " ".join(utterances)

                # print(text)
                yield text

src/lm_datasets/datasets/pl/pl_parliamentary_corpus.py
malteos/lm-datasets
from lm_datasets.datasets.base import Availability, BILLION
from lm_datasets.datasets.hf_dataset import HFDataset


class BrWacDataset(HFDataset):
    """
    Download instructions:

    You need to
    1. Manually download `brwac.vert.gz` from https://www.inf.ufrgs.br/pln/wiki/index.php?title=BrWaC
    2. Extract the brwac.vert.gz in; this will result in the file brwac.vert in a folder <path/to/folder>
    The <path/to/folder> can e.g. be `~/Downloads`.
    BrWaC can then be loaded using the following command `datasets.load_dataset("brwac", data_dir="<path/to/folder>")`.

    """

    DATASET_ID = "brwac"

    TITLE = "Brazilian Portuguese Web as Corpus"
    HOMEPAGE = "https://huggingface.co/datasets/brwac"
    AVAILIBILITY = Availability.ON_REQUEST

    LANGUAGES = ["pt"]

    DESCRIPTION = (
        "The BrWaC (Brazilian Portuguese Web as Corpus) is a large corpus constructed following"
        "the Wacky framework, which was made public for research purposes."
        "The current corpus version, released in January 2017, is composed by 3.53 million documents,"
        "2.68 billion tokens and 5.79 million types. Please note that this resource is available"
        "solely for academic research purposes, and you agreed not to use it for any commercial applications."
    )

    TOKENS = 2.6 * BILLION

    HF_DATASET_ID = "brwac"
    HF_DATASET_SPLIT = "train"
    HF_DATA_DIR = "/netscratch/ortiz/corpora/ELE/pt/"

    text_column_name = "text"
    remove_columns = ["doc_id", "uri"]

    title_column_name = "title"

    def get_text_from_item(self, item) -> str:
        return item["title"] + self.title_delimiter + (self.paragraph_delimiter.join(item["text"]["paragraphs"]))

src/lm_datasets/datasets/pt/brwac.py
malteos/lm-datasets
from lm_datasets.datasets.base import Availability, GB
from lm_datasets.datasets.hf_dataset import HFDataset

import logging

logger = logging.getLogger(__name__)


class ParlamentoPtDataset(HFDataset):
    DATASET_ID = "parlamento_pt"
    TITLE = "ParlamentoPT"
    HOMEPAGE = "https://huggingface.co/datasets/PORTULAN/parlamento-pt"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["pt"]

    DESCRIPTION = """
    The ParlamentoPT is a Portuguese language data set obtained by collecting publicly
    available documents containing transcriptions of debates in the Portuguese Parliament. The data was collected from
    the Portuguese Parliament portal in accordance with its open data policy.
    """

    BYTES = 2.6 * GB

    HF_DATASET_ID = "PORTULAN/parlamento-pt"
    HF_DATASET_SPLIT = "train"

    HF_DATASET_CONFIGS = ["PORTULAN--parlamento-pt"]

    text_column_name = "text"
    remove_columns = []


#    title_column_name = "title"

src/lm_datasets/datasets/pt/parlamento_pt.py
malteos/lm-datasets
import logging
import zipfile
from tqdm.auto import tqdm
from lm_datasets.datasets.base import BaseDataset, MILLION, Genre

logger = logging.getLogger(__name__)


class MarcellLegislativeSubcorpusV2Dataset(BaseDataset):
    DATASET_ID = "marcell_legislative_subcorpus_v2"
    TITLE = "MARCELL Romanian legislative subcorpus v2"
    HOMEPAGE = "https://elrc-share.eu/repository/browse/marcell-romanian-legislative-subcorpus-v2/2da548428b9d11eb9c1a00155d026706ce94a6b59ffc4b0e9fb5cd9cebe6889e/"  # noqa
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["ro"]
    GENRES = [Genre.LEGAL]

    DESCRIPTION = (
        "The Romanian corpus contains 163,274 files, which represent the body of national legislation ranging from 1881"
        " to 2021. This corpus includes mainly: governmental decisions, ministerial orders, decisions, decrees and"
        " laws. All the texts were obtained via crawling from the public Romanian legislative portal. This corpus"
        " resulted from the MARCELL project. Alternate download location: https://relate.racai.ro/marcell/new/"
    )
    PII = "No"
    LICENSE = "public domain"

    TOKENS = 31 * MILLION

    def get_texts(self):
        archive_fp = self.get_dataset_file_paths(single_file=True, needed_suffix=".zip")
        logger.info(f"Reading from archive: {archive_fp}")

        with zipfile.ZipFile(archive_fp, "r") as zf:
            # zf.printdir()
            member_fns = [fn for fn in zf.namelist() if fn in {"ro-raw.zip"}]

            for member_fn in member_fns:
                logger.info(f"Reading from archive member: {member_fn}")

                with zf.open(member_fn) as member_f:  #
                    with zipfile.ZipFile(member_f) as zf_2:
                        needed_suffix_2 = ".txt"
                        members_fns_2 = zf_2.namelist()
                        members_fns_2 = [fn for fn in members_fns_2 if fn.endswith(needed_suffix_2)]

                        for member_fn_2 in tqdm(members_fns_2, desc="Extracting txt files"):
                            # logger.info(f"Reading archive member from {member_fn_2}")
                            with zf_2.open(member_fn_2) as member_f_2:
                                text = member_f_2.read().decode()
                                yield text

src/lm_datasets/datasets/ro/marcell_legislative_subcorpus_v2.py
malteos/lm-datasets
from typing import List
from lm_datasets.datasets.base import BaseDataset, Genre


class SKCourtDecisionsDataset(BaseDataset):
    DATASET_ID = "sk_court_decisions"
    TITLE = "od-justice 2.0"
    HOMEPAGE = "https://www.juls.savba.sk/data/od-justice/od-justice-2.0.ver.xz"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["sk"]
    GENRES = [Genre.LEGAL]

    DESCRIPTION = "Slovak court decisions"
    PII = "No"
    LICENSE = "public domain"

    TOKENS = 10_618_105_036

    def parse_doc_lines(self, doc_lines: List[str]):
        # Parse a full <doc> item
        # extract URL from doc_lines[0]:
        # <doc url="https://obcan.justice.sk/content/public/item/48712d97-58f9-4033-ab41-313d0e8e6631" court="Okresný súd Stará Ľubovňa" zn="4Ps/4/2013" date="2014-01-13" tokcountdd="1422" tokcount="1422" >\  # noqa
        # print("".join(doc_lines))

        doc_text = ""
        sent = ""
        for doc_line in doc_lines:
            if doc_line.startswith("</block") or doc_line.startswith("</p"):
                # Each block is a paragraph -> paragraph completed
                # yield doc_text.rstrip()  # remove last white space
                # doc_text = ""
                # return doc_text
                doc_text += self.paragraph_delimiter
            elif doc_line.startswith("<s "):
                # beginning of sentence
                sent = ""
            elif doc_line.startswith("</s>"):
                # parse sentence
                s = ""
                last_line = "<g/>"
                for sent_line in sent.splitlines():
                    sent_line = sent_line.strip()

                    if sent_line.startswith("<"):
                        pass
                    else:
                        if last_line != "<g/>":
                            s += " "

                        cols = sent_line.split("\t", 2)

                        if len(cols) > 1:  # is valid entry (CONLLU-like format)
                            original = cols[0]  # first column contains original string
                            s += original

                    last_line = sent_line

                # sentence parsed
                doc_text += s + self.sentence_delimiter
            else:
                sent += doc_line

        return doc_text

    def get_texts(self):
        import lzma as xz

        filepath = self.get_dataset_file_paths(needed_suffix=".xz", single_file=True)

        with xz.open(open(filepath, "rb"), "rt", encoding="utf-8") as f_in:
            doc_lines = []
            doc_counter = 0

            # Read file line by line
            for i, line in enumerate(f_in):
                # line = line.strip()

                doc_lines.append(line)

                # Extract <doc> ... </doc>
                if line.startswith("<doc "):
                    doc_lines = [line]
                elif line.startswith("</doc>"):
                    # End of doc -> process
                    doc_counter += 1

                    text = self.parse_doc_lines(doc_lines)

                    if text:
                        yield text

src/lm_datasets/datasets/sk/sk_court_decisions.py
malteos/lm-datasets
from typing import List
from lm_datasets.datasets.base import BaseDataset, Genre


class SKLawsDataset(BaseDataset):
    DATASET_ID = "sk_laws"
    TITLE = "Korpus slovenských právnych predpisov v1.9"
    HOMEPAGE = "https://www.juls.savba.sk/data/marcell/legal-sk-20220322-1.9.ver.xz"
    AVAILIBILITY = "Yes - it has a direct download link or links"

    LANGUAGES = ["sk"]
    GENRES = [Genre.LEGAL]

    DESCRIPTION = "Slovak body of laws (1955-2022)"
    PII = "No"
    LICENSE = "public domain"

    TOKENS = 45_000_000

    def parse_doc_lines(self, doc_lines: List[str]):
        # Parse a full <doc> item
        # extract URL from doc_lines[0]:
        # <doc url="https://obcan.justice.sk/content/public/item/48712d97-58f9-4033-ab41-313d0e8e6631" court="Okresný súd Stará Ľubovňa" zn="4Ps/4/2013" date="2014-01-13" tokcountdd="1422" tokcount="1422" >\  # noqa
        # print("".join(doc_lines))

        doc_text = ""
        sent = ""
        for doc_line in doc_lines:
            if doc_line.startswith("</block") or doc_line.startswith("</p"):
                # Each block is a paragraph -> paragraph completed
                # yield doc_text.rstrip()  # remove last white space
                # doc_text = ""
                # return doc_text
                doc_text += self.paragraph_delimiter
            elif doc_line.startswith("<s "):
                # beginning of sentence
                sent = ""
            elif doc_line.startswith("</s>"):
                # parse sentence
                s = ""
                last_line = "<g/>"
                for sent_line in sent.splitlines():
                    sent_line = sent_line.strip()

                    if sent_line.startswith("<"):
                        pass
                    else:
                        if last_line != "<g/>":
                            s += " "

                        cols = sent_line.split("\t", 2)

                        if len(cols) > 1:  # is valid entry (CONLLU-like format)
                            original = cols[0]  # first column contains original string
                            s += original

                    last_line = sent_line

                # sentence parsed
                doc_text += s + self.sentence_delimiter
            else:
                sent += doc_line

        return doc_text

    def get_texts(self):
        import lzma as xz

        filepath = self.get_dataset_file_paths(needed_suffix=".xz", single_file=True)

        with xz.open(open(filepath, "rb"), "rt", encoding="utf-8") as f_in:
            doc_lines = []
            doc_counter = 0

            # Read file line by line
            for i, line in enumerate(f_in):
                # line = line.strip()

                doc_lines.append(line)

                # Extract <doc> ... </doc>
                if line.startswith("<doc "):
                    doc_lines = [line]
                elif line.startswith("</doc>"):
                    # End of doc -> process
                    doc_counter += 1

                    text = self.parse_doc_lines(doc_lines)

                    if text:
                        yield text

src/lm_datasets/datasets/sk/sk_laws.py
malteos/lm-datasets
"""
https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1448/kas.json.tar.gz

academic_slovene_kas

"""

import json
import logging

from lm_datasets.datasets.base import BaseDataset, Availability

from tqdm.auto import tqdm
from smart_open import open


logger = logging.getLogger(__name__)


class AcademicSloveneKASDataset(BaseDataset):
    """
    The KAS corpus of Slovene academic writing consists of almost 65,000 BSc/BA, 16,000 MSc/MA and 1,600 PhD theses
    (82 thousand texts, 5 million pages or 1,5 billion tokens) written 2000 - 2018 and gathered from the digital
    libraries of Slovene higher education institutions via the Slovene Open Science portal (http://openscience.si/).
    """

    DATASET_ID = "academic_slovene_kas"
    TITLE = "Corpus of academic Slovene KAS 2.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1448"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["sl"]

    LOCAL_DIRS = ["pegasus:/netscratch/ortiz/corpora/ELE/sl/academic_slovene_kas"]

    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1448/kas.json.tar.gz"]

    TOKENS = 1_496_079_001

    def decompress(self):
        # tar gz
        pass

    def get_texts(self):
        # TODO read directly from tar gz

        # Read from many JSON files
        for fp in tqdm(self.get_dataset_file_paths(needed_suffix=".json", subdirectories=True), desc="Reading files"):
            # logger.info(f"Reading from {fp}")

            with open(fp) as f:
                doc = json.load(f)
                text = ""
                for line, paragraphs in doc.items():
                    for p in paragraphs:
                        text += p.replace("\n", " ").strip() + "\n\n"

                    # text += "\n"
                # print(text.strip())
                # print("#######")

                yield text.strip()

src/lm_datasets/datasets/sl/academic_slovene_kas.py
malteos/lm-datasets
from lm_datasets.datasets.base import BaseDataset, Availability


class CCGigaFidaDataset(BaseDataset):
    """ """

    DATASET_ID = "cc_gigafida"
    TITLE = "Written corpus ccGigafida 1.0"
    HOMEPAGE = "https://www.clarin.si/repository/xmlui/handle/11356/1035"
    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    LANGUAGES = ["sl"]

    DOWNLOAD_URLS = ["https://www.clarin.si/repository/xmlui/bitstream/handle/11356/1035/ccGigafida-text.zip"]

    TOKENS = 126919097

    def decompress(self):
        # decompress single zip file
        pass

    def get_texts(self):
        # TODO read directly from zip

        file_paths = self.get_dataset_file_paths(needed_suffix=".txt")
        for fp in file_paths:
            with open(fp) as f:
                text = f.read()

                yield text

src/lm_datasets/datasets/sl/cc_gigafida.py
malteos/lm-datasets
from typing import List
from lm_datasets.datasets.base import BILLION, BaseDataset, Availability
from lm_datasets.utils import generate_texts_from_tab_columns_in_xml
from smart_open import open


class SLWaCWebDataset(BaseDataset):
    DATASET_ID = "slwac_web"
    HOMEPAGE = "http://nlp.ffzg.hr/resources/corpora/slwac/"
    DOWNLOAD_URLS = ["http://nlp.ffzg.hr/data/corpora/slwac2.0.gz"]
    LANGUAGES = ["sl"]

    TITLE = "slWaC web corpus"
    DESCRIPTION = (
        "slWaC is a web corpus collected from the .si top-level domain in 2011 and 2014. The corpus is"
        " tokenized and annotated with the lemma and the morphosyntax layer."
    )

    TOKENS = 1.2 * BILLION
    LICENSE = "open license"
    AVAILIBITIY = Availability.DIRECT_DOWNLOAD
    WEB_CRAWLED = True

    def parse_doc_lines(self, doc_lines: List[str]):
        # Parse a full p item

        doc_text = ""
        sent = ""
        for doc_line in doc_lines:
            if doc_line.startswith("</block") or doc_line.startswith("</p"):
                # Each block is a paragraph -> paragraph completed
                # yield doc_text.rstrip()  # remove last white space
                # doc_text = ""
                # return doc_text
                doc_text += self.paragraph_delimiter
            elif doc_line.startswith("<s "):
                # beginning of sentence
                sent = ""
            elif doc_line.startswith("</s>"):
                # parse sentence
                s = ""
                last_line = "<g/>"
                for sent_line in sent.splitlines():
                    sent_line = sent_line.strip()

                    if sent_line.startswith("<"):
                        pass
                    else:
                        if last_line != "<g/>":
                            s += " "

                        cols = sent_line.split("\t", 2)

                        if len(cols) > 1:  # is valid entry (CONLLU-like format)
                            original = cols[0]  # first column contains original string
                            s += original

                    last_line = sent_line

                # sentence parsed
                doc_text += s + self.sentence_delimiter
            else:
                sent += doc_line

        return doc_text

    def get_texts(self):
        # read from slwac2.0.gz
        fp = self.get_dataset_file_paths(single_file=True, needed_suffix=".gz")

        with open(fp) as f:
            yield from generate_texts_from_tab_columns_in_xml(
                f,
                doc_start="<text ",
                doc_end="</text>",
                sentence_delimiter=self.sentence_delimiter,
                paragraph_delimiter=self.paragraph_delimiter,
            )

        raise NotImplementedError

src/lm_datasets/datasets/sl/slwac_web.py
malteos/lm-datasets
import json
from typing import Iterable
from lm_datasets.datasets.base import BaseDataset, Availability, QualityWarning
import zipfile
import logging

logger = logging.getLogger(__name__)


class SrpKorDataset(BaseDataset):
    DATASET_ID = "srpkor"
    TITLE = "SrpKorSubset (news, legal, academic, conversation, literary)"

    AVAILIBILITY = Availability.ON_REQUEST

    LANGUAGES = ["sr"]
    QUALITY_WARNINGS = [QualityWarning.SHORT_TEXT]

    def get_texts(self) -> Iterable[str]:
        # read from OpenGptx-JeRTeh.zip
        zip_fp = self.get_dataset_file_paths(needed_suffix=".zip", single_file=True)
        logger.info(f"Extracting from {zip_fp}")
        with zipfile.ZipFile(zip_fp) as zf:
            member_fns = zf.namelist()
            for fn in member_fns:
                if fn.endswith(".json"):
                    logger.info(f"Reading from {fn}")

                    with zf.open(fn) as member_f:
                        obj = json.load(member_f)

                        for sent in obj["sents"]:
                            # yield str(len(obj["sents"]))
                            yield sent

                # break

src/lm_datasets/datasets/sr/srpkor.py
malteos/lm-datasets
from io import BytesIO
import logging

import bz2

import tarfile

from lm_datasets.datasets.base import BaseDataset, Availability, BILLION


logger = logging.getLogger(__name__)


class SVGigawordDataset(BaseDataset):
    DATASET_ID = "sv_gigaword"
    TITLE = "The Swedish Culturomics Gigaword Corpus"
    HOMEPAGE = "https://spraakbanken.gu.se/en/resources/gigaword"
    DESCRIPTION = (
        "One billion Swedish words from 1950 and onwards. Code to extract data from the corpus, as well as usage"
        " instructions, can be downloaded from https://svn.spraakdata.gu.se/sb-arkiv/tools/gigaword/"
    )

    LANGUAGES = ["sv"]
    TOKENS = 1 * BILLION

    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    DOWNLOAD_URLS = [
        "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-1950-59.tar",
        "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-1960-69.tar",
        "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-1970-79.tar",
        "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-1980-89.tar",
        "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-1990-99.tar",
        "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-2000-09.tar",
        "https://spraakbanken.gu.se/lb/resurser/meningsmangder/gigaword-2010-15.tar",
    ]

    def is_downloaded(self):
        return len(self.get_dataset_file_paths(needed_suffix=".tar")) == len(self.DOWNLOAD_URLS)

    def get_texts(self):
        from lxml import etree

        if not self.is_downloaded():
            self.download()

        # read from tar files
        tar_file_paths = self.get_dataset_file_paths(needed_suffix=".tar")
        for tar_fp in tar_file_paths:
            logger.info(f"Extracting from {tar_fp}")

            with tarfile.open(tar_fp) as tar_f:
                members = [m for m in tar_f.getmembers() if m.name.endswith(".xml.bz2")]

                # Read from XML bz2 files
                for member in members:
                    logger.info(f"Decompress: {member.name}")

                    decompressed_member = bz2.decompress(tar_f.extractfile(member).read())

                    logger.info(f"Parse: {member.name}")

                    # Iterate over "text" tags
                    for event, element in etree.iterparse(BytesIO(decompressed_member), tag="text"):
                        text = ""
                        # TODO: whitespaces before punctutation
                        for el in element.itertext():
                            text += el.replace("\n", " ")
                        text = text.replace("  ", " ").strip()

                        # print(text)

                        yield text

src/lm_datasets/datasets/sv/sv_gigaword.py
malteos/lm-datasets
import logging
from smart_open import open

from lm_datasets.datasets.base import Availability, BaseDataset


logger = logging.getLogger(__name__)


class UKLawsDataset(BaseDataset):
    DATASET_ID = "uk_laws"
    TITLE = "Corpus of laws and legal acts of Ukraine"
    HOMEPAGE = "https://lang.org.ua/en/corpora/#anchor7"

    LANGUAGES = ["uk"]
    TOKENS = 578_988_264

    AVAILIBILITY = Availability.DIRECT_DOWNLOAD

    DOWNLOAD_URLS = ["https://lang.org.ua/static/downloads/corpora/laws.txt.tokenized.bz2"]

    def is_downloaded(self):
        return len(self.get_dataset_file_paths(needed_suffix=".tokenized.bz2")) == len(self.DOWNLOAD_URLS)

    def get_texts(self):
        if not self.is_downloaded():
            self.download()

        file_path = self.get_dataset_file_paths(needed_suffix=".tokenized.bz2", single_file=True)

        logger.info(f"Reading from {file_path}")

        with open(file_path) as f:
            text = ""
            for i, line in enumerate(f):
                # print(line)
                # print("'", line, "'")

                if line.strip().isnumeric():
                    # new doc
                    # print(text)
                    yield text.strip()

                    text = ""
                else:
                    text += line

src/lm_datasets/datasets/uk/uk_laws.py
malteos/lm-datasets
import argparse

import logging

from lm_datasets.utils import get_auto_workers, get_bytes_from_int_or_string

from .datasets.dataset_registry import get_registered_dataset_classes
from .datasets.base import BaseDataset


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("datasets", help="Name of datasets to process (comma separated)")
    parser.add_argument(
        "output_dir",
        help="Output is saved in this directory (<language code>/<source name>.<jsonl/parquet>)",
    )
    parser.add_argument("--override", action="store_true", help="Override existing output files")
    parser.add_argument(
        "--ignore_errors",
        action="store_true",
        help="Ignore dataset-level errors (use when processing multiple datasets)",
    )
    parser.add_argument("--json_ensure_ascii", action="store_true", help="Escape non-ASCII characters in JSON output")
    parser.add_argument("--output_format", default="jsonl", type=str, help="Output format (jsonl,parquet)")
    parser.add_argument(
        "--output_compression",
        default=None,
        type=str,
        help="""Compression of output (jsonl: "gzip"; parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
    )
    parser.add_argument(
        "--output_batch_size",
        default=1000,
        type=int,
        help="""Write batch size; smaller batch size = more accurate splitts but slower""",
    )
    parser.add_argument(
        "--max_output_chunk_uncompressed_bytes",
        default="10GB",
        type=str,
        help="Chunks are splitted if they exceed this byte count (<n>, <n>KB, <n>MB, or <n>GB)",
    )
    parser.add_argument(
        "--workers",
        default=1,
        type=int,
        help="Number of workers for parallel processing",
    )
    parser.add_argument(
        "--raw_datasets_dir",
        default=None,
        type=str,
        help="Dataset files are read from this directory",
    )
    parser.add_argument(
        "--extra_dataset_registries",
        default=None,
        type=str,
        help="List of Python packages to load dataset registries",
    )
    parser.add_argument(
        "--log_file",
        default=None,
        type=str,
        help="Log file is saved at this path",
    )
    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging (log level = debug)")
    parser.add_argument("--limit", default=0, type=int, help="Limit dataset size (for debugging)")
    parser.add_argument(
        "--skip_items",
        default=0,
        type=int,
        help="Skip N items (depending on dataset: directories, subsets, files, documents) (for debugging)",
    )
    parser.add_argument("--hf_auth_token", default=None, type=str, help="HuggingFace auth token")
    args = parser.parse_args()

    log_handlers = [logging.StreamHandler()]

    if args.log_file:
        log_handlers.append(logging.FileHandler(args.log_file))

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.DEBUG if args.verbose else logging.INFO,
        handlers=log_handlers,
    )
    logger = logging.getLogger(__name__)

    id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes(args.extra_dataset_registries)}

    datasets_list = args.datasets.split(",")

    if len(datasets_list) == 1 and datasets_list[0] == "all":
        # Get list of all non-dummy datasets
        datasets_list = id_to_dataset_class.keys()

    # Iterate over datasets
    for i, dataset_id in enumerate(datasets_list, 1):
        if dataset_id in id_to_dataset_class:
            dataset_cls = id_to_dataset_class[dataset_id]
        else:
            raise ValueError(f"Unknown dataset ID: {dataset_id} (available: {id_to_dataset_class.keys()})")

        logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")

        try:
            dataset: BaseDataset = dataset_cls(
                raw_datasets_dir=args.raw_datasets_dir,
                output_dir=args.output_dir,
                workers=get_auto_workers(args.workers),
                limit=args.limit,
                override_output=args.override,
                output_format=args.output_format,
                output_compression=args.output_compression,
                output_batch_size=args.output_batch_size,
                json_ensure_ascii=args.json_ensure_ascii,
                skip_items=args.skip_items,
                max_output_chunk_uncompressed_bytes=get_bytes_from_int_or_string(
                    args.max_output_chunk_uncompressed_bytes
                ),
            )

            if dataset.is_dummy():
                logger.warning(f"Skipping {dataset_id} (cannot extract from dummy datasets)")
                continue

            try:
                dataset.extract_plaintext()
            except FileExistsError as e:
                logger.error(f"Cannot extract text from {dataset_id}: {e}")
        except KeyboardInterrupt as e:
            logger.error(f"Stopping... {e}")
            break
        except BaseException as e:
            if args.ignore_errors:
                logger.error(f"Unexpected error occured with {dataset_id}: {e}")
            else:
                raise e

    logger.info("Done")

src/lm_datasets/extract_plaintext.py
malteos/lm-datasets
import typing as T
import logging

logger = logging.getLogger(__name__)


def get_texts_from_conllu_file(
    file_handler: T.TextIO,
    title_delimiter: str = ":\n\n",
    sentence_delimiter: str = " ",
    skip_sentence_prefixes: T.Optional[T.Iterable[str]] = None,
):
    """
    Reads CONLLU and CONLLU-Plus format: https://universaldependencies.org/ext-format.html

    https://github.com/EmilStenstrom/conllu/

    """
    import conllu
    from conllu.parser import (
        _FieldParserType,
        _MetadataParserType,
        parse_conllu_plus_fields,
        parse_sentences,
        parse_token_and_metadata,
    )
    from conllu.exceptions import ParseException

    # Custom parse function that handles errors
    def conllu_parse_incr(
        in_file: T.TextIO,
        fields: T.Optional[T.Sequence[str]] = None,
        field_parsers: T.Optional[T.Dict[str, _FieldParserType]] = None,
        metadata_parsers: T.Optional[T.Dict[str, _MetadataParserType]] = None,
        # skip_sentence_prefixes: T.Optional[T.Iterable[str]] = None,
    ) -> conllu.SentenceGenerator:
        if not hasattr(in_file, "read"):
            raise FileNotFoundError("Invalid file, 'parse_incr' needs an opened file as input")

        if not fields:
            fields = parse_conllu_plus_fields(in_file, metadata_parsers=metadata_parsers)

        def generator():
            for sentence in parse_sentences(in_file):
                if skip_sentence_prefixes:
                    skip = False
                    for prefix in skip_sentence_prefixes:
                        if sentence.startswith(prefix):
                            logger.debug(f"Skip sentence: {sentence}")
                            skip = True
                            break

                    if skip:
                        continue

                try:
                    yield parse_token_and_metadata(
                        sentence, fields=fields, field_parsers=field_parsers, metadata_parsers=metadata_parsers
                    )
                except ParseException as e:
                    logger.error(f"Cannot parse sentence: {sentence}; {e}")

        return conllu.SentenceGenerator(generator())

    # from conllu.exceptions import ParseException

    text = None

    # try:
    for sentence in conllu_parse_incr(file_handler):
        if "newdoc id" in sentence.metadata:
            if text is not None:
                # doc completed
                yield text
            text = ""  # init empty document

        # append text to doc
        if "text" in sentence.metadata:
            if not text:
                text = ""  # some conllu are not using doc ids -> force init
            else:
                text += sentence_delimiter  # whitespace betweeen sentences

            text += sentence.metadata["text"]

        if "title" in sentence.metadata:
            text += title_delimiter

    # yield last document
    if text:
        yield text

src/lm_datasets/io/conllu_file.py
malteos/lm-datasets
import os
import pyarrow.parquet as pq
import logging

import itertools
from typing import Any, Generator, Tuple

import pyarrow as pa

logger = logging.getLogger(__name__)


def open_parquet_file_with_retries(file_path, retries: int = 2):
    """
    A little hack to avoid the "[Errno 14] Error reading bytes from file. Detail: [errno 14] Bad address"
    """
    for retry in range(1, retries):
        try:
            f = pq.ParquetFile(file_path)
            return f
        except OSError as e:
            logger.error(f"Could not open parquet file due to `Bad address` error (retry {retry}/{retries}): {e}")
            pass

    # last try that does not catch the error
    f = pq.ParquetFile(file_path)
    return f


def get_parquet_batches(texts: Generator[str, Any, None], schema, batch_size):
    while True:
        arr = pa.array(itertools.islice(texts, batch_size))
        batch = pa.RecordBatch.from_arrays([arr], schema=schema)

        if not batch:
            break
        yield batch


def save_texts_to_parquet_chunks(
    texts: Generator[str, Any, None],
    schema,
    max_chunk_uncompressed_bytes: int,
    output_path_func: callable,
    compression: str = "ZSTD",
    batch_size: int = 1024,
    print_write_progress: int = 10_000,
    limit: int = 0,
) -> Tuple[int, int]:
    max_chunks = 9999
    total_rows = 0
    total_bytes = 0
    batch_iter = get_parquet_batches(texts, schema=schema, batch_size=batch_size)
    limit_reached = False

    for chunk_i in range(1, max_chunks + 1):
        chunk_rows = 0
        chunk_nbytes = 0
        chunk_buffer_size = 0
        chunk_fp = output_path_func(chunk_i) if max_chunk_uncompressed_bytes > 0 else output_path_func()

        logger.info(f"Writing to {chunk_fp}")

        with pq.ParquetWriter(chunk_fp, schema=schema, compression=compression) as writer:
            try:
                while True:
                    batch = next(batch_iter)
                    writer.write_batch(batch)
                    total_rows += len(batch)
                    total_bytes = batch.nbytes
                    chunk_rows += len(batch)
                    chunk_nbytes += batch.nbytes
                    chunk_buffer_size += batch.get_total_buffer_size()

                    if total_rows > 0 and (total_rows % print_write_progress) == 0:
                        logger.info(f"Written {total_rows:,} rows ...")

                    if limit > 0 and total_rows >= limit:
                        logger.warning(f"Limit reached ({total_rows:,} rows)")
                        limit_reached = True
                        break

                    if max_chunk_uncompressed_bytes > 0 and chunk_nbytes >= max_chunk_uncompressed_bytes:
                        # if chunk_buffer_size >= max_chunk_bytes_with_safety:
                        logger.info(
                            f"Chunk {chunk_i} completed (rows: {chunk_rows:,}; nbytes: {chunk_nbytes:,}; buffer size:"
                            f" {chunk_buffer_size:,})"
                        )
                        # logger.info(f"Chunk size on disk: {os.stat(chunk_fp).st_size:,} bytes")
                        break

            except StopIteration:
                logger.info(f"All rows written ({total_rows=}; {total_bytes=})")
                break

        if limit_reached:
            # break outer loop
            break

    total_chunks = chunk_i

    if max_chunk_uncompressed_bytes > 0:
        # Rename files with total number of chunks
        for chunk_i in range(1, total_chunks + 1):
            logger.info(f"Renaming to {output_path_func(chunk_i, total_chunks)}")
            os.rename(output_path_func(chunk_i), output_path_func(chunk_i, total_chunks))

    return total_rows, total_chunks

src/lm_datasets/io/parquet.py
malteos/lm-datasets
import prevert
import warnings

from typing import TextIO


class PrevertFile(prevert.dataset):
    """
    Work-around for `prevert` dataset be directly read from file handlers
    """

    def __init__(self, file: TextIO, xml=True):
        self.file = file
        if xml:
            line = self.file.readline()
            if not line.startswith("<corpus"):
                warnings.warn(
                    "Warning: your prevert might not have a XML header. Define the xml second parameter as False in the"
                    " constructor if this is the case."
                )

src/lm_datasets/io/prevert_file.py
malteos/lm-datasets
import argparse

import logging

import pandas as pd

from .datasets.dataset_registry import get_registered_dataset_classes
from .datasets.base import BaseDataset

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger(__name__)


def get_datasets_as_dataframe(
        output_dir=None,
        output_format="jsonl",
        raw_datasets_dir=None,
        extra_dataset_registries=None
    ):

    # convert datasets to table rows
    rows = []

    for ds_cls in get_registered_dataset_classes(extra_dataset_registries=extra_dataset_registries):
        ds: BaseDataset = ds_cls(
            output_dir=output_dir,
            output_format=output_format,
            raw_datasets_dir=raw_datasets_dir,
        )
        row = dict(
            dataset_id=ds.DATASET_ID,
            source_id=ds.get_source_id(),
            title=ds.TITLE,
            homepage=ds.HOMEPAGE,
            language=ds.get_language_code(),
            tokens=ds.get_tokens(),
            bytes=ds.BYTES,
            is_downloaded=None,
            web_crawled=1 if ds.WEB_CRAWLED else 0,
            availibility=ds.AVAILIBILITY,
            dummy=1 if ds.DUMMY else 0,
            has_output_file=None,
        )

        if raw_datasets_dir:
            row["is_downloaded"] = 1 if ds.is_downloaded() else 0

        if output_dir:
            row["has_output_file"] = 1 if ds.has_output_files(min_file_size=1024) else 0

        rows.append(row)

    df = pd.DataFrame(rows)

    return df


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # parser.add_argument("dataset", help="Name of dataset to process")
    parser.add_argument("--print_format", default="csv", type=str, help="Print format (tsv,csv,md)")
    parser.add_argument(
        "--raw_datasets_dir",
        default=None,
        type=str,
        help="Dataset files are read from this directory (needed for `is_downloaded` field)",
    )
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        help="Processed dataset are saved in this directory (need for `has_output_file` field)",
    )
    parser.add_argument(
        "--output_format",
        default="jsonl",
        type=str,
        help="Format of processed dataset (need for `has_output_file` field)",
    )
    parser.add_argument(
        "--extra_dataset_registries",
        default=None,
        type=str,
        help="List of Python packages to load dataset registries",
    )
    args = parser.parse_args()
    print_format = args.print_format

    df = get_datasets_as_dataframe(
        output_dir=args.output_dir if args.output_dir else "/dev/null",
        output_format=args.output_format,
        raw_datasets_dir=args.raw_datasets_dir,
        extra_dataset_registries=args.extra_dataset_registries,
    )

    # to_markdown
    # to_csv
    to_kwargs = dict(index=False)
    if print_format == "tsv":
        out = df.to_csv(sep="\t", **to_kwargs)
    elif print_format == "csv":
        out = df.to_csv(**to_kwargs)
    elif print_format == "md":
        out = df.to_markdown(**to_kwargs)
    else:
        raise ValueError("Unsupported output format: %s" % print_format)

    print(out)

src/lm_datasets/print_stats.py
malteos/lm-datasets
"""

input: parquet files, configs

shuffle + train/test split

"""
import argparse

import os
import logging


from pathlib import Path

from .datasets.dataset_registry import get_registered_dataset_classes
from .datasets.base import BaseDataset, GB

from datasets import load_dataset

from tqdm.auto import tqdm

import pyarrow.parquet as pq

import polars as pl


DEFAULT_MIN_FILE_SIZE_FOR_BUFFERED_SHUFFLING = 5 * GB


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("datasets", help="Name of datasets to shuffle (comma separated)")
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        help="Unshuffled datasets are loaded from this directory",
    )
    parser.add_argument(
        "--shuffled_output_dir",
        default=None,
        type=str,
        help="Shuffled dataset are saved in this directory",
    )
    parser.add_argument(
        "--output_format",
        default="parquet",
        type=str,
        help="Format of processed dataset",
    )
    parser.add_argument(
        "--output_compression",
        default=None,
        type=str,
        help="""Compression of output (jsonl: "gzip"; parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
    )
    parser.add_argument(
        "--seed",
        default=0,
        type=int,
        help="Random seed",
    )
    parser.add_argument(
        "--shuffle_buffer_size",
        default=1_000_000,
        type=int,
        help="Number of items in buffer to be shuffled at once (larger buffer = more memory but better shuffing)",
    )
    parser.add_argument(
        "--min_file_size_for_buffered_shuffling",
        default=DEFAULT_MIN_FILE_SIZE_FOR_BUFFERED_SHUFFLING,
        type=int,
        help="Min. file size bytes for buffered shuffling (default: 5GB; set to 0 to disable)",
    )
    parser.add_argument(
        "--log_file",
        default=None,
        type=str,
        help="Log file is saved at this path",
    )
    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging (log level = debug)")
    parser.add_argument("--override", action="store_true", help="Override existing output files")
    parser.add_argument(
        "--skip_large_datasets",
        action="store_true",
        help="Skip datasets with bytes > --min_file_size_for_buffered_shuffling",
    )
    args = parser.parse_args()

    log_handlers = [logging.StreamHandler()]

    if args.log_file:
        log_handlers.append(logging.FileHandler(args.log_file))

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.DEBUG if args.verbose else logging.INFO,
        handlers=log_handlers,
    )
    logger = logging.getLogger(__name__)

    if args.output_format != "parquet":
        raise ValueError(f"Output format is not supported (currently only parquet is supported; {args.output_format=})")

    seed = args.seed
    logger.info(f"Seed: {seed}")

    shuffle_buffer_size = args.shuffle_buffer_size

    id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes()}

    datasets_list = args.datasets.split(",")

    if len(datasets_list) == 1 and datasets_list[0] == "all":
        # Get list of all non-dummy datasets
        datasets_list = id_to_dataset_class.keys()

    min_file_size_for_buffered_shuffling = args.min_file_size_for_buffered_shuffling

    # Iterate over datasets
    for i, dataset_id in enumerate(datasets_list, 1):
        if dataset_id in id_to_dataset_class:
            dataset_cls = id_to_dataset_class[dataset_id]
        else:
            raise ValueError(f"Unknown dataset ID: {dataset_id} (available: {id_to_dataset_class.keys()})")

        logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")

        dataset: BaseDataset = dataset_cls(
            output_dir=args.output_dir,
            output_format=args.output_format,
            output_batch_size=shuffle_buffer_size,
            output_compression=args.output_compression,
        )

        if not dataset.has_output_files():
            logger.warning(f"Skipping {dataset_id}: cannot shuffle dataset without processed output files")
            continue

        output_fps = dataset.get_output_file_paths()

        for i, output_fp in enumerate(sorted(output_fps), i):
            logger.info(f"Shuffling {output_fp} ({i} / {len(output_fps)} files of {dataset_id})")

            output_file_name = Path(output_fp).name

            shuffled_output_file_path = os.path.join(
                args.shuffled_output_dir, output_file_name.replace(".parquet", ".shuffled.parquet")
            )

            assert shuffled_output_file_path != output_fp

            if os.path.exists(shuffled_output_file_path):
                if args.override:
                    logger.warning("Overriding existing shuffled output file")
                else:
                    logger.warning(
                        f"Skipping {dataset_id}: Shuffled output file exist already ({shuffled_output_file_path})"
                    )
                    continue

            # File size
            file_stats = os.stat(output_fp)

            if min_file_size_for_buffered_shuffling > 0 and file_stats.st_size > min_file_size_for_buffered_shuffling:
                # File is too large to be shuffled all at once => shuffle in chunks
                if args.skip_large_datasets:
                    logger.info(f"Skip because too large dataset ({file_stats.st_size=})")
                    continue

                # Reading meta from parquet (for progress bar)
                logger.info("Reading metadata ...")
                metadata = pq.read_metadata(output_fp)
                docs_count = metadata.num_rows

                logger.info("Initializing HF streaming dataset ...")
                hf_dataset = load_dataset("parquet", data_files={"train": output_fp}, split="train", streaming=True)

                logger.info("Shuffling and writing to new file ...")

                def generate_text():
                    for item in tqdm(hf_dataset.shuffle(buffer_size=shuffle_buffer_size, seed=seed), total=docs_count):
                        yield str(item[dataset.output_text_field])  # force to str

                # Writer
                shuffled_docs_count = dataset.save_texts_to_parquet(
                    generate_text(), file_path=shuffled_output_file_path, apply_filter=False
                )

                if docs_count != shuffled_docs_count:
                    logger.error(
                        f"Original and shuffled docs count do not match: {docs_count=} vs {shuffled_docs_count=}"
                    )

            else:
                # Shuffle in memory
                logger.info("Initializing PL in-memory dataframe ...")
                df = pl.read_parquet(output_fp)
                logger.info("Shuffling and writing to new file ...")
                df = df.sample(fraction=1, shuffle=True, seed=seed).write_parquet(
                    shuffled_output_file_path, compression="zstd"
                )

    logger.info("Done")

src/lm_datasets/shuffle.py
malteos/lm-datasets
import argparse

import os
import logging
from pathlib import Path


from .datasets.dataset_registry import get_registered_dataset_classes
from .datasets.base import BaseDataset

from .utils import get_bytes_from_int_or_string, get_parquet_compression

import pyarrow.parquet as pq


def validate_original_and_parts():
    # TODO make sure original and parts are the same
    pass


def iter_parquet(pq_file, batch_size):
    # tbl = pq.ParquetFile(file_path)
    for batch in pq_file.iter_batches(batch_size):
        yield batch


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("datasets", help="Name of datasets to shuffle (comma separated)")
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        help="Unshuffled datasets are loaded from this directory",
    )
    parser.add_argument(
        "--output_format",
        default="parquet",
        type=str,
        help="Format of processed dataset",
    )
    parser.add_argument(
        "--max_uncompressed_bytes_per_chunk",
        default="5GB",
        type=str,
        help="Chunks are splitted if they exceed this byte count (<n>, <n>KB, <n>MB, or <n>GB)",
    )
    parser.add_argument(
        "--safety_factor",
        default=0.975,
        type=float,
        help="Max. chunk size is multiplied with this factor (accounts for inaccurate chunk sizes due to batching)",
    )
    parser.add_argument(
        "--output_compression",
        default=None,
        type=str,
        help="""Compression of output (default: compression of existing output; jsonl: "gzip"; """
        + """parquet: "NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD")""",
    )
    parser.add_argument(
        "--batch_size",
        default=16,
        type=int,
        help="""Read and write batch size; smaller batch size = more accurate splitts but slower""",
    )
    parser.add_argument(
        "--log_file",
        default=None,
        type=str,
        help="Log file is saved at this path",
    )
    parser.add_argument("--verbose", action="store_true", help="Enable verbose logging (log level = debug)")
    parser.add_argument("--override", action="store_true", help="Override existing output files")
    parser.add_argument("--rename_original", action="store_true", help="Renames orignal output file")
    args = parser.parse_args()

    log_handlers = [logging.StreamHandler()]

    if args.log_file:
        log_handlers.append(logging.FileHandler(args.log_file))

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
        level=logging.DEBUG if args.verbose else logging.INFO,
        handlers=log_handlers,
    )
    logger = logging.getLogger(__name__)

    if args.output_format != "parquet":
        raise ValueError("Output format is not supported (currently only parquet is supported)")

    id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes()}

    datasets_list = args.datasets.split(",")

    if len(datasets_list) == 1 and datasets_list[0] == "all":
        # Get list of all non-dummy datasets
        datasets_list = id_to_dataset_class.keys()

    max_uncompressed_bytes_per_chunk = get_bytes_from_int_or_string(args.max_uncompressed_bytes_per_chunk)
    max_chunk_bytes_with_safety = int(max_uncompressed_bytes_per_chunk * args.safety_factor)
    logger.info(f"Max. chunk size with {args.safety_factor=}: {max_chunk_bytes_with_safety:,} bytes")

    # Iterate over datasets
    for i, dataset_id in enumerate(datasets_list, 1):
        if dataset_id in id_to_dataset_class:
            dataset_cls = id_to_dataset_class[dataset_id]
        else:
            raise ValueError(f"Unknown dataset ID: {dataset_id} (available: {id_to_dataset_class.keys()})")

        logger.info(f"Dataset ID: {dataset_id} ({i} / {len(datasets_list)})")

        dataset: BaseDataset = dataset_cls(
            output_dir=args.output_dir,
            output_format=args.output_format,
            # output_compression=args.output_compression,
        )

        if dataset.has_chunked_output_files():
            logger.warning(f"Skipping {dataset_id}: dataset has already chunked output")

        if not dataset.has_single_output_file():
            logger.warning(f"Skipping {dataset_id}: cannot split dataset without processed output")
            continue

        logger.info(f"Splitting {dataset.get_single_output_file_path()}")

        # File size
        file_stats = os.stat(dataset.get_single_output_file_path())

        if file_stats.st_size > max_uncompressed_bytes_per_chunk:
            # File is too large => split into chunks
            output_file_path = Path(dataset.get_single_output_file_path())
            output_dir = output_file_path.parent

            # Check for existing chunks
            existing_chunk_fp_pattern = str(output_file_path.name).replace(".parquet", ".part-*.parquet")
            existing_chunk_fps = list(output_dir.glob(existing_chunk_fp_pattern))

            if existing_chunk_fps:
                if args.override:
                    logger.info("Removing existing part files ...")
                    for fp in existing_chunk_fps:
                        logger.info(f"Removing {fp} ...")
                        os.unlink(fp)
                else:
                    logger.warning("Skipping because part files already exist.")

            # Open parquet file
            # pq_file = pq.ParquetFile(dataset.get_output_file_path())
            with open(dataset.get_single_output_file_path(), "rb") as file_handler:
                # pq_file = open_parquet_file_with_retries(dataset.get_output_file_path(), retries=2)
                pq_file = pq.ParquetFile(file_handler)
                batch_iter = iter_parquet(pq_file, args.batch_size)  # iterate over row groups and batches

                # Auto-detect compression from source file
                if args.output_compression is None:
                    compression = pq_file.metadata.row_group(0).column(0).compression
                else:
                    compression = get_parquet_compression(args.output_compression)

                logger.info(f"Unsplitted file size: {file_stats.st_size:,} bytes ({compression=})")

                max_chunks = 9999  # max number of chunks
                total_docs_count = 0
                chunk_fps = []

                for chunk_i in range(1, max_chunks + 1):
                    # if total_docs_count >= pq_file.metadata.num_rows:
                    #     logger.warning("All docs written")
                    #     break
                    chunk_docs_count = 0
                    chunk_nbytes = 0
                    chunk_buffer_size = 0
                    chunk_fp = dataset.get_single_output_file_path().replace(".parquet", f".part-{chunk_i:04d}.parquet")
                    chunk_fps.append(chunk_fp)

                    logger.info(f"Writing to {chunk_fp}")

                    with pq.ParquetWriter(
                        chunk_fp, schema=pq_file.schema.to_arrow_schema(), compression=compression
                    ) as writer:
                        try:
                            while True:
                                batch = next(batch_iter)
                                writer.write_batch(batch)
                                total_docs_count += len(batch)
                                chunk_docs_count += len(batch)
                                chunk_nbytes += batch.nbytes
                                chunk_buffer_size += batch.get_total_buffer_size()

                                if chunk_nbytes >= max_chunk_bytes_with_safety:
                                    # if chunk_buffer_size >= max_chunk_bytes_with_safety:
                                    logger.info(
                                        f"Chunk {chunk_i} completed (docs: {chunk_docs_count:,}; nbytes:"
                                        f" {chunk_nbytes:,}; buffer size: {chunk_buffer_size:,})"
                                    )
                                    # logger.info(f"Chunk size on disk: {os.stat(chunk_fp).st_size:,} bytes")
                                    break

                        except StopIteration:
                            logger.info(f"All rows written ({total_docs_count=})")
                            break

                if total_docs_count != pq_file.metadata.num_rows:
                    logger.error(
                        f"Written rows do not match input rows: {total_docs_count:,} (expected:"
                        f" {pq_file.metadata.num_rows:,})"
                    )

                # Rename part files with total number of chunks
                total_chunks = len(chunk_fps)

                for chunk_fp in chunk_fps:
                    new_chunk_fp = chunk_fp.replace(".parquet", f"-of-{total_chunks:04d}.parquet")

                    logger.info(f"Renaming {new_chunk_fp}")
                    os.rename(chunk_fp, new_chunk_fp)

                # Keep original: by renaming to ".orignal"
                if args.rename_original:
                    new_path = str(output_file_path).replace(".parquet", ".parquet.original")
                    logger.info(f"Moving original file to {new_path}")
                    os.rename(str(output_file_path), new_path)

        else:
            # Small fie
            logger.info("File is small no split neded.")

    logger.info("Done")

src/lm_datasets/split_by_size.py
malteos/lm-datasets
import os
import logging
from enum import Enum
from typing import List


logger = logging.getLogger(__name__)


class System(Enum):
    DEFAULT = "default"
    PEGASUS = "pegasus"
    TEST_RUNNER = "test_runner"


def get_current_system(allow_default: bool = True) -> System:
    if os.path.exists("/netscratch/"):
        return System.PEGASUS
    elif os.path.exists("/home/runner/"):
        return System.TEST_RUNNER
    elif allow_default:
        return System.DEFAULT
    else:
        raise ValueError("Cannot determine system")


def get_path_by_system(possible_paths: List[str], allow_default: bool = True, default_path: str = "/dev/null") -> str:
    current_system = get_current_system(allow_default)

    for raw_path in possible_paths:
        system_name, path = raw_path.split(":", 2)

        if system_name == current_system.value:
            return path

    if allow_default:
        logger.warning("Using default path because current system cannot be matched!")

        return default_path
    else:
        raise ValueError(f"Cannot assign possible paths to current system: {possible_paths} {current_system=}")

src/lm_datasets/systems.py
malteos/lm-datasets
import re
import logging
from typing import Iterable, List, Union


logger = logging.getLogger(__name__)


def remove_whitespaces_before_punctuation(input_text: str) -> str:
    # Taken from [1]
    # [1]: https://stackoverflow.com/questions/18878936/how-to-strip-whitespace-from-before-but-not-after-punctuation-in-python  # noqa
    return re.sub(r"\s([?.!,:](?:\s|$))", r"\1", input_text).replace("( ", "(").replace(" )", ")")


def get_text_from_tab_columns_in_xml(doc_lines: List[str], sentence_delimiter, paragraph_delimiter) -> str:
    # Parse a full <doc> item
    # extract URL from doc_lines[0]:
    # <doc url="https://obcan.justice.sk/content/public/item/48712d97-58f9-4033-ab41-313d0e8e6631" court="Okresný súd Stará Ľubovňa" zn="4Ps/4/2013" date="2014-01-13" tokcountdd="1422" tokcount="1422" >\  # noqa
    # print("".join(doc_lines))

    doc_text = ""
    sent = ""
    for doc_line in doc_lines:
        if doc_line.startswith("</block") or doc_line.startswith("</p"):
            # Each block is a paragraph -> paragraph completed
            # yield doc_text.rstrip()  # remove last white space
            # doc_text = ""
            # return doc_text
            doc_text += paragraph_delimiter
        elif doc_line.startswith("<s ") or doc_line.startswith("<s>"):
            # beginning of sentence
            sent = ""
        elif doc_line.startswith("</s>"):
            # parse sentence
            s = ""
            last_line = "<g/>"
            for sent_line in sent.splitlines():
                sent_line = sent_line.strip()

                if sent_line.startswith("<"):
                    pass
                else:
                    if last_line != "<g/>":
                        s += " "

                    cols = sent_line.split("\t", 2)

                    if len(cols) > 1:  # is valid entry (CONLLU-like format)
                        original = cols[0]  # first column contains original string
                        s += original

                last_line = sent_line

            # sentence parsed
            doc_text += s + sentence_delimiter
        else:
            sent += doc_line

    # print(doc_text)
    # print(doc_lines[0])
    return doc_text


def generate_texts_from_tab_columns_in_xml(
    line_iterator: Iterable[str],
    doc_start: str = "<doc ",
    doc_end: str = "</doc>",
    sentence_delimiter: str = " ",
    paragraph_delimiter: str = "\n\n",
):
    doc_lines = []
    doc_counter = 0

    # Read file line by line
    for i, line in enumerate(line_iterator):
        # line = line.strip()

        doc_lines.append(line)

        # Extract <doc> ... </doc>
        if line.startswith(doc_start):
            doc_lines = [line]

        elif line.startswith(doc_end):
            # End of doc -> process
            doc_counter += 1

            text = get_text_from_tab_columns_in_xml(doc_lines, sentence_delimiter, paragraph_delimiter)

            if text:
                yield text


def get_auto_workers(workers: int) -> int:
    import multiprocessing

    if workers < 1:
        cpu_cores = multiprocessing.cpu_count()

        if workers == 0:
            workers = cpu_cores
            logger.info(f"Using all available CPU cores: {workers}")
        else:
            workers = cpu_cores - abs(workers)
            logger.info(f"Using workers: {workers} / (available cores {cpu_cores})")

    return workers


def get_parquet_compression(request_compression: str) -> str:
    supported_parquet_compressions = {"NONE", "SNAPPY", "GZIP", "BROTLI", "LZ4", "ZSTD"}

    if request_compression is None:
        return "NONE"
    elif request_compression in supported_parquet_compressions:
        return request_compression
    else:
        if request_compression.upper() in supported_parquet_compressions:
            return request_compression.upper()
        else:
            raise ValueError(
                f"Unsupported parquet compression: {request_compression} ({supported_parquet_compressions=})"
            )


def get_bytes_from_int_or_string(bytes_int_or_str: Union[int, str]) -> int:
    from lm_datasets.datasets.base import GB, KB, MB

    if isinstance(bytes_int_or_str, int):
        return bytes_int_or_str
    elif isinstance(bytes_int_or_str, str):
        if bytes_int_or_str.endswith("KB"):
            return int(bytes_int_or_str.rstrip("KB")) * KB
        elif bytes_int_or_str.endswith("MB"):
            return int(bytes_int_or_str.rstrip("MB")) * MB
        elif bytes_int_or_str.endswith("GB"):
            return int(bytes_int_or_str.rstrip("GB")) * GB
        else:
            return int(bytes_int_or_str)
    else:
        raise ValueError(f"Invalid type: {type(bytes_int_or_str)=}")

src/lm_datasets/utils/__init__.py
malteos/lm-datasets
# This file contains the settings for the project.
import logging


LANGUAGES = {
    # 24 languages
    # bg,cs,da,de,el,en,es,et,fi,fr,ga,hr,hu,it,lt,lv,mt,nl,pl,pt,ro,sk,sl,sv
    # Bulgarian, Czech, Danish, German, Greek, English, Spanish, Estonian, Finnish, French, Irish, Croatian, Hungarian, Italian, Lithuanian, Latvian, Maltese, Dutch, Polish, Portuguese, Romanian, Slovak, Slovenian, Swedish  # noqa
    "bg": "Bulgarian",
    "cs": "Czech",
    "da": "Danish",
    "de": "German",
    "el": "Greek",
    "en": "English",
    "es": "Spanish",
    "et": "Estonian",
    "fi": "Finnish",
    "fr": "French",
    "ga": "Irish",
    "hr": "Croatian",
    "hu": "Hungarian",
    "it": "Italian",
    "lt": "Lithuanian",
    "lv": "Latvian",
    "mt": "Maltese",
    "nl": "Dutch",
    "pl": "Polish",
    "pt": "Portuguese",
    "ro": "Romanian",
    "sk": "Slovak",
    "sl": "Slovenian",
    "sv": "Swedish",
}


HIGH_LANGUAGES = set("en,de,fr,es".split(","))
# MID_LANGUAGES = set(''.split(','))
# LOW_LANGUAGES = set(''.split(','))

# Max number of dumps to download at once (hard limit for Wikimedia dumps)
MAX_DOWNLOADS = 2

# wiki,wikibooks,wikiquote,wikinews,wikisource,wikivoyage
WIKI_TYPES = [
    "wiki",  # Wikipedia
    "wikibooks",  # Wikibooks
    "wikiquote",  # Wikiquote
    "wikinews",  # Wikinews
    "wikisource",  # Wikisource
    "wikivoyage",  # Wikivoyage
]

MIN_TEXT_LENGTH_BY_WIKI_TYPE = {
    "wiki": 5_000,
    "wikibooks": 1_000,
    "wikiquote": 500,
    "wikinews": 500,
    "wikisource": 500,
    "wikivoyage": 500,
}
WIKI_MIN_TEXT_LENGTH = 5_000

OSCAR_MIN_HARMFUL_PP_BY_LANGUAGE = {
    "en": 10.0,
    "fr": 10.0,
    "es": 10.0,
    "de": 10.0,
}
DEFAULT_OSCAR_MIN_HARMFUL_PP = 20.0

OSCAR_MIN_LANGUAGE_PROB = 0.85

LOGGING_KWARGS = dict(
    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
    handlers=[logging.StreamHandler()],
)

# GPT-2 1.5B is trained with 40GB of Internet text, which is roughly 10 Billion tokens
# BILLION_TOKENS_PER_GIGA_BYTE = 4
BILLION_TOKENS_PER_GIGA_BYTE = 0.2  # (based on German OSCAR dataset)

DEFAULT_VOCAB_SIZE = 159900
DEFAULT_TOKENIZER_RATIO = 0.10
DEFAULT_TEST_RATIO = 0.01

src/lm_datasets/utils/settings.py
malteos/lm-datasets
"""
Take from https://raw.githubusercontent.com/YaYaB/shuffle-big-file/master/shuffle_big_file/shuffle_big_file.py

This tool helps you shuffle by line a big file that does not fit in memory. Given the
batch_size that your machine can put in memory it will shuffle the whole file by reading
as many times as necessary.
"""
import os
from datetime import datetime, timedelta
import math
import random
import argparse
import sys
import logging

RANDOM_SEED = 1337

logger = logging.getLogger(__name__)


def get_args():
    parser = argparse.ArgumentParser("Shuffle a huge file without loading in fully in RAM")
    parser.add_argument("--input_file", type=str, help="Path to input file")
    parser.add_argument("--batch_size", type=int, help="Batch size that can fit in memory")
    parser.add_argument("--output_file", type=str, help="Path to output shuffled file")
    parser.add_argument("--seed", type=int, help="Seed for random generator", default=RANDOM_SEED)

    args = parser.parse_args()
    return args


def compute_nb_read(path_file, batch_size):
    """
    :param path_file: path file for which we count the number of lines
    :param batch_size: batch size that the machine can handle
    :return: number of lines and number of read necessary
    """
    # Read all file for the first time to know how many read we need
    with open(path_file) as f:
        for i, l in enumerate(f):
            pass

    nb_lines = i + 1
    nb_read = math.ceil(nb_lines / batch_size)

    return nb_read, nb_lines


def compute_shuffled_index(nb_lines):
    """
    :param nb_lines: number of lines
    :return: shuffled index of the number of lines
    """
    # Create index
    idx_shuffled = list(range(0, nb_lines))

    # Shuffle the index
    random.shuffle(idx_shuffled)

    return idx_shuffled


def generate_random_string(
    length=500, alphabet="azertyuiop^$qsdfghjklmù*wxcvbn,;:!1234567890°+¨£%µ§/.?<>AZERTYUIOPQSDFGHJKLMWXCVBN&é'(-è_çà)="
):
    """
    :param length: length of string wanted
    :param alphabet: alphabet from witch characters are sampled
    :return: random string of size `length`
    """
    return "".join(random.choice(alphabet) for x in range(length))


def generate_random_file(path_file="./random_file.txt", nb_lines=10000, max_line_length=1000):
    """
    :param path_file: path file that will be created
    :param nb_lines: number of lines wanted
    :param max_lines_length: maximum number of characters per line
    """
    with open(path_file, "w") as f:
        for i in range(0, nb_lines):
            f.write(generate_random_string(max_line_length) + "\n")


def generate_random_file_cli():
    parser = argparse.ArgumentParser("generate random file containing a string per line")
    parser.add_argument("--path_file", type=str, default="./random_file.txt", help="Path to file that will be created")
    parser.add_argument("--nb_lines", type=int, default=10000, help="Number of lines wanted for the file")
    parser.add_argument("--max_line_length", type=int, default=1000, help="Max number of characters per line")
    opt = parser.parse_args()

    generate_random_file(opt.path_file, opt.nb_lines, opt.max_line_length)


def shuffle(input_file, output_file, idx_shuffled, nb_read, batch_size):
    """
    :param input_file: path file that will be created
    :param output_file: number of lines wanted
    :param idx_shuffled: index shuffled
    :param nb_read: number of read necessary
    :param batch_size: batch size used

    """
    with open(output_file, "a") as f:
        for i in range(nb_read):
            startTime = datetime.now()
            elements = []
            idx = idx_shuffled[i * batch_size : (i + 1) * batch_size]
            # get index of batch in order
            idx_sort = sorted(idx)
            waiting_line = idx_sort[0]

            j = 0
            with open(input_file) as in_f:
                for k, l in enumerate(in_f):
                    # If line in batch add it to list of elements
                    if k == waiting_line:
                        elements.append(l)
                        j += 1
                        # Try to access element to the idx sort
                        try:
                            waiting_line = idx_sort[j]
                        # If fails that means it exceds the max_line
                        except:
                            break

            # Shuffle the batch created
            random.shuffle(elements)

            # Append output file
            for elem in elements:
                f.write(elem)

            logger.info(
                f"{nb_read - (i+1)} read are left. It will take"
                f" {timedelta(seconds=(nb_read - (i+1)) * (datetime.now() - startTime).total_seconds())}"
            )


def shuffle_big_file(input_file, output_file, batch_size, seed=RANDOM_SEED):
    """
    :param input_file: path of the inputfile
    :param batch_size: batch size used
    :param output_file: path of the output file
    :param seed: seed for the random generator
    """

    assert os.path.exists(input_file)
    assert not os.path.exists(output_file)

    random.seed(seed)

    # Compute number of read necessary
    startTime_ = datetime.now()
    nb_read, nb_lines = compute_nb_read(input_file, batch_size)
    one_read = datetime.now() - startTime_

    logger.info(
        f"File was read in {one_read} and {nb_read} reads are necessary. It will take"
        f" {timedelta(seconds=nb_read * one_read.total_seconds())}"
    )

    # Shuffle the index of the number of lines
    idx_shuffled = compute_shuffled_index(nb_lines)

    # Shuffle file
    shuffle(input_file, output_file, idx_shuffled, nb_read, batch_size)

    logger.info(f"File was shuffled in {datetime.now() - startTime_}")

    return 0


def shuffle_big_file_cli():
    """
    Main function of the program
    """

    # Get arguments passed in CLI
    opt = get_args()

    return shuffle_big_file(opt.input_file, opt.output_file, opt.batch_size, opt.seed)


if __name__ == "__main__":
    try:
        sys.exit(shuffle_big_file_cli())
    except Exception as e:
        print("[ERR] Uncaught error waiting for scripts to finish")
        print(e)
        raise

src/lm_datasets/utils/shuffle_big_file.py
malteos/lm-datasets
"""
Taken from https://huggingface.co/datasets/wikipedia/blob/main/wikipedia.py
"""
import re

# Source: https://en.wikipedia.org/wiki/List_of_Wikipedias (accessed 3/1/2019)
# Removed because no articles: hz.
WIKIPEDIA_LANGUAGES = [
    "aa",
    "ab",
    "ace",
    "ady",
    "af",
    "ak",
    "als",
    "am",
    "an",
    "ang",
    "ar",
    "arc",
    "arz",
    "as",
    "ast",
    "atj",
    "av",
    "ay",
    "az",
    "azb",
    "ba",
    "bar",
    "bat-smg",
    "bcl",
    "be",
    "be-x-old",
    "bg",
    "bh",
    "bi",
    "bjn",
    "bm",
    "bn",
    "bo",
    "bpy",
    "br",
    "bs",
    "bug",
    "bxr",
    "ca",
    "cbk-zam",
    "cdo",
    "ce",
    "ceb",
    "ch",
    "cho",
    "chr",
    "chy",
    "ckb",
    "co",
    "cr",
    "crh",
    "cs",
    "csb",
    "cu",
    "cv",
    "cy",
    "da",
    "de",
    "din",
    "diq",
    "dsb",
    "dty",
    "dv",
    "dz",
    "ee",
    "el",
    "eml",
    "en",
    "eo",
    "es",
    "et",
    "eu",
    "ext",
    "fa",
    "ff",
    "fi",
    "fiu-vro",
    "fj",
    "fo",
    "fr",
    "frp",
    "frr",
    "fur",
    "fy",
    "ga",
    "gag",
    "gan",
    "gd",
    "gl",
    "glk",
    "gn",
    "gom",
    "gor",
    "got",
    "gu",
    "gv",
    "ha",
    "hak",
    "haw",
    "he",
    "hi",
    "hif",
    "ho",
    "hr",
    "hsb",
    "ht",
    "hu",
    "hy",
    "ia",
    "id",
    "ie",
    "ig",
    "ii",
    "ik",
    "ilo",
    "inh",
    "io",
    "is",
    "it",
    "iu",
    "ja",
    "jam",
    "jbo",
    "jv",
    "ka",
    "kaa",
    "kab",
    "kbd",
    "kbp",
    "kg",
    "ki",
    "kj",
    "kk",
    "kl",
    "km",
    "kn",
    "ko",
    "koi",
    "krc",
    "ks",
    "ksh",
    "ku",
    "kv",
    "kw",
    "ky",
    "la",
    "lad",
    "lb",
    "lbe",
    "lez",
    "lfn",
    "lg",
    "li",
    "lij",
    "lmo",
    "ln",
    "lo",
    "lrc",
    "lt",
    "ltg",
    "lv",
    "mai",
    "map-bms",
    "mdf",
    "mg",
    "mh",
    "mhr",
    "mi",
    "min",
    "mk",
    "ml",
    "mn",
    "mr",
    "mrj",
    "ms",
    "mt",
    "mus",
    "mwl",
    "my",
    "myv",
    "mzn",
    "na",
    "nah",
    "nap",
    "nds",
    "nds-nl",
    "ne",
    "new",
    "ng",
    "nl",
    "nn",
    "no",
    "nov",
    "nrm",
    "nso",
    "nv",
    "ny",
    "oc",
    "olo",
    "om",
    "or",
    "os",
    "pa",
    "pag",
    "pam",
    "pap",
    "pcd",
    "pdc",
    "pfl",
    "pi",
    "pih",
    "pl",
    "pms",
    "pnb",
    "pnt",
    "ps",
    "pt",
    "qu",
    "rm",
    "rmy",
    "rn",
    "ro",
    "roa-rup",
    "roa-tara",
    "ru",
    "rue",
    "rw",
    "sa",
    "sah",
    "sat",
    "sc",
    "scn",
    "sco",
    "sd",
    "se",
    "sg",
    "sh",
    "si",
    "simple",
    "sk",
    "sl",
    "sm",
    "sn",
    "so",
    "sq",
    "sr",
    "srn",
    "ss",
    "st",
    "stq",
    "su",
    "sv",
    "sw",
    "szl",
    "ta",
    "tcy",
    "te",
    "tet",
    "tg",
    "th",
    "ti",
    "tk",
    "tl",
    "tn",
    "to",
    "tpi",
    "tr",
    "ts",
    "tt",
    "tum",
    "tw",
    "ty",
    "tyv",
    "udm",
    "ug",
    "uk",
    "ur",
    "uz",
    "ve",
    "vec",
    "vep",
    "vi",
    "vls",
    "vo",
    "wa",
    "war",
    "wo",
    "wuu",
    "xal",
    "xh",
    "xmf",
    "yi",
    "yo",
    "za",
    "zea",
    "zh",
    "zh-classical",
    "zh-min-nan",
    "zh-yue",
    "zu",
]

# Source: for each Wikipedia language code (example shown for "ab"), aliases for namespaces -2 and 6 accessed
# via this API call: [1] (accessed 12/21/2021)
# [1]: https://ab.wikipedia.org/w/api.php?action=query&meta=siteinfo&siprop=namespacealiases|namespaces&format=json&formatversion=2 # noqa
MEDIA_ALIASES = {
    "ab": ["Медиа", "Файл", "Афаил", "Амедиа", "Изображение"],
    "ace": ["Beureukaih", "Gambar", "Alat", "Berkas"],
    "ady": ["Медиа"],
    "af": ["Lêer", "Beeld"],
    "als": ["Medium", "Datei", "Bild"],
    "am": ["ፋይል", "ስዕል"],
    "an": ["Imachen", "Imagen"],
    "ang": ["Ymele", "Biliþ"],
    "ar": ["ميديا", "صورة", "وسائط", "ملف"],
    "arc": ["ܠܦܦܐ", "ܡܝܕܝܐ"],
    "arz": ["ميديا", "صورة", "وسائط", "ملف"],
    "as": ["চিত্ৰ", "चित्र", "চিত্র", "মাধ্যম"],
    "ast": ["Imaxen", "Ficheru", "Imaxe", "Archivu", "Imagen", "Medios"],
    "atj": ["Tipatcimoctakewin", "Natisinahikaniwoc"],
    "av": ["Медиа", "Файл", "Изображение"],
    "ay": ["Medio", "Archivo", "Imagen"],
    "az": ["Mediya", "Şəkil", "Fayl"],
    "azb": ["رسانه", "تصویر", "مدیا", "فایل", "رسانه‌ای"],
    "ba": ["Медиа", "Рәсем", "Файл", "Изображение"],
    "bar": ["Medium", "Datei", "Bild"],
    "bat-smg": ["Vaizdas", "Medėjė", "Abruozdielis"],
    "bcl": ["Medio", "Ladawan"],
    "be": ["Мультымедыя", "Файл", "Выява"],
    "be-x-old": ["Мэдыя", "Файл", "Выява"],
    "bg": ["Медия", "Файл", "Картинка"],
    "bh": ["मीडिया", "चित्र"],
    "bjn": ["Barakas", "Gambar", "Berkas"],
    "bm": ["Média", "Fichier"],
    "bn": ["চিত্র", "মিডিয়া"],
    "bpy": ["ছবি", "মিডিয়া"],
    "br": ["Skeudenn", "Restr"],
    "bs": ["Mediji", "Slika", "Datoteka", "Medija"],
    "bug": ["Gambar", "Berkas"],
    "bxr": ["Файл", "Меди", "Изображение"],
    "ca": ["Fitxer", "Imatge"],
    "cbk-zam": ["Medio", "Archivo", "Imagen"],
    "cdo": ["文件", "媒體", "圖像", "檔案"],
    "ce": ["Хlум", "Медиа", "Сурт", "Файл", "Медйа", "Изображение"],
    "ceb": ["Payl", "Medya", "Imahen"],
    "ch": ["Litratu"],
    "ckb": ["میدیا", "پەڕگە"],
    "co": ["Immagine"],
    "crh": ["Медиа", "Resim", "Файл", "Fayl", "Ресим"],
    "cs": ["Soubor", "Média", "Obrázok"],
    "csb": ["Òbrôzk", "Grafika"],
    "cu": ["Видъ", "Ви́дъ", "Дѣло", "Срѣдьства"],
    "cv": ["Медиа", "Ӳкерчĕк", "Изображение"],
    "cy": ["Delwedd"],
    "da": ["Billede", "Fil"],
    "de": ["Medium", "Datei", "Bild"],
    "din": ["Ciɛl", "Apamduööt"],
    "diq": ["Medya", "Dosya"],
    "dsb": ["Wobraz", "Dataja", "Bild", "Medija"],
    "dty": ["चित्र", "मिडिया"],
    "dv": ["ފައިލު", "މީޑިއާ", "ފައިލް"],
    "el": ["Εικόνα", "Αρχείο", "Μέσο", "Μέσον"],
    "eml": ["Immagine"],
    "eo": ["Dosiero", "Aŭdvidaĵo"],
    "es": ["Medio", "Archivo", "Imagen"],
    "et": ["Pilt", "Fail", "Meedia"],
    "eu": ["Irudi", "Fitxategi"],
    "ext": ["Archivu", "Imagen", "Mediu"],
    "fa": ["رسانه", "تصویر", "مدیا", "پرونده", "رسانه‌ای"],
    "ff": ["Média", "Fichier"],
    "fi": ["Kuva", "Tiedosto"],
    "fiu-vro": ["Pilt", "Meediä"],
    "fo": ["Miðil", "Mynd"],
    "fr": ["Média", "Fichier"],
    "frp": ["Émâge", "Fichiér", "Mèdia"],
    "frr": ["Medium", "Datei", "Bild"],
    "fur": ["Immagine", "Figure"],
    "fy": ["Ofbyld"],
    "ga": ["Íomhá", "Meán"],
    "gag": ["Mediya", "Medya", "Resim", "Dosya", "Dosye"],
    "gan": ["媒体文件", "文件", "文檔", "档案", "媒體", "图像", "圖像", "媒体", "檔案"],
    "gd": ["Faidhle", "Meadhan"],
    "gl": ["Imaxe", "Ficheiro", "Arquivo", "Imagem"],
    "glk": ["رسانه", "تصویر", "پرونده", "فاىل", "رسانه‌ای", "مديا"],
    "gn": ["Medio", "Imagen", "Ta'ãnga"],
    "gom": ["माध्यम", "मिडिया", "फायल"],
    "gor": ["Gambar", "Berkas"],
    "got": ["𐍆𐌴𐌹𐌻𐌰"],
    "gu": ["દ્રશ્ય-શ્રાવ્ય (મિડિયા)", "દ્રશ્ય-શ્રાવ્ય_(મિડિયા)", "ચિત્ર"],
    "gv": ["Coadan", "Meanyn"],
    "hak": ["文件", "媒體", "圖像", "檔案"],
    "haw": ["Kiʻi", "Waihona", "Pāpaho"],
    "he": ["תמונה", "קו", "מדיה", "קובץ"],
    "hi": ["मीडिया", "चित्र"],
    "hif": ["file", "saadhan"],
    "hr": ["Mediji", "DT", "Slika", "F", "Datoteka"],
    "hsb": ["Wobraz", "Dataja", "Bild"],
    "ht": ["Imaj", "Fichye", "Medya"],
    "hu": ["Kép", "Fájl", "Média"],
    "hy": ["Պատկեր", "Մեդիա"],
    "ia": ["Imagine", "Multimedia"],
    "id": ["Gambar", "Berkas"],
    "ig": ["Nká", "Midia", "Usòrò", "Ákwúkwó orünotu", "Ákwúkwó_orünotu"],
    "ii": ["媒体文件", "文件", "档案", "图像", "媒体"],
    "ilo": ["Midia", "Papeles"],
    "inh": ["Медиа", "Файл", "Изображение"],
    "io": ["Imajo", "Arkivo"],
    "is": ["Miðill", "Mynd"],
    "it": ["Immagine"],
    "ja": ["メディア", "ファイル", "画像"],
    "jbo": ["velsku", "datnyvei"],
    "jv": ["Barkas", "Medhia", "Gambar", "Médhia"],
    "ka": ["მედია", "სურათი", "ფაილი"],
    "kaa": ["Swret", "Таспа", "سۋرەت", "Taspa", "Su'wret", "Сурет", "تاسپا"],
    "kab": ["Tugna"],
    "kbd": ["Медиа", "Файл"],
    "kbp": ["Média", "Fichier"],
    "kg": ["Fisye"],
    "kk": ["Swret", "سۋرەت", "Таспа", "Taspa", "Сурет", "تاسپا"],
    "kl": ["Billede", "Fiileq", "Fil"],
    "km": ["ឯកសារ", "រូបភាព", "មេឌា", "មីឌា"],
    "kn": ["ಚಿತ್ರ", "ಮೀಡಿಯ"],
    "ko": ["미디어", "파일", "그림"],
    "koi": ["Медиа", "Файл", "Изображение"],
    "krc": ["Медиа", "Файл", "Изображение"],
    "ks": ["میڈیا", "فَیِل"],
    "ksh": ["Beld", "Meedije", "Medie", "Belld", "Medium", "Datei", "Meedijum", "Bild"],
    "ku": ["میدیا", "پەڕگە", "Medya", "Wêne"],
    "kv": ["Медиа", "Файл", "Изображение"],
    "kw": ["Restren"],
    "ky": ["Медиа", "Файл"],
    "la": ["Imago", "Fasciculus"],
    "lad": ["Dossia", "Medya", "Archivo", "Dosya", "Imagen", "Meddia"],
    "lb": ["Fichier", "Bild"],
    "lbe": ["Медиа", "Сурат", "Изображение"],
    "lez": ["Медиа", "Mediya", "Файл", "Şəkil", "Изображение"],
    "lfn": ["Fix"],
    "li": ["Afbeelding", "Plaetje", "Aafbeilding"],
    "lij": ["Immaggine", "Immagine"],
    "lmo": ["Immagine", "Imàjine", "Archivi"],
    "ln": ["Média", "Fichier"],
    "lo": ["ສື່ອ", "ສື່", "ຮູບ"],
    "lrc": ["رسانه", "تصویر", "رسانه‌ای", "جانیا", "أسگ", "ڤارئسگأر"],
    "lt": ["Vaizdas", "Medija"],
    "ltg": ["Medeja", "Fails"],
    "lv": ["Attēls"],
    "mai": ["मेडिया", "फाइल"],
    "map-bms": ["Barkas", "Medhia", "Gambar", "Médhia"],
    "mdf": ["Медиа", "Няйф", "Изображение"],
    "mg": ["Rakitra", "Sary", "Média"],
    "mhr": ["Медиа", "Файл", "Изображение"],
    "min": ["Gambar", "Berkas"],
    "mk": ["Податотека", "Медија", "Медиум", "Слика"],
    "ml": ["പ്രമാണം", "ചി", "മീഡിയ", "പ്ര", "ചിത്രം"],
    "mn": ["Медиа", "Файл", "Зураг"],
    "mr": ["चित्र", "मिडिया"],
    "mrj": ["Медиа", "Файл", "Изображение"],
    "ms": ["Fail", "Imej"],
    "mt": ["Midja", "Medja", "Stampa"],
    "mwl": ["Multimédia", "Fexeiro", "Ficheiro", "Arquivo", "Imagem"],
    "my": ["ဖိုင်", "မီဒီယာ"],
    "myv": ["Медия", "Артовкс", "Изображение"],
    "mzn": ["رسانه", "تصویر", "مه‌دیا", "مدیا", "پرونده", "رسانه‌ای"],
    "nah": ["Mēdiatl", "Īxiptli", "Imagen"],
    "nap": ["Fiùra", "Immagine"],
    "nds": ["Datei", "Bild"],
    "nds-nl": ["Ofbeelding", "Afbeelding", "Bestaand"],
    "ne": ["मीडिया", "चित्र"],
    "new": ["किपा", "माध्यम"],
    "nl": ["Bestand", "Afbeelding"],
    "nn": ["Fil", "Bilde", "Filpeikar"],
    "no": ["Fil", "Medium", "Bilde"],
    "nov": [],
    "nrm": ["Média", "Fichier"],
    "nso": ["Seswantšho"],
    "nv": ["Eʼelyaaígíí"],
    "oc": ["Imatge", "Fichièr", "Mèdia"],
    "olo": ["Kuva", "Medii", "Failu"],
    "or": ["ମାଧ୍ୟମ", "ଫାଇଲ"],
    "os": ["Ныв", "Медиа", "Файл", "Изображение"],
    "pa": ["ਤਸਵੀਰ", "ਮੀਡੀਆ"],
    "pcd": ["Média", "Fichier"],
    "pdc": ["Medium", "Datei", "Bild", "Feil"],
    "pfl": ["Dadai", "Medium", "Datei", "Bild"],
    "pi": ["मीडिया", "पटिमा"],
    "pl": ["Plik", "Grafika"],
    "pms": ["Figura", "Immagine"],
    "pnb": ["میڈیا", "تصویر", "فائل"],
    "pnt": ["Εικόνα", "Αρχείον", "Εικόναν", "Μέσον"],
    "ps": ["انځور", "رسنۍ", "دوتنه"],
    "pt": ["Multimédia", "Ficheiro", "Arquivo", "Imagem"],
    "qu": ["Midya", "Imagen", "Rikcha"],
    "rm": ["Multimedia", "Datoteca"],
    "rmy": ["Fişier", "Mediya", "Chitro", "Imagine"],
    "ro": ["Fişier", "Imagine", "Fișier"],
    "roa-rup": ["Fişier", "Imagine", "Fișier"],
    "roa-tara": ["Immagine"],
    "ru": ["Медиа", "Файл", "Изображение"],
    "rue": ["Медіа", "Медиа", "Файл", "Изображение", "Зображення"],
    "rw": ["Dosiye", "Itangazamakuru"],
    "sa": ["चित्रम्", "माध्यमम्", "सञ्चिका", "माध्यम", "चित्रं"],
    "sah": ["Миэдьийэ", "Ойуу", "Билэ", "Изображение"],
    "sat": ["ᱨᱮᱫ", "ᱢᱤᱰᱤᱭᱟ"],
    "sc": ["Immàgini"],
    "scn": ["Immagine", "Mmàggini", "Mèdia"],
    "sd": ["عڪس", "ذريعات", "فائل"],
    "se": ["Fiila"],
    "sg": ["Média", "Fichier"],
    "sh": ["Mediji", "Slika", "Медија", "Datoteka", "Medija", "Слика"],
    "si": ["රූපය", "මාධ්‍යය", "ගොනුව"],
    "sk": ["Súbor", "Obrázok", "Médiá"],
    "sl": ["Slika", "Datoteka"],
    "sq": ["Figura", "Skeda"],
    "sr": ["Датотека", "Medij", "Slika", "Медија", "Datoteka", "Медиј", "Medija", "Слика"],
    "srn": ["Afbeelding", "Gefre"],
    "stq": ["Bielde", "Bild"],
    "su": ["Média", "Gambar"],
    "sv": ["Fil", "Bild"],
    "sw": ["Faili", "Picha"],
    "szl": ["Plik", "Grafika"],
    "ta": ["படிமம்", "ஊடகம்"],
    "tcy": ["ಮಾದ್ಯಮೊ", "ಫೈಲ್"],
    "te": ["ఫైలు", "దస్త్రం", "బొమ్మ", "మీడియా"],
    "tet": ["Imajen", "Arquivo", "Imagem"],
    "tg": ["Акс", "Медиа"],
    "th": ["ไฟล์", "สื่อ", "ภาพ"],
    "ti": ["ፋይል", "ሜድያ"],
    "tk": ["Faýl"],
    "tl": ["Midya", "Talaksan"],
    "tpi": ["Fail"],
    "tr": ["Medya", "Resim", "Dosya", "Ortam"],
    "tt": ["Медиа", "Рәсем", "Файл", "Räsem", "Изображение"],
    "ty": ["Média", "Fichier"],
    "tyv": ["Медиа", "Файл", "Изображение"],
    "udm": ["Медиа", "Файл", "Суред", "Изображение"],
    "ug": ["ۋاسىتە", "ھۆججەت"],
    "uk": ["Медіа", "Медиа", "Файл", "Изображение", "Зображення"],
    "ur": ["میڈیا", "تصویر", "وسیط", "زریعہ", "فائل", "ملف"],
    "uz": ["Mediya", "Tasvir", "Fayl"],
    "vec": ["Immagine", "Imàjine", "Mèdia"],
    "vep": ["Pilt", "Fail"],
    "vi": ["Phương_tiện", "Tập_tin", "Hình", "Tập tin", "Phương tiện"],
    "vls": ["Afbeelding", "Ofbeeldienge"],
    "vo": ["Ragiv", "Magod", "Nünamakanäd"],
    "wa": ["Imådje"],
    "war": ["Medya", "Fayl", "Paypay"],
    "wo": ["Xibaarukaay", "Dencukaay"],
    "wuu": ["文件", "档案", "图像", "媒体"],
    "xal": ["Аһар", "Боомг", "Изображение", "Зург"],
    "xmf": ["მედია", "სურათი", "ფაილი"],
    "yi": ["מעדיע", "תמונה", "טעקע", "בילד"],
    "yo": ["Fáìlì", "Amóhùnmáwòrán", "Àwòrán"],
    "za": ["媒体文件", "文件", "档案", "图像", "媒体"],
    "zea": ["Afbeelding", "Plaetje"],
    "zh": ["媒体文件", "F", "文件", "媒體", "档案", "图像", "圖像", "媒体", "檔案"],
    "zh-classical": ["文件", "媒體", "圖像", "檔案"],
    "zh-min-nan": ["tóng-àn", "文件", "媒體", "Mûi-thé", "圖像", "檔案"],
    "zh-yue": ["檔", "档", "文件", "图", "媒體", "圖", "档案", "图像", "圖像", "媒体", "檔案"],
}

# Source: for each Wikipedia language code (example shown for "ab"), aliases for namespace 14 accessed
# via this API call: [1]
# [1]: https://ab.wikipedia.org/w/api.php?action=query&meta=siteinfo&siprop=namespacealiases|namespaces&format=json&formatversion=2 # noqa
CAT_ALIASES = {
    "ab": ["Категория", "Акатегориа"],
    "ace": ["Kawan", "Kategori"],
    "af": ["Kategorie"],
    "ak": ["Nkyekyem"],
    "als": ["Kategorie"],
    "am": ["መደብ"],
    "an": ["Categoría"],
    "ang": ["Flocc"],
    "ar": ["تصنيف"],
    "arc": ["ܣܕܪܐ"],
    "arz": ["تصنيف"],
    "as": ["CAT", "শ্ৰেণী", "श्रेणी", "শ্রেণী"],
    "ast": ["Categoría"],
    "atj": ["Tipanictawin"],
    "av": ["Категория"],
    "ay": ["Categoría"],
    "az": ["Kateqoriya"],
    "azb": ["بؤلمه"],
    "ba": ["Төркөм", "Категория"],
    "bar": ["Kategorie"],
    "bat-smg": ["Kategorija", "Kateguorėjė"],
    "bcl": ["Kategorya"],
    "be": ["Катэгорыя"],
    "be-x-old": ["Катэгорыя"],
    "bg": ["Категория"],
    "bh": ["श्रेणी"],
    "bjn": ["Tumbung", "Kategori"],
    "bm": ["Catégorie"],
    "bn": ["বিষয়শ্রেণী", "വിഭാഗം"],
    "bpy": ["থাক"],
    "br": ["Rummad"],
    "bs": ["Kategorija"],
    "bug": ["Kategori"],
    "bxr": ["Категори", "Категория"],
    "ca": ["Categoria"],
    "cbk-zam": ["Categoría"],
    "cdo": ["分類"],
    "ce": ["Категори", "Тоба", "Кадегар"],
    "ceb": ["Kategoriya"],
    "ch": ["Katigoria"],
    "ckb": ["پ", "پۆل"],
    "co": ["Categoria"],
    "crh": ["Категория", "Kategoriya"],
    "cs": ["Kategorie"],
    "csb": ["Kategòrëjô"],
    "cu": ["Катигорї", "Категория", "Катигорїꙗ"],
    "cv": ["Категори"],
    "cy": ["Categori"],
    "da": ["Kategori"],
    "de": ["Kategorie"],
    "din": ["Bekätakthook"],
    "diq": ["Kategoriye", "Kategori"],
    "dsb": ["Kategorija"],
    "dty": ["श्रेणी"],
    "dv": ["ޤިސްމު"],
    "el": ["Κατηγορία"],
    "eml": ["Categoria"],
    "eo": ["Kategorio"],
    "es": ["CAT", "Categoría"],
    "et": ["Kategooria"],
    "eu": ["Kategoria"],
    "ext": ["Categoría", "Categoria"],
    "fa": ["رده"],
    "ff": ["Catégorie"],
    "fi": ["Luokka"],
    "fiu-vro": ["Katõgooria"],
    "fo": ["Bólkur"],
    "fr": ["Catégorie"],
    "frp": ["Catègorie"],
    "frr": ["Kategorie"],
    "fur": ["Categorie"],
    "fy": ["Kategory"],
    "ga": ["Rang", "Catagóir"],
    "gag": ["Kategori", "Kategoriya"],
    "gan": ["分類", "分类"],
    "gd": ["Roinn-seòrsa"],
    "gl": ["Categoría"],
    "glk": ["جرگه", "رده"],
    "gn": ["Ñemohenda"],
    "gom": ["वर्ग", "श्रेणी"],
    "gor": ["Dalala"],
    "got": ["𐌷𐌰𐌽𐍃𐌰"],
    "gu": ["શ્રેણી", "CAT", "શ્રે"],
    "gv": ["Ronney"],
    "hak": ["分類"],
    "haw": ["Māhele"],
    "he": ["קטגוריה", "קט"],
    "hi": ["श्र", "श्रेणी"],
    "hif": ["vibhag"],
    "hr": ["CT", "KT", "Kategorija"],
    "hsb": ["Kategorija"],
    "ht": ["Kategori"],
    "hu": ["Kategória"],
    "hy": ["Կատեգորիա"],
    "ia": ["Categoria"],
    "id": ["Kategori"],
    "ie": ["Categorie"],
    "ig": ["Ébéonọr", "Òtù"],
    "ii": ["分类"],
    "ilo": ["Kategoria"],
    "inh": ["ОагӀат"],
    "io": ["Kategorio"],
    "is": ["Flokkur"],
    "it": ["CAT", "Categoria"],
    "ja": ["カテゴリ"],
    "jbo": ["klesi"],
    "jv": ["Kategori"],
    "ka": ["კატეგორია"],
    "kaa": ["Sanat", "Kategoriya", "Санат", "سانات"],
    "kab": ["Taggayt"],
    "kbd": ["Категория", "Категориэ"],
    "kbp": ["Catégorie"],
    "kg": ["Kalasi"],
    "kk": ["Sanat", "Санат", "سانات"],
    "kl": ["Sumut_atassuseq", "Kategori", "Sumut atassuseq"],
    "km": ["ចំនាត់ថ្នាក់ក្រុម", "ចំណាត់ក្រុម", "ចំណាត់ថ្នាក់ក្រុម"],
    "kn": ["ವರ್ಗ"],
    "ko": ["분류"],
    "koi": ["Категория"],
    "krc": ["Категория"],
    "ks": ["زٲژ"],
    "ksh": ["Saachjropp", "Saachjrop", "Katejori", "Kategorie", "Saachjrupp", "Kattejori", "Sachjrop"],
    "ku": ["Kategorî", "پۆل"],
    "kv": ["Категория"],
    "kw": ["Class", "Klass"],
    "ky": ["Категория"],
    "la": ["Categoria"],
    "lad": ["Kateggoría", "Katēggoría", "Categoría"],
    "lb": ["Kategorie"],
    "lbe": ["Категория"],
    "lez": ["Категория"],
    "lfn": ["Categoria"],
    "li": ["Categorie", "Kategorie"],
    "lij": ["Categorîa", "Categoria"],
    "lmo": ["Categuria", "Categoria"],
    "ln": ["Catégorie"],
    "lo": ["ໝວດ"],
    "lrc": ["دأسە"],
    "lt": ["Kategorija"],
    "ltg": ["Kategoreja"],
    "lv": ["Kategorija"],
    "mai": ["CA", "श्रेणी"],
    "map-bms": ["Kategori"],
    "mdf": ["Категорие", "Категория"],
    "mg": ["Sokajy", "Catégorie"],
    "mhr": ["Категория", "Категорий"],
    "min": ["Kategori"],
    "mk": ["Категорија"],
    "ml": ["വിഭാഗം", "വി", "വർഗ്ഗം", "വ"],
    "mn": ["Ангилал"],
    "mr": ["वर्ग"],
    "mrj": ["Категори", "Категория"],
    "ms": ["Kategori"],
    "mt": ["Kategorija"],
    "mwl": ["Catadorie", "Categoria"],
    "my": ["ကဏ္ဍ"],
    "myv": ["Категория"],
    "mzn": ["رج", "رده"],
    "nah": ["Neneuhcāyōtl", "Categoría"],
    "nap": ["Categurìa", "Categoria"],
    "nds": ["Kategorie"],
    "nds-nl": ["Categorie", "Kattegerie", "Kategorie"],
    "ne": ["श्रेणी"],
    "new": ["पुचः"],
    "nl": ["Categorie"],
    "nn": ["Kategori"],
    "no": ["Kategori"],
    "nrm": ["Catégorie"],
    "nso": ["Setensele"],
    "nv": ["Tʼááłáhági_átʼéego", "Tʼááłáhági átʼéego"],
    "oc": ["Categoria"],
    "olo": ["Kategourii"],
    "or": ["ବିଭାଗ", "ଶ୍ରେଣୀ"],
    "os": ["Категори"],
    "pa": ["ਸ਼੍ਰੇਣੀ"],
    "pcd": ["Catégorie"],
    "pdc": ["Abdeeling", "Kategorie"],
    "pfl": ["Kadegorie", "Sachgrubb", "Kategorie"],
    "pi": ["विभाग"],
    "pl": ["Kategoria"],
    "pms": ["Categorìa"],
    "pnb": ["گٹھ"],
    "pnt": ["Κατηγορίαν"],
    "ps": ["وېشنيزه"],
    "pt": ["Categoria"],
    "qu": ["Katiguriya"],
    "rm": ["Categoria"],
    "rmy": ["Shopni"],
    "ro": ["Categorie"],
    "roa-rup": ["Categorie"],
    "roa-tara": ["Categoria"],
    "ru": ["Категория", "К"],
    "rue": ["Категория", "Катеґорія"],
    "rw": ["Ikiciro"],
    "sa": ["वर्गः"],
    "sah": ["Категория"],
    "sat": ["ᱛᱷᱚᱠ"],
    "sc": ["Categoria"],
    "scn": ["Catigurìa"],
    "sd": ["زمرو"],
    "se": ["Kategoriija"],
    "sg": ["Catégorie"],
    "sh": ["Kategorija", "Категорија"],
    "si": ["ප්‍රවර්ගය"],
    "sk": ["Kategória"],
    "sl": ["Kategorija"],
    "sq": ["Kategoria", "Kategori"],
    "sr": ["Kategorija", "Категорија"],
    "srn": ["Categorie", "Guru"],
    "stq": ["Kategorie"],
    "su": ["Kategori"],
    "sv": ["Kategori"],
    "sw": ["Jamii"],
    "szl": ["Kategoryjo", "Kategoria"],
    "ta": ["பகுப்பு"],
    "tcy": ["ವರ್ಗೊ"],
    "te": ["వర్గం"],
    "tet": ["Kategoría", "Kategoria"],
    "tg": ["Гурӯҳ"],
    "th": ["หมวดหมู่"],
    "ti": ["መደብ"],
    "tk": ["Kategoriýa"],
    "tl": ["Kategorya", "Kaurian"],
    "tpi": ["Grup"],
    "tr": ["Kategori", "KAT"],
    "tt": ["Төркем", "Törkem", "Категория"],
    "ty": ["Catégorie"],
    "tyv": ["Аңгылал", "Категория"],
    "udm": ["Категория"],
    "ug": ["تۈر"],
    "uk": ["Категория", "Категорія"],
    "ur": ["زمرہ"],
    "uz": ["Turkum", "Kategoriya"],
    "vec": ["Categoria"],
    "vep": ["Kategorii"],
    "vi": ["Thể_loại", "Thể loại"],
    "vls": ["Categorie"],
    "vo": ["Klad"],
    "wa": ["Categoreye"],
    "war": ["Kaarangay"],
    "wo": ["Wàll", "Catégorie"],
    "wuu": ["分类"],
    "xal": ["Янз", "Әәшл"],
    "xmf": ["კატეგორია"],
    "yi": ["קאטעגאריע", "קאַטעגאָריע"],
    "yo": ["Ẹ̀ka"],
    "za": ["分类"],
    "zea": ["Categorie"],
    "zh": ["分类", "分類", "CAT"],
    "zh-classical": ["分類", "CAT"],
    "zh-min-nan": ["分類", "Lūi-pia̍t"],
    "zh-yue": ["分类", "分類", "类", "類"],
}


def get_namespace(tag):
    """Get the namespace of tag.

    Parameters
    ----------
    tag : str
        Namespace or tag.

    Returns
    -------
    str
        Matched namespace or tag.

    """
    m = re.match("^{(.*?)}", tag)
    namespace = m.group(1) if m else ""
    if not namespace.startswith("http://www.mediawiki.org/xml/export-"):
        raise ValueError("%s not recognized as MediaWiki dump namespace" % namespace)
    return namespace


def parse_and_clean_wikicode(raw_content, parser, language):
    """Strips formatting and unwanted sections from raw page content."""
    wikicode = parser.parse(raw_content)

    # Filters for magic words that are parser instructions -- e.g., __NOTOC__
    re_rm_magic = re.compile("__[A-Z]*__", flags=re.UNICODE)

    # Filters for file/image links.
    media_prefixes = "|".join(["File", "Image", "Media"] + MEDIA_ALIASES.get(language, []))
    re_rm_wikilink = re.compile(f"^(?:{media_prefixes}):", flags=re.IGNORECASE | re.UNICODE)

    def rm_wikilink(obj):
        return bool(re_rm_wikilink.match(str(obj.title)))

    # Filters for references and tables
    def rm_tag(obj):
        return str(obj.tag) in {"ref", "table"}

    # Leave category links in-place but remove the category prefixes
    cat_prefixes = "|".join(["Category"] + CAT_ALIASES.get(language, []))
    re_clean_wikilink = re.compile(f"^(?:{cat_prefixes}):", flags=re.IGNORECASE | re.UNICODE)

    def is_category(obj):
        return bool(re_clean_wikilink.match(str(obj.title)))

    def clean_wikilink(obj):
        text = obj.__strip__()
        text = re.sub(re_clean_wikilink, "", text)
        obj.text = text

    def try_replace_obj(obj):
        try:
            clean_wikilink(obj)
        except ValueError:
            # For unknown reasons, objects are sometimes not found.
            pass

    def try_remove_obj(obj, section):
        try:
            section.remove(obj)
        except ValueError:
            # For unknown reasons, objects are sometimes not found.
            pass

    section_text = []
    # Filter individual sections to clean.
    for section in wikicode.get_sections(flat=True, include_lead=True, include_headings=True):
        for obj in section.ifilter_wikilinks(recursive=True):
            if rm_wikilink(obj):
                try_remove_obj(obj, section)
            elif is_category(obj):
                try_replace_obj(obj)
        for obj in section.ifilter_tags(matches=rm_tag, recursive=True):
            try_remove_obj(obj, section)

        section_text.append(re.sub(re_rm_magic, "", section.strip_code().strip()))
    return "\n\n".join(section_text)

src/lm_datasets/utils/wikimedia.py
malteos/lm-datasets
import os
import pytest
from dataclasses import dataclass
from lm_datasets.datasets.dataset_registry import get_registered_dataset_classes
from lm_datasets.datasets.base import BaseDataset

import pyarrow.parquet as pq

SHUFFLED_OUTPUT_DIR = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts_shuffled"


@pytest.mark.skipif(not os.path.exists(SHUFFLED_OUTPUT_DIR), reason="test file not exists")
def test_iterate_over_datasets():
    id_to_dataset_class = {cls.DATASET_ID: cls for cls in get_registered_dataset_classes()}
    dataset_cls = id_to_dataset_class["openlegaldata"]

    @dataclass
    class Arguments(object):
        output_dir = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts"
        shuffled_output_dir = SHUFFLED_OUTPUT_DIR
        output_format = "parquet"
        output_compression = "zst"

        skip_ratio = 0.01  # use these as vaidation set
        skip_n = 1000

        train_ratio = 0.05  # number of rows to use for tokenizer training

    args = Arguments()

    dataset: BaseDataset = dataset_cls(
        output_dir=args.output_dir,
        output_format=args.output_format,
        output_compression=args.output_compression,
        shuffled_output_dir=args.shuffled_output_dir,
    )

    batch_size = 1000
    # text_column_name = "text"

    def generate_texts():
        for fp in sorted(dataset.get_chunked_output_file_paths(shuffled=True)):
            pq_fp = fp

            with open(pq_fp, "rb") as file_handler:
                # pq_file = open_parquet_file_with_retries(dataset.get_output_file_path(), retries=2)
                pq_file = pq.ParquetFile(file_handler)

                print(f"rows {pq_file.metadata.num_rows}")

                rows = 0
                max_rows = int(pq_file.metadata.num_rows * args.train_ratio)
                skip_rows = min(int(pq_file.metadata.num_rows * args.skip_ratio), args.skip_n)

                print(f"reading {max_rows} rows")
                print(f"skipping {skip_rows} rows")

                for batch in pq_file.iter_batches(batch_size):
                    for text in batch[dataset.output_text_field]:
                        yield text

                        rows += 1

                print()
            break

    for text in generate_texts():
        print(text)

    print("x")

tests/test_iterate_over_shuffled_datasets.py
malteos/lm-datasets
from smart_open import open
import json
import pytest
import os

DATASET_FILE_PATH = "/netscratch/mostendorff/experiments/eulm/data/ele/de/openlegaldata/merged_v1_v2.jsonl.gz"


@pytest.mark.skipif(not os.path.exists(DATASET_FILE_PATH), reason="test file not exists")
def test_openlegaldata():
    docs_count = 1_658_611
    docs = []

    # after 250/1659 row groups
    # rows per group = 1000

    with open(DATASET_FILE_PATH) as f:
        for idx, line in enumerate(f):
            if idx > docs_count - 1_000:
                doc = json.loads(line)
                docs.append(doc)

    print("done")

tests/test_openlegaldata.py
malteos/lm-datasets
from pathlib import Path
import pyarrow.parquet as pq
import os
import pytest

from tqdm.auto import tqdm


@pytest.mark.skip(reason="requires dataset files")
def test_read_pq():
    fp = (  # macocu_el.parquet
        "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/opengptx_costep_de.parquet"  # noqa
    )
    fp = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/opengptx_dissertations_de.parquet"  # noqa
    fp = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/opengptx_dissertations_de.parquet"  # noqa
    fp = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts_shuffled/opengptx_bt_drucksachen.shuffled.parquet"  # noqa

    assert os.path.exists(fp)

    metadata = pq.read_metadata(fp)
    print(metadata)

    parquet_file = pq.ParquetFile(fp)

    rg = parquet_file.read_row_group(0)
    texts = rg["text"]

    print(str(texts[0])[:100])

    print(metadata)


@pytest.mark.skip(reason="requires dataset files")
def test_polars_read():
    fps = [
        "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts_shuffled/opengptx_bt_drucksachen.shuffled.parquet"  # noqa
    ]  # 6445
    fps = Path("/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/").glob("*.parquet")
    fps = list(sorted(fps))

    for i, fp in enumerate(tqdm(fps)):
        fp = str(fp)

        print(i, fp)

        if "__bak" in fp or "part-" in fp or i < 137:
            print("-- skip")
            continue

        fp = "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/pes2o.parquet"

    print("done")


@pytest.mark.skip(reason="requires dataset files")
def test_parquet_os_error():
    fps = ["/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/danewsroom.parquet"]

    fps = Path("/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts/").glob("*.parquet")
    fps = list(sorted(fps))

    for i, fp in enumerate(tqdm(fps)):
        fp = str(fp)

        print(i, fp)

        if "__bak" in fp or "part-" in fp:  # or i < 137:
            print("-- skip")
            continue

        # if ""
        print(fp)

        # import pyarrow as pa
        # t = pq.read_table(fp, memory_map=True)

        ok = 0
        errors = 0

        for i in range(4):
            try:
                # with open(fp, "rb") as f:
                #     pq.read_schema(f)

                with open(fp, "rb") as f:
                    parquet_file = pq.ParquetFile(f)
                    for _, record_batch in enumerate(parquet_file.iter_batches(batch_size=100)):
                        # pa_table = pa.Table.from_batches([record_batch])
                        for text in record_batch["text"]:
                            text = str(text)
                            break
                        break

                # tab = pq.read_table(fp, memory_map=True)
                # for batch in tab.to_batches():
                #     for text in batch["text"]:
                #         text = str(text)
                #         break
                #     break

                # from datasets import load_dataset
                # ds = load_dataset("parquet", data_files={"train": fp+"1"}, split="train", streaming=True)
                # next(iter(ds))

                # with  pq.ParquetFile(fp) as parquet_file:
                #     rg = parquet_file.read_row_group(0)  # columns=["text"]
                #     texts = rg["text"]
                #     text = str(texts[0])

                # parquet_file = pq.ParquetFile(fp)

                # texts = rg["text"]
                # text = str(texts[0])

                # df = pl.scan_parquet(fp)
                # text = df.row(0)

                # print(parquet_file.closed)

                # for batch in parquet_file.iter_batches(row_groups=[0]):
                #     text = str(batch["text"][0])
                #     break

                # if "x" in text:
                #     pass

                ok += 1
                # print(i, "OK")
            except OSError as e:
                errors += 1
                print(i, "ERROR", fp)
                raise e

        print("stats = ", ok, errors, fp)

        if errors:
            print("#### error")

        print("")

    print("done")


def test_read_groups():
    # TODO loop over files and read random groups
    pass


if __name__ == "__main__":
    # test_read_pq()
    test_parquet_os_error()
    # test_polars_read()

tests/test_read_parquet.py
malteos/lm-datasets
import os
import pytest

DATASET_FILE_PATH = (
    "/netscratch/mostendorff/experiments/eulm/data/lm_datasets_texts_shuffled/opengptx_bt_drucksachen.shuffled.parquet"
)


@pytest.mark.skipif(not os.path.exists(DATASET_FILE_PATH), reason="test file not exists")
def test_shuffle_parquet():
    import polars as pl

    fp = DATASET_FILE_PATH

    df = pl.scan_parquet(fp)

    df = df.with_columns(pl.all().shuffle(seed=1)).with_row_count()
    train_fraction: float = 0.75

    df_train = df.filter(pl.col("row_nr") < pl.col("row_nr").max() * train_fraction)
    df_test = df.filter(pl.col("row_nr") >= pl.col("row_nr").max() * train_fraction)

    df_train.collect().write_parquet("./shuffled_train.parquet")

    df_train.sink_parquet("./shuffled_train.parquet")
    df_test.write_parquet("./shuffled_train.parquet")

    df = df.sample(fraction=1, shuffle=True)

    print("done")


# df.write_parquet(shuffled_output_file_path)

tests/test_shuffle_parquet.py
malteos/lm-datasets
from textwrap import TextWrapper

wrapper = TextWrapper(width=100, max_lines=5)

text = "".join(["Foobar..."] * 1000)

wrapped_lines = wrapper.wrap(text)
wrapper_text = """\n<mark>&gt;wrap&lt;</mark> """.join(wrapped_lines)
print(wrapper_text)

print(text)

print("done")

tests/test_textwrap.py
malteos/lm-datasets
import itertools
import pyarrow as pa
import pyarrow.parquet as pq


def test_write_and_shuffle_parquet():
    def generate_texts(n):
        for i in range(n):
            yield f"Some text {i}"

    def get_batches(rows_iterable, chunk_size, schema):
        rows_it = iter(rows_iterable)
        while True:
            arr = pa.array(itertools.islice(rows_it, chunk_size))
            batch = pa.RecordBatch.from_arrays([arr], schema=schema)

            if not batch:
                break
            yield batch

    schema = pa.schema(
        [
            ("text", pa.string()),
        ]
    )

    print("start")

    batches = get_batches(generate_texts(1_00), chunk_size=10, schema=schema)

    print("writing ...")
    # Write the batches
    out_fp = "example.parquet"

    with pq.ParquetWriter(out_fp, schema=schema) as writer:
        for batch in batches:
            print("write batch")
            writer.write_batch(batch)

    print("done")

    # read again
    table = pq.read_table("example.parquet", columns=["text"], memory_map=True)

    assert str(table["text"][99]) == "Some text 99"

    from datasets.arrow_dataset import Dataset

    ds = Dataset(arrow_table=table)
    ds.shuffle(seed=1, keep_in_memory=False).to_parquet("example.shuffled.parquet", batch_size=10)

    shuffled_table = pq.read_table("example.shuffled.parquet", columns=["text"], memory_map=True)
    print(str(shuffled_table["text"][99]))

    assert str(shuffled_table["text"][99]) != "Some text 99"

    print("reading done.")


if __name__ == "__main__":
    test_write_and_shuffle_parquet()

tests/test_write_parquet.py
malteos/lm-datasets
import random
import string
import logging

import os
import pyarrow as pa
from lm_datasets.io.parquet import save_texts_to_parquet_chunks

from lm_datasets.utils.settings import LOGGING_KWARGS

from pathlib import Path
from random import randint
import tempfile

logging.basicConfig(**LOGGING_KWARGS)

logger = logging.getLogger(__name__)

# TMP_DIR = Path(__file__).parent.parent / "data/tmp"


def get_texts(n: int, min_len: int, max_len: int):
    for i in range(n):
        random_len = randint(min_len, max_len)
        random_str = "".join(random.choices(string.ascii_lowercase, k=random_len))

        yield random_str


def test_write_parquet_chunks():
    random.seed(0)

    with tempfile.TemporaryDirectory() as TMP_DIR:
        TMP_DIR = Path(TMP_DIR)

        print(f"Temp dir = {TMP_DIR}")

        def get_output_path(part=None, total_parts=None) -> str:
            if total_parts:
                return TMP_DIR / f"output.part-{part:04d}-of-{total_parts:04d}.parquet"
            elif part:
                return TMP_DIR / f"output.part-{part:04d}.parquet"
            else:
                return TMP_DIR / "output.parquet"

        if not os.path.exists(TMP_DIR):
            os.makedirs(TMP_DIR)

        # remove all tmp files
        for fp in TMP_DIR.glob("*.parquet"):
            os.remove(fp)
            print(f"Removed {fp}")

        output_text_field = "text"
        schema = pa.schema(
            [
                (output_text_field, pa.string()),
            ]
        )

        saved_docs, saved_chunks = save_texts_to_parquet_chunks(
            texts=get_texts(n=10_000, min_len=250, max_len=10_000),  # approx 50 MB  (none compression)
            schema=schema,
            max_chunk_uncompressed_bytes=10 * 1024 * 1024,  # 10 MB
            output_path_func=get_output_path,
            compression="ZSTD",
            batch_size=24,
        )

        assert saved_docs == 10_000
        assert saved_chunks == 5

        limited_saved_docs, limited_saved_chunks = save_texts_to_parquet_chunks(
            texts=get_texts(n=10_000, min_len=250, max_len=10_000),  # approx 50 MB  (none compression)
            schema=schema,
            max_chunk_uncompressed_bytes=10 * 1024 * 1024,  # 10 MB
            output_path_func=get_output_path,
            compression="ZSTD",
            batch_size=24,
            limit=1000,
        )

        assert limited_saved_docs == 1008  # more than limit due to batches
        assert limited_saved_chunks == 1

        single_file_saved_docs, single_file_saved_chunks = save_texts_to_parquet_chunks(
            texts=get_texts(n=1_000, min_len=250, max_len=10_000),  # approx 5 MB (none compression)
            schema=schema,
            max_chunk_uncompressed_bytes=0,
            output_path_func=get_output_path,
            compression="ZSTD",
            batch_size=24,
        )

        assert single_file_saved_docs == 1_000
        assert single_file_saved_chunks == 1

        print("done")


if __name__ == "__main__":
    test_write_parquet_chunks()

tests/test_write_parquet_chunks.py
malteos/lm-datasets