default_wiki_pipeline.py

# directly copied from https://huggingface.co/datasets/wikimedia/wikipedia/blob/script/wikipedia.py
"""Wikipedia dataset containing cleaned articles of all languages."""


import bz2
import csv
import io
import itertools
import json
import re
import xml.etree.ElementTree as etree
from urllib.parse import quote

import mwparserfromhell

import datasets


logger = datasets.logging.get_logger(__name__)


_HOMEPAGE = "https://dumps.wikimedia.org"

_CITATION = """\
@ONLINE {wikidump,
    author = {Wikimedia Foundation},
    title  = {Wikimedia Downloads},
    url    = {https://dumps.wikimedia.org}
}
"""

_DESCRIPTION = """\
Wikipedia dataset containing cleaned articles of all languages.
The datasets are built from the Wikipedia dump
(https://dumps.wikimedia.org/) with one split per language. Each example
contains the content of one full Wikipedia article with cleaning to strip
markdown and unwanted sections (references, etc.).
"""

_LICENSE = (
    "This work is licensed under the Creative Commons Attribution-ShareAlike "
    "3.0 Unported License. To view a copy of this license, visit "
    "http://creativecommons.org/licenses/by-sa/3.0/ or send a letter to "
    "Creative Commons, PO Box 1866, Mountain View, CA 94042, USA."
)

# Source: https://meta.wikimedia.org/wiki/List_of_Wikipedias
# Number: 326 = 339 - 13 (retrieved: 2023-11-17)
WIKIPEDIA_LANGUAGES = [
    "ab",
    "ace",
    "ady",
    "af",
    "als",
    "alt",
    "am",
    "ami",
    "an",
    "ang",
    "anp",
    "ar",
    "arc",
    "ary",
    "arz",
    "as",
    "ast",
    "atj",
    "av",
    "avk",
    "awa",
    "ay",
    "az",
    "azb",
    "ba",
    "ban",
    "bar",
    "bat-smg",
    "bbc",
    "bcl",
    "be",
    "be-tarask",
    "bg",
    "bh",
    "bi",
    "bjn",
    "blk",
    "bm",
    "bn",
    "bo",
    "bpy",
    "br",
    "bs",
    "bug",
    "bxr",
    "ca",
    "cbk-zam",
    "cdo",
    "ce",
    "ceb",
    "ch",
    "chr",
    "chy",
    "ckb",
    "co",
    "cr",
    "crh",
    "cs",
    "csb",
    "cu",
    "cv",
    "cy",
    "da",
    "dag",
    "de",
    "dga",
    "din",
    "diq",
    "dsb",
    "dty",
    "dv",
    "dz",
    "ee",
    "el",
    "eml",
    "en",
    "eo",
    "es",
    "et",
    "eu",
    "ext",
    "fa",
    "fat",
    "ff",
    "fi",
    "fiu-vro",
    "fj",
    "fo",
    "fon",
    "fr",
    "frp",
    "frr",
    "fur",
    "fy",
    "ga",
    "gag",
    "gan",
    "gcr",
    "gd",
    "gl",
    "glk",
    "gn",
    "gom",
    "gor",
    "got",
    "gpe",
    "gu",
    "guc",
    "gur",
    "guw",
    "gv",
    "ha",
    "hak",
    "haw",
    "he",
    "hi",
    "hif",
    "hr",
    "hsb",
    "ht",
    "hu",
    "hy",
    "hyw",
    "ia",
    "id",
    "ie",
    "ig",
    "ik",
    "ilo",
    "inh",
    "io",
    "is",
    "it",
    "iu",
    "ja",
    "jam",
    "jbo",
    "jv",
    "ka",
    "kaa",
    "kab",
    "kbd",
    "kbp",
    "kcg",
    "kg",
    "ki",
    "kk",
    "kl",
    "km",
    "kn",
    "ko",
    "koi",
    "krc",
    "ks",
    "ksh",
    "ku",
    "kv",
    "kw",
    "ky",
    "la",
    "lad",
    "lb",
    "lbe",
    "lez",
    "lfn",
    "lg",
    "li",
    "lij",
    "lld",
    "lmo",
    "ln",
    "lo",
    "lt",
    "ltg",
    "lv",
    "mad",
    "mai",
    "map-bms",
    "mdf",
    "mg",
    "mhr",
    "mi",
    "min",
    "mk",
    "ml",
    "mn",
    "mni",
    "mnw",
    "mr",
    "mrj",
    "ms",
    "mt",
    "mwl",
    "my",
    "myv",
    "mzn",
    "nah",
    "nap",
    "nds",
    "nds-nl",
    "ne",
    "new",
    "nia",
    "nl",
    "nn",
    "no",
    "nov",
    "nqo",
    "nrm",
    "nso",
    "nv",
    "ny",
    "oc",
    "olo",
    "om",
    "or",
    "os",
    "pa",
    "pag",
    "pam",
    "pap",
    "pcd",
    "pcm",
    "pdc",
    "pfl",
    "pi",
    "pih",
    "pl",
    "pms",
    "pnb",
    "pnt",
    "ps",
    "pt",
    "pwn",
    "qu",
    "rm",
    "rmy",
    "rn",
    "ro",
    "roa-rup",
    "roa-tara",
    "ru",
    "rue",
    "rw",
    "sa",
    "sah",
    "sat",
    "sc",
    "scn",
    "sco",
    "sd",
    "se",
    "sg",
    "sh",
    "shi",
    "shn",
    "si",
    "simple",
    "sk",
    "skr",
    "sl",
    "sm",
    "smn",
    "sn",
    "so",
    "sq",
    "sr",
    "srn",
    "ss",
    "st",
    "stq",
    "su",
    "sv",
    "sw",
    "szl",
    "szy",
    "ta",
    "tay",
    "tcy",
    "te",
    "tet",
    "tg",
    "th",
    "ti",
    "tk",
    "tl",
    "tly",
    "tn",
    "to",
    "tpi",
    "tr",
    "trv",
    "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",
    "zgh",
    "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:
# https://ab.wikipedia.org/w/api.php?action=query&meta=siteinfo&siprop=namespacealiases|namespaces&format=json&formatversion=2
# Retrieved: 2023-11-17
MEDIA_ALIASES = {
    "ab": ["Амедиа", "Афаил", "Изображение", "Медиа", "Файл"],
    "ace": ["Alat", "Berkas", "Beureukaih", "Gambar"],
    "ady": ["Медиа"],
    "af": ["Beeld", "Lêer"],
    "als": ["Bild", "Datei", "Medium"],
    "alt": ["Изображение", "Медиа", "Файл"],
    "am": ["ስዕል", "ፋይል"],
    "ami": ["Faylo", "Mitiya", "图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "档案", "檔案"],
    "an": ["Imachen", "Imagen"],
    "ang": ["Biliþ", "Ymele"],
    "anp": ["फाईल", "मीडिया"],
    "ar": ["صورة", "ملف", "ميديا", "وسائط"],
    "arc": ["ܠܦܦܐ", "ܡܝܕܝܐ"],
    "ary": ["صورة", "فيشي", "ملف", "ميديا", "وسائط"],
    "arz": ["صورة", "ملف", "ميديا", "وسائط"],
    "as": ["चित्र", "চিত্র", "চিত্ৰ", "মাধ্যম"],
    "ast": ["Archivu", "Ficheru", "Imagen", "Imaxe", "Imaxen", "Medios"],
    "atj": ["Natisinahikaniwoc", "Tipatcimoctakewin"],
    "av": ["Изображение", "Медиа", "Файл"],
    "avk": ["Ewava", "Imagen", "Iyeltak", "Kanaca", "Mamind", "Изображение"],
    "awa": ["फाइल", "मीडिया"],
    "ay": ["Archivo", "Imagen", "Medio"],
    "az": ["Fayl", "Mediya", "Şəkil"],
    "azb": ["تصویر", "رسانه", "رسانه‌ای", "فایل", "مدیا"],
    "ba": ["Изображение", "Медиа", "Рәсем", "Файл"],
    "ban": ["Berkas", "Gambar", "Média"],
    "bar": ["Bild", "Datei", "Medium"],
    "bat-smg": ["Abruozdielis", "Medėjė", "Vaizdas"],
    "bbc": ["Ugasan"],
    "bcl": ["Ladawan", "Medio"],
    "be": ["Выява", "Мультымедыя", "Файл"],
    "be-tarask": ["Выява", "Мэдыя", "Файл"],
    "bg": ["Картинка", "Медия", "Файл"],
    "bh": ["चित्र", "मीडिया"],
    "bjn": ["Barakas", "Berkas", "Gambar"],
    "blk": ["ဖုဲင်", "မီဒီယာ"],
    "bm": ["Fichier", "Média"],
    "bn": ["চিত্র", "মিডিয়া"],
    "bpy": ["ছবি", "মিডিয়া"],
    "br": ["Restr", "Skeudenn"],
    "bs": ["Datoteka", "Medija", "Mediji", "Slika"],
    "bug": ["Berkas", "Gambar"],
    "bxr": ["Изображение", "Меди", "Файл"],
    "ca": ["Fitxer", "Imatge"],
    "cbk-zam": ["Archivo", "Imagen", "Medio"],
    "cdo": ["图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "档案", "檔案"],
    "ce": ["Изображение", "Медиа", "Медйа", "Сурт", "Файл", "Хlум"],
    "ceb": ["Imahen", "Medya", "Payl"],
    "ch": ["Litratu"],
    "ckb": ["میدیا", "پەڕگە"],
    "co": ["Immagine"],
    "crh": ["Fayl", "Resim", "Медиа", "Ресим", "Файл"],
    "cs": ["Média", "Obrázok", "Soubor"],
    "csb": ["Grafika", "Òbrôzk"],
    "cu": ["Ви́дъ", "Видъ", "Дѣло", "Срѣдьства"],
    "cv": ["Изображение", "Медиа", "Ӳкерчĕк"],
    "cy": ["Delwedd"],
    "da": ["Billede", "Fil"],
    "dag": ["Lahabali kɔligu", "Miidiya"],
    "de": ["Bild", "Datei", "Medium"],
    "dga": ["Duoro bimbu zie", "Duoro kɔre"],
    "din": ["Apamduööt", "Ciɛl"],
    "diq": ["Dosya", "Medya"],
    "dsb": ["Bild", "Dataja", "Medija", "Wobraz"],
    "dty": ["चित्र", "मिडिया"],
    "dv": ["މީޑިއާ", "ފައިލު", "ފައިލް"],
    "el": ["Αρχείο", "Εικόνα", "Μέσο", "Μέσον"],
    "eml": ["Immagine"],
    "eo": ["Aŭdvidaĵo", "Dosiero"],
    "es": ["Archivo", "Imagen", "Medio"],
    "et": ["Fail", "Meedia", "Pilt"],
    "eu": ["Fitxategi", "Irudi"],
    "ext": ["Archivu", "Imagen", "Mediu"],
    "fa": ["تصویر", "رسانه", "رسانه‌ای", "مدیا", "پرونده"],
    "fat": ["Fael"],
    "ff": ["Fichier", "Média"],
    "fi": ["Kuva", "Tiedosto"],
    "fiu-vro": ["Meediä", "Pilt"],
    "fo": ["Miðil", "Mynd"],
    "fon": ["Wékpo", "Yɛwliɖonuji"],
    "fr": ["Fichier", "Média"],
    "frp": ["Fichiér", "Mèdia", "Émâge"],
    "frr": ["Bild", "Datei", "Medium"],
    "fur": ["Figure", "Immagine"],
    "fy": ["Ofbyld"],
    "ga": ["Meán", "Íomhá"],
    "gag": ["Dosya", "Dosye", "Mediya", "Medya", "Resim"],
    "gan": ["图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "文檔", "档案", "檔案"],
    "gcr": ["Fiché", "Médja"],
    "gd": ["Faidhle", "Meadhan"],
    "gl": ["Arquivo", "Ficheiro", "Imagem", "Imaxe"],
    "glk": ["تصویر", "رسانه", "رسانه‌ای", "فاىل", "مديا", "پرونده"],
    "gn": ["Imagen", "Medio", "Ta'ãnga"],
    "gom": ["फायल", "माध्यम", "मिडिया"],
    "gor": ["Berkas", "Gambar"],
    "got": ["𐍆𐌴𐌹𐌻𐌰"],
    "gu": ["ચિત્ર", "દ્રશ્ય-શ્રાવ્ય (મિડિયા)"],
    "guc": ["Anaajaalaa", "Ayaakuwapülee", "Imagen"],
    "gur": ["Faali", "Miidiya"],
    "guw": ["Wepo"],
    "gv": ["Coadan", "Meanyn"],
    "hak": ["图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "档案", "檔案"],
    "haw": ["Kiʻi", "Pāpaho", "Waihona"],
    "he": ["מדיה", "קו", "קובץ", "תמונה"],
    "hi": ["चित्र", "मीडिया"],
    "hif": ["file", "saadhan"],
    "hr": ["DT", "Datoteka", "Mediji", "Slika"],
    "hsb": ["Bild", "Dataja", "Wobraz"],
    "ht": ["Fichye", "Imaj", "Medya"],
    "hu": ["Fájl", "Kép", "Média"],
    "hy": ["Մեդիա", "Պատկեր"],
    "hyw": ["Մեդիա", "Պատկեր"],
    "ia": ["Imagine", "Multimedia"],
    "id": ["Berkas", "Gambar"],
    "ig": ["Midia", "Nká", "Usòrò", "Ákwúkwó orünotu"],
    "ilo": ["Midia", "Papeles"],
    "inh": ["Изображение", "Медиа", "Файл"],
    "io": ["Arkivo", "Imajo"],
    "is": ["Miðill", "Mynd"],
    "it": ["Immagine"],
    "ja": ["ファイル", "メディア", "画像"],
    "jbo": ["datnyvei", "velsku"],
    "jv": ["Barkas", "Gambar", "Medhia", "Médhia"],
    "ka": ["მედია", "სურათი", "ფაილი"],
    "kaa": ["Fayl", "Su'wret", "Swret", "Taspa", "Сурет", "Таспа", "تاسپا", "سۋرەت"],
    "kab": ["Tugna"],
    "kbd": ["Медиа", "Файл"],
    "kbp": ["Fichier", "Média"],
    "kcg": ["Fail"],
    "kg": ["Fisye"],
    "kk": ["Swret", "Taspa", "Сурет", "Таспа", "تاسپا", "سۋرەت"],
    "kl": ["Billede", "Fiileq", "Fil"],
    "km": ["មីឌា", "មេឌា", "រូបភាព", "ឯកសារ"],
    "kn": ["ಚಿತ್ರ", "ಮೀಡಿಯ"],
    "ko": ["그림", "미디어", "파일"],
    "koi": ["Изображение", "Медиа", "Файл"],
    "krc": ["Изображение", "Медиа", "Файл"],
    "ks": ["فَیِل", "میڈیا"],
    "ksh": ["Beld", "Belld", "Bild", "Datei", "Medie", "Medium", "Meedije", "Meedijum"],
    "ku": ["Medya", "Wêne", "میدیا", "پەڕگە"],
    "kv": ["Изображение", "Медиа", "Файл"],
    "kw": ["Restren"],
    "ky": ["Медиа", "Файл"],
    "la": ["Fasciculus", "Imago"],
    "lad": ["Archivo", "Dossia", "Dosya", "Imagen", "Meddia", "Medya"],
    "lb": ["Bild", "Fichier"],
    "lbe": ["Изображение", "Медиа", "Сурат"],
    "lez": ["Mediya", "Şəkil", "Изображение", "Медиа", "Файл"],
    "lfn": ["Fix"],
    "li": ["Aafbeilding", "Afbeelding", "Plaetje"],
    "lij": ["Immaggine", "Immagine"],
    "lld": ["Immagine"],
    "lmo": ["Archivi", "Immagine", "Imàjine"],
    "ln": ["Fichier", "Média"],
    "lo": ["ສື່", "ສື່ອ", "ຮູບ"],
    "lt": ["Medija", "Vaizdas"],
    "ltg": ["Fails", "Medeja"],
    "lv": ["Attēls"],
    "mad": ["Bhengkek", "Gambar", "Mèḍia"],
    "mai": ["फाइल", "मेडिया"],
    "map-bms": ["Barkas", "Gambar", "Medhia", "Médhia"],
    "mdf": ["Изображение", "Медиа", "Няйф"],
    "mg": ["Média", "Rakitra", "Sary"],
    "mhr": ["Изображение", "Медиа", "Файл"],
    "min": ["Berkas", "Gambar"],
    "mk": ["Медиум", "Медија", "Податотека", "Слика"],
    "ml": ["ചി", "ചിത്രം", "പ്ര", "പ്രമാണം", "മീഡിയ"],
    "mn": ["Зураг", "Медиа", "Файл"],
    "mni": ["ꯃꯦꯗꯤꯌꯥ", "ꯐꯥꯏꯜ"],
    "mnw": ["မဳဒဳယာ", "ဝှာင်"],
    "mr": ["चित्र", "मिडिया"],
    "mrj": ["Изображение", "Медиа", "Файл"],
    "ms": ["Fail", "Imej"],
    "mt": ["Medja", "Midja", "Stampa"],
    "mwl": ["Arquivo", "Fexeiro", "Ficheiro", "Imagem", "Multimédia"],
    "my": ["ဖိုင်", "မီဒီယာ"],
    "myv": ["Артовкс", "Изображение", "Медия"],
    "mzn": ["تصویر", "رسانه", "رسانه‌ای", "مدیا", "مه‌دیا", "پرونده"],
    "nah": ["Imagen", "Mēdiatl", "Īxiptli"],
    "nap": ["Fiùra", "Immagine"],
    "nds": ["Bild", "Datei"],
    "nds-nl": ["Afbeelding", "Bestaand", "Ofbeelding"],
    "ne": ["चित्र", "मीडिया"],
    "new": ["किपा", "माध्यम"],
    "nia": ["Berkas", "Gambar"],
    "nl": ["Afbeelding", "Bestand"],
    "nn": ["Bilde", "Fil", "Filpeikar"],
    "no": ["Bilde", "Fil", "Medium"],
    "nqo": ["ߞߐߕߐ߮", "ߟߊߛߋߢߊߥߙߍ"],
    "nrm": ["Fichier", "Média"],
    "nso": ["Seswantšho"],
    "nv": ["Eʼelyaaígíí"],
    "oc": ["Fichièr", "Imatge", "Mèdia"],
    "olo": ["Failu", "Kuva", "Medii"],
    "or": ["ଫାଇଲ", "ମାଧ୍ୟମ"],
    "os": ["Изображение", "Медиа", "Ныв", "Файл"],
    "pa": ["ਤਸਵੀਰ", "ਮੀਡੀਆ"],
    "pcd": ["Fichier", "Média"],
    "pdc": ["Bild", "Datei", "Feil", "Medium"],
    "pfl": ["Bild", "Dadai", "Datei", "Medium"],
    "pi": ["पटिमा", "मीडिया"],
    "pl": ["Grafika", "Plik"],
    "pms": ["Figura", "Immagine"],
    "pnb": ["تصویر", "فائل", "میڈیا"],
    "pnt": ["Αρχείον", "Εικόνα", "Εικόναν", "Μέσον"],
    "ps": ["انځور", "دوتنه", "رسنۍ"],
    "pt": ["Arquivo", "Ficheiro", "Imagem", "Multimédia"],
    "pwn": [
        "mitiya",
        "sineqetj a vecik",
        "图像",
        "图片",
        "圖像",
        "圖片",
        "媒体",
        "媒体文件",
        "媒体档案",
        "媒體",
        "媒體文件",
        "媒體檔案",
        "文件",
        "档案",
        "檔案",
    ],
    "qu": ["Imagen", "Midya", "Rikcha"],
    "rm": ["Bild", "Datoteca", "Multimedia"],
    "rmy": ["Chitro", "Fişier", "Imagine", "Mediya"],
    "rn": ["Dosiye"],
    "ro": ["Fişier", "Fișier", "Imagine"],
    "roa-rup": ["Fişier", "Fișier", "Imagine"],
    "roa-tara": ["Immagine"],
    "ru": ["Изображение", "Медиа", "Файл"],
    "rue": ["Зображення", "Изображение", "Медиа", "Медіа", "Файл"],
    "rw": ["Dosiye", "Itangazamakuru"],
    "sa": ["चित्रं", "चित्रम्", "माध्यम", "माध्यमम्", "सञ्चिका"],
    "sah": ["Билэ", "Изображение", "Миэдьийэ", "Ойуу"],
    "sat": ["ᱢᱤᱰᱤᱭᱟ", "ᱨᱮᱫ"],
    "sc": ["Immàgini"],
    "scn": ["Immagine", "Mmàggini", "Mèdia"],
    "sd": ["ذريعات", "عڪس", "فائل"],
    "se": ["Bilde", "Fiila", "Kuva"],
    "sg": ["Fichier", "Média"],
    "sh": ["DT", "Datoteka", "Fotografija", "Medija", "Mediji", "Slika", "ДТ", "Датотека", "Медиј"],
    "shi": ["Afaylu", "Midya"],
    "shn": ["သိုဝ်ႇၶၢဝ်ႇ", "ၾၢႆႇ"],
    "si": ["ගොනුව", "මාධ්‍යය", "රූපය"],
    "sk": ["Médiá", "Obrázok", "Súbor"],
    "skr": ["فائل", "میڈیا"],
    "sl": ["Datoteka", "Slika"],
    "smn": ["Kuva", "Tiätuvuárkká"],
    "sq": ["Figura", "Skeda"],
    "sr": [
        "Datoteka",
        "Fotografija",
        "Medij",
        "Medija",
        "Slika",
        "Датотека",
        "Медиј",
        "Медија",
        "Слика",
        "Фотографија",
    ],
    "srn": ["Afbeelding", "Gefre"],
    "stq": ["Bielde", "Bild"],
    "su": ["Gambar", "Média"],
    "sv": ["Bild", "Fil"],
    "sw": ["Faili", "Picha"],
    "szl": ["Grafika", "Plik"],
    "szy": ["myiti", "tangan", "图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "档案", "檔案"],
    "ta": ["ஊடகம்", "படிமம்"],
    "tay": [
        "biru' na zayzyuwaw",
        "biru’ na zayzyuwaw",
        "media",
        "图像",
        "图片",
        "圖像",
        "圖片",
        "媒体",
        "媒体文件",
        "媒体档案",
        "媒體",
        "媒體文件",
        "媒體檔案",
        "文件",
        "档案",
        "檔案",
    ],
    "tcy": ["ಫೈಲ್", "ಮಾದ್ಯಮೊ"],
    "te": ["దస్త్రం", "ఫైలు", "బొమ్మ", "మీడియా"],
    "tet": ["Arquivo", "Imagem", "Imajen"],
    "tg": ["Акс", "Медиа"],
    "th": ["ภาพ", "สื่อ", "ไฟล์"],
    "ti": ["ሜድያ", "ፋይል"],
    "tk": ["Faýl"],
    "tl": ["Midya", "Talaksan"],
    "tly": ["Fajl", "Medja"],
    "tn": ["Pego", "Setshwantsho"],
    "tpi": ["Fail"],
    "tr": ["Dosya", "Medya", "Ortam", "Resim"],
    "trv": [
        "Meyti",
        "Patas bntasan",
        "图像",
        "图片",
        "圖像",
        "圖片",
        "媒体",
        "媒体文件",
        "媒体档案",
        "媒體",
        "媒體文件",
        "媒體檔案",
        "文件",
        "档案",
        "檔案",
    ],
    "tt": ["Räsem", "Изображение", "Медиа", "Рәсем", "Файл"],
    "ty": ["Fichier", "Média"],
    "tyv": ["Изображение", "Медиа", "Файл"],
    "udm": ["Изображение", "Медиа", "Суред", "Файл"],
    "ug": ["ھۆججەت", "ۋاسىتە"],
    "uk": ["Зображення", "Изображение", "Медиа", "Медіа", "Файл"],
    "ur": ["تصویر", "زریعہ", "فائل", "ملف", "میڈیا", "وسیط"],
    "uz": ["Fayl", "Mediya", "Tasvir"],
    "vec": ["Immagine", "Imàjine", "Mèdia"],
    "vep": ["Fail", "Pilt"],
    "vi": ["Hình", "Phương tiện", "Tập tin"],
    "vls": ["Afbeelding", "Ofbeeldienge"],
    "vo": ["Magod", "Nünamakanäd", "Ragiv"],
    "wa": ["Imådje"],
    "war": ["Fayl", "Medya", "Paypay"],
    "wo": ["Dencukaay", "Xibaarukaay"],
    "wuu": ["图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "档案", "檔案"],
    "xal": ["Аһар", "Боомг", "Зург", "Изображение"],
    "xmf": ["მედია", "სურათი", "ფაილი"],
    "yi": ["בילד", "טעקע", "מעדיע", "תמונה"],
    "yo": ["Amóhùnmáwòrán", "Fáìlì", "Àwòrán"],
    "za": ["图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "档案", "檔案"],
    "zea": ["Afbeelding", "Plaetje"],
    "zgh": ["ⴰⴼⴰⵢⵍⵓ", "ⵎⵉⴷⵢⴰ"],
    "zh": ["图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "档案", "檔案"],
    "zh-classical": ["图像", "图片", "圖像", "圖片", "媒体", "媒体文件", "媒体档案", "媒體", "媒體文件", "媒體檔案", "文件", "档案", "檔案"],
    "zh-min-nan": [
        "Mûi-thé",
        "tóng-àn",
        "图像",
        "图片",
        "圖像",
        "圖片",
        "媒体",
        "媒体文件",
        "媒体档案",
        "媒體",
        "媒體文件",
        "媒體檔案",
        "文件",
        "档案",
        "檔案",
    ],
    "zh-yue": ["图", "图像", "圖", "圖像", "媒体", "媒體", "文件", "档", "档案", "檔", "檔案"],
}

# Source: for each Wikipedia language code (example shown for "ab"), aliases for namespace 14 accessed via this API call:
# https://ab.wikipedia.org/w/api.php?action=query&meta=siteinfo&siprop=namespacealiases|namespaces&format=json&formatversion=2
# Retrieved: 2023-11-17
CAT_ALIASES = {
    "ab": ["Акатегориа", "Категория"],
    "ace": ["Kategori", "Kawan"],
    "af": ["Kategorie"],
    "als": ["Kategorie"],
    "alt": ["Категория"],
    "am": ["መደብ"],
    "ami": ["Kasasiwasiw", "分类", "分類"],
    "an": ["Categoría"],
    "ang": ["Flocc"],
    "anp": ["श्रेणी"],
    "ar": ["تصنيف"],
    "arc": ["ܣܕܪܐ"],
    "ary": ["تصنيف"],
    "arz": ["تصنيف"],
    "as": ["CAT", "श्रेणी", "শ্রেণী", "শ্ৰেণী"],
    "ast": ["Categoría"],
    "atj": ["Tipanictawin"],
    "av": ["Категория"],
    "avk": ["Loma"],
    "awa": ["श्रेणी"],
    "ay": ["Categoría"],
    "az": ["Kateqoriya"],
    "azb": ["بؤلمه"],
    "ba": ["Категория", "Төркөм"],
    "ban": ["Kategori"],
    "bar": ["Kategorie"],
    "bat-smg": ["Kategorija", "Kateguorėjė"],
    "bbc": ["Horong"],
    "bcl": ["Kategorya"],
    "be": ["Катэгорыя"],
    "be-tarask": ["Катэгорыя"],
    "bg": ["Категория"],
    "bh": ["श्रेणी"],
    "bjn": ["Kategori", "Tumbung"],
    "blk": ["ကဏ္ဍ"],
    "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"],
    "dag": ["Pubu"],
    "de": ["Kategorie"],
    "dga": ["Gbuli"],
    "din": ["Bekätakthook"],
    "diq": ["Kategori", "Kategoriye"],
    "dsb": ["Kategorija"],
    "dty": ["श्रेणी"],
    "dv": ["ޤިސްމު"],
    "el": ["Κατηγορία"],
    "eml": ["Categoria"],
    "eo": ["Kategorio"],
    "es": ["CAT", "Categoría"],
    "et": ["Kategooria"],
    "eu": ["Kategoria"],
    "ext": ["Categoria", "Categoría"],
    "fa": ["رده"],
    "fat": ["Nkyekyεmu"],
    "ff": ["Catégorie"],
    "fi": ["Luokka"],
    "fiu-vro": ["Katõgooria"],
    "fo": ["Bólkur"],
    "fon": ["Akpaxwé"],
    "fr": ["Catégorie"],
    "frp": ["Catègorie"],
    "frr": ["Kategorie"],
    "fur": ["Categorie"],
    "fy": ["Kategory"],
    "ga": ["Catagóir", "Rang"],
    "gag": ["Kategori", "Kategoriya"],
    "gan": ["分类", "分類"],
    "gcr": ["Katégori"],
    "gd": ["Roinn-seòrsa"],
    "gl": ["Categoría"],
    "glk": ["جرگه", "رده"],
    "gn": ["Ñemohenda"],
    "gom": ["वर्ग", "श्रेणी"],
    "gor": ["Dalala"],
    "got": ["𐌷𐌰𐌽𐍃𐌰"],
    "gu": ["CAT", "શ્રે", "શ્રેણી"],
    "guc": ["Akotchajülee sünülia"],
    "gur": ["Buuri buuri"],
    "guw": ["Adà"],
    "gv": ["Ronney"],
    "hak": ["分类", "分類"],
    "haw": ["Māhele"],
    "he": ["קט", "קטגוריה"],
    "hi": ["श्र", "श्रेणी"],
    "hif": ["vibhag"],
    "hr": ["CT", "KT", "Kategorija"],
    "hsb": ["Kategorija"],
    "ht": ["Kategori"],
    "hu": ["Kategória"],
    "hy": ["Կատեգորիա"],
    "hyw": ["Ստորոգութիւն"],
    "ia": ["Categoria"],
    "id": ["Kategori"],
    "ie": ["Categorie"],
    "ig": ["Ébéonọr", "Òtù"],
    "ilo": ["Kategoria"],
    "inh": ["ОагӀат"],
    "io": ["Kategorio"],
    "is": ["Flokkur"],
    "it": ["CAT", "Categoria"],
    "ja": ["カテゴリ"],
    "jbo": ["klesi"],
    "jv": ["Kategori"],
    "ka": ["კატეგორია"],
    "kaa": ["Kategoriya", "Sanat", "Санат", "سانات"],
    "kab": ["Taggayt"],
    "kbd": ["Категориэ", "Категория"],
    "kbp": ["Catégorie"],
    "kcg": ["Sa"],
    "kg": ["Kalasi"],
    "kk": ["Sanat", "Санат", "سانات"],
    "kl": ["Kategori", "Sumut atassuseq"],
    "km": ["ចំណាត់ក្រុម", "ចំណាត់ថ្នាក់ក្រុម", "ចំនាត់ថ្នាក់ក្រុម"],
    "kn": ["ವರ್ಗ"],
    "ko": ["분류"],
    "koi": ["Категория"],
    "krc": ["Категория"],
    "ks": ["زٲژ"],
    "ksh": ["Kategorie", "Katejori", "Kattejori", "Saachjrop", "Saachjropp", "Saachjrupp", "Sachjrop"],
    "ku": ["Kategorî", "پۆل"],
    "kv": ["Категория"],
    "kw": ["Class", "Klass"],
    "ky": ["Категория"],
    "la": ["Categoria"],
    "lad": ["Categoría", "Kateggoría", "Katēggoría"],
    "lb": ["Kategorie"],
    "lbe": ["Категория"],
    "lez": ["Категория"],
    "lfn": ["Categoria"],
    "li": ["Categorie", "Kategorie"],
    "lij": ["Categoria", "Categorîa"],
    "lld": ["Categoria"],
    "lmo": ["Categoria", "Categuria"],
    "ln": ["Catégorie"],
    "lo": ["ໝວດ"],
    "lt": ["Kategorija"],
    "ltg": ["Kategoreja"],
    "lv": ["Kategorija"],
    "mad": ["Bhângsa"],
    "mai": ["CA", "श्रेणी"],
    "map-bms": ["Kategori"],
    "mdf": ["Категорие", "Категория"],
    "mg": ["Catégorie", "Sokajy"],
    "mhr": ["Категорий", "Категория"],
    "min": ["Kategori"],
    "mk": ["Категорија"],
    "ml": ["വ", "വി", "വിഭാഗം", "വർഗ്ഗം"],
    "mn": ["Ангилал"],
    "mni": ["ꯃꯆꯥꯈꯥꯏꯕ"],
    "mnw": ["ကဏ္ဍ"],
    "mr": ["वर्ग"],
    "mrj": ["Категори", "Категория"],
    "ms": ["Kategori"],
    "mt": ["Kategorija"],
    "mwl": ["Catadorie", "Categoria"],
    "my": ["ကဏ္ဍ"],
    "myv": ["Категория"],
    "mzn": ["رج", "رده"],
    "nah": ["Categoría", "Neneuhcāyōtl"],
    "nap": ["Categoria", "Categurìa"],
    "nds": ["Kategorie"],
    "nds-nl": ["Categorie", "Kategorie", "Kattegerie"],
    "ne": ["श्रेणी"],
    "new": ["पुचः"],
    "nia": ["Kategori"],
    "nl": ["Categorie"],
    "nn": ["Kategori"],
    "no": ["Kategori"],
    "nqo": ["ߦߌߟߡߊ"],
    "nrm": ["Catégorie"],
    "nso": ["Setensele"],
    "nv": ["Tʼááłáhági átʼéego"],
    "oc": ["Categoria"],
    "olo": ["Kategourii"],
    "or": ["ବିଭାଗ", "ଶ୍ରେଣୀ"],
    "os": ["Категори"],
    "pa": ["ਸ਼੍ਰੇਣੀ"],
    "pcd": ["Catégorie"],
    "pcm": ["Katigori"],
    "pdc": ["Abdeeling", "Kategorie"],
    "pfl": ["Kadegorie", "Kategorie", "Sachgrubb"],
    "pi": ["विभाग"],
    "pl": ["Kategoria"],
    "pms": ["Categorìa"],
    "pnb": ["گٹھ"],
    "pnt": ["Κατηγορίαν"],
    "ps": ["وېشنيزه"],
    "pt": ["Categoria"],
    "pwn": ["pinapapilipiliqan", "分类", "分類"],
    "qu": ["Katiguriya"],
    "rm": ["Categoria"],
    "rmy": ["Shopni"],
    "rn": ["Umuce"],
    "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": ["KAT", "KT", "Kategorija", "КАТ", "КТ", "Категорија"],
    "shi": ["Taggayt"],
    "shn": ["ပိူင်ထၢၼ်ႈ"],
    "si": ["ප්‍රවර්ගය"],
    "sk": ["Kategória"],
    "skr": ["ونکی"],
    "sl": ["Kategorija"],
    "smn": ["Luokka"],
    "sq": ["Kategori", "Kategoria"],
    "sr": ["Kategorija", "Категорија"],
    "srn": ["Categorie", "Guru"],
    "stq": ["Kategorie"],
    "su": ["Kategori"],
    "sv": ["Kategori"],
    "sw": ["Jamii"],
    "szl": ["Kategoria", "Kategoryjo"],
    "szy": ["kakuniza", "分类", "分類"],
    "ta": ["பகுப்பு"],
    "tay": ["zyuwaw na", "分类", "分類"],
    "tcy": ["ವರ್ಗೊ"],
    "te": ["వర్గం"],
    "tet": ["Kategoria", "Kategoría"],
    "tg": ["Гурӯҳ"],
    "th": ["หมวดหมู่"],
    "ti": ["መደብ"],
    "tk": ["Kategoriýa"],
    "tl": ["Kategorya", "Kaurian"],
    "tly": ["Tispir"],
    "tn": ["Karolo"],
    "tpi": ["Grup"],
    "tr": ["KAT", "Kategori"],
    "trv": ["Snakun", "分类", "分類"],
    "tt": ["Törkem", "Категория", "Төркем"],
    "tw": ["Nkyekyεmu"],
    "ty": ["Catégorie"],
    "tyv": ["Аңгылал", "Категория"],
    "udm": ["Категория"],
    "ug": ["تۈر"],
    "uk": ["Категория", "Категорія"],
    "ur": ["زمرہ"],
    "uz": ["Kategoriya", "Turkum"],
    "vec": ["Categoria"],
    "vep": ["Kategorii"],
    "vi": ["Thể loại"],
    "vls": ["Categorie"],
    "vo": ["Klad"],
    "wa": ["Categoreye"],
    "war": ["Kaarangay"],
    "wo": ["Catégorie", "Wàll"],
    "wuu": ["分类", "分類"],
    "xal": ["Янз", "Әәшл"],
    "xmf": ["კატეგორია"],
    "yi": ["קאַטעגאָריע", "קאטעגאריע"],
    "yo": ["Ẹ̀ka"],
    "za": ["分类", "分類"],
    "zea": ["Categorie"],
    "zgh": ["ⴰⵙⵎⵉⵍ"],
    "zh": ["CAT", "分类", "分類"],
    "zh-classical": ["CAT", "分类", "分類"],
    "zh-min-nan": ["Lūi-pia̍t", "分类", "分類"],
    "zh-yue": ["分类", "分類", "类", "類"],
}

# Note that Wikimedia servers have rate limited downloaders and they are capping the number of per-ip connections to 2
# Their mirror sites do not have this cap: https://dumps.wikimedia.org/mirrors.html
_HOST = "https://dumps.wikimedia.org"
_URL_PATH_SEGMENT = "/{lang}wiki/{date}/"
_INFO_FILE = "dumpstatus.json"


_VERSION = datasets.Version("4.0.0", "")


class WikipediaConfig(datasets.BuilderConfig):
    """BuilderConfig for Wikipedia."""

    def __init__(self, language=None, date=None, host=_HOST, version=_VERSION, **kwargs):
        """BuilderConfig for Wikipedia.
        Args:
            language (str): Language code for the Wikipedia dump to use.
            date (str): Date of the Wikipedia dump in YYYYMMDD format. A list of
                available dates can be found at https://dumps.wikimedia.org/enwiki/.
            host (str, defaults to 'https://dumps.wikimedia.org'): URL of the server that hosts the Wikipedia dump.
                It defaults to the official Wikimedia host, but this has rate limited downloaders and is capping the
                number of per-IP connections to 2.
                To parallelize data download, use a mirror: https://dumps.wikimedia.org/mirrors.html
            **kwargs: Keyword arguments forwarded to super.
        """
        super().__init__(
            name=f"{date}.{language}",
            description=f"Wikipedia dataset for {language}, parsed from {date} dump.",
            version=version,
            **kwargs,
        )
        self.date = date
        self.language = language
        self.host = host.rstrip("/")


class Wikipedia(datasets.GeneratorBasedBuilder):
    """Wikipedia dataset."""

    # Use mirror (your.org) to avoid download caps.
    BUILDER_CONFIG_CLASS = WikipediaConfig

    def _info(self):
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=datasets.Features(
                {
                    "id": datasets.Value("string"),
                    "url": datasets.Value("string"),
                    "title": datasets.Value("string"),
                    "text": datasets.Value("string"),
                }
            ),
            homepage=_HOMEPAGE,
            citation=_CITATION,
        )

    def _get_base_url(self):
        return self.config.host + _URL_PATH_SEGMENT.format(
            lang=self.config.language.replace("-", "_"), date=self.config.date
        )

    def _split_generators(self, dl_manager):
        # Download dump status info file
        info_url = self._get_base_url() + _INFO_FILE
        info_path = dl_manager.download_and_extract(info_url)
        with open(info_path, encoding="utf-8") as f:
            dump_info = json.load(f)
        multistream_dump_info = dump_info["jobs"]["articlesmultistreamdump"]
        if multistream_dump_info["status"] != "done":
            raise FileNotFoundError(
                f"Specified dump ({self._get_base_url()}) multistream status is not 'done':"
                f" {multistream_dump_info['status']}"
            )

        # Download index and multistream XML data files
        total_bytes = 0
        index_urls = []
        xml_urls = []
        for fname, info in multistream_dump_info["files"].items():
            if ".txt" in fname:
                index_urls.append(self._get_base_url() + fname)
            elif ".xml" in fname:
                total_bytes += info["size"]
                xml_urls.append(self._get_base_url() + fname)
        data_urls = [*zip(sorted(index_urls), sorted(xml_urls))]  # Parallelize data downloading
        data_paths = dl_manager.download(data_urls)

        # Parallelize over concatenated multiple compressed streams (with 100 pages each)
        filepaths, starts, ends = [], [], []
        for index_path, xml_path in data_paths:
            index = _extract_index(index_path)
            for start, end in _pairwise(index):
                filepaths.append(xml_path)
                starts.append(start)
                ends.append(end)

        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                gen_kwargs={
                    "filepaths": filepaths,
                    "starts": starts,
                    "ends": ends,
                },
            )
        ]

    def _generate_examples(self, filepaths, starts, ends):
        # No shuffle anymore
        for filepath, start, end in zip(filepaths, starts, ends):
            for id_, title, raw_content in _extract_content(filepath, start, end):
                yield from _clean_content((id_, title, raw_content), self.config.language)


def _extract_index(filepath):
    with open(filepath, "rb") as compressed_file:
        binary_file = bz2.BZ2File(filename=compressed_file)
        text_file = io.TextIOWrapper(binary_file, encoding="utf-8", newline="")
        return sorted({int(row[0]) for row in csv.reader(text_file, delimiter=":")})


def _pairwise(index):
    starts, ends = itertools.tee(index)
    _ = next(ends, None)
    return itertools.zip_longest(starts, ends)


def _extract_content(filepath, start, end):
    """Extract articles from a single stream of a multistream WikiMedia XML file."""
    logger.info("generating examples from = %s", filepath)
    with open(filepath, "rb") as compressed_file:
        compressed_file.seek(start)
        compressed_data = compressed_file.read(end - start if end else -1)
        binary_data = bz2.BZ2Decompressor().decompress(compressed_data)
        # Enclose within a single root node to avoid ParseError: junk after document element
        binary_data = b"<mediawiki>" + binary_data + b"</mediawiki>"
        with io.StringIO(binary_data.decode(encoding="utf-8")) as text_stream:
            for _, elem in etree.iterparse(text_stream):
                if not elem.tag.endswith("page"):
                    continue
                namespace = elem.tag[:-4]
                ns = elem.find(f"./{namespace}ns").text
                redirect = elem.find(f"./{namespace}redirect")
                # Filter pages that are not in the "main" namespace or that are redirects
                if ns != "0" or redirect is not None:
                    elem.clear()
                    continue
                id_ = elem.find(f"./{namespace}id").text
                title = elem.find(f"./{namespace}title").text
                raw_content = elem.find(f"./{namespace}revision/{namespace}text").text
                elem.clear()
                # Filter empty pages
                if raw_content is None:
                    continue
                yield id_, title, raw_content


def _clean_content(inputs, language):
    """Clean raw wikicode to extract text."""
    id_, title, raw_content = inputs
    try:
        text = _parse_and_clean_wikicode(raw_content, parser=mwparserfromhell, language=language)
    except mwparserfromhell.parser.ParserError as e:
        logger.error("mwparserfromhell ParseError: %s", e)
        return
    if not text:
        return
    url = _construct_url(title, language)
    yield id_, {"id": id_, "url": url, "title": title, "text": text}


def _parse_and_clean_wikicode(raw_content, parser, language):
    """Strip 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)


def _construct_url(title, language):
    # See: https://meta.wikimedia.org/wiki/Help:URL
    return f"https://{language}.wikipedia.org/wiki/{quote(title)}"