import bz2
import codecs
import os
import re
import subprocess
import urllib.parse as up
import xml.etree.cElementTree as etree
import mwparserfromhell
from datasets import Dataset
from data_juicer.download.downloader import (
DocumentDownloader,
DocumentExtractor,
DocumentIterator,
download_and_extract,
get_wikipedia_urls,
)
from data_juicer.utils.file_utils import expand_outdir_and_mkdir
# The majority of this code is taken from the HuggingFace
# implementation of the Wikipedia dataset preparation:
# https://github.com/huggingface/datasets/blob/7e30308f49f8c85dc7a2ab5aafbff04b5d2f38e2/datasets/wikipedia/wikipedia.py
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": [
"檔",
"档",
"文件",
"图",
"媒體",
"圖",
"档案",
"图像",
"圖像",
"媒体",
"檔案",
],
}
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": ["分类", "分類", "类", "類"],
}
[docs]
class WikipediaDownloader(DocumentDownloader):
[docs]
def __init__(self, download_dir, verbose=False):
super().__init__()
self._download_dir = download_dir
self._verbose = verbose
[docs]
def download(self, url):
urlpath = up.urlparse(url).path[1:]
output_name = urlpath.replace("/", "-")
output_file = os.path.join(self._download_dir, output_name)
if os.path.exists(output_file):
print(f"bz2 file: {output_file} exists. Not downloading")
else:
print(f"Downloading {url} and writing to {output_file}")
# Download with either wget or s5cmd (aws)
cmd = ["wget", url, "-O", output_file]
if self._verbose:
stdout, stderr = None, None
else:
stdout, stderr = subprocess.DEVNULL, subprocess.DEVNULL
p = subprocess.run(
cmd,
stdout=stdout,
stderr=stderr,
)
if p.returncode != 0:
print(f"Failed to download {url} to {output_file}")
return output_file
[docs]
class WikipediaIterator(DocumentIterator):
[docs]
def __init__(self, language="en", log_frequency=1000):
super().__init__()
self._language = language
self._log_frequency = log_frequency
self._counter = 0
[docs]
def iterate(self, file_path):
self._counter = 0
bname = os.path.split(file_path)[-1]
input_file = bz2.BZ2File(filename=file_path)
utf_f = codecs.getreader("utf-8")(input_file)
context = etree.iterparse(utf_f, events=("end",))
for i, (unused_event, elem) in enumerate(context):
if not elem.tag.endswith("page"):
continue
if self._counter > 0 and self._counter % self._log_frequency == 0:
print(f"Extracted {self._counter} articles from {file_path}")
self._counter += 1
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")
url = f"https://{self._language}.wikipedia.org/wiki/{up.quote(title)}" # noqa: E501
# 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()
# Filter redirects.
if raw_content is None or red_ is not None:
continue
yield {
"title": title,
"id": id_,
"url": url,
"language": self._language,
"source_id": f"{bname}",
}, raw_content
[docs]
def download_wikipedia(
output_path: str,
language: str = "en",
dump_date=None,
output_type: str = "jsonl",
raw_download_dir=None,
keep_raw_download=False,
force_download=False,
url_limit=None,
item_limit=None,
) -> Dataset:
"""
Downloads the latest Wikipedia dumps and extracts them using mwparserfromhell
Args:
output_path: The path to the root directory of the files
language: The language of the Wikipedia articles to download
dump_date: A string formatted as "YYYYMMDD" for the wikipedia dump to use.
If None, latest dump is used.
output_type: The file type to save the data as.
raw_download_dir: Path to store the raw download files for intermediate processing.
If None, they are stored in a folder named "downloads" under output_path.
keep_raw_download: If True, keeps the bz2 files that have not been extracted.
force_download: If False, will skip processing all files in output_paths that already exist and
directly read from them instead.
url_limit: The maximum number of raw files to download from the snapshot. If None, all
files from the range of snapshots are downloaded.
""" # noqa: E501
wikipedia_urls = get_wikipedia_urls(language=language, dump_date=dump_date)
if url_limit:
wikipedia_urls = wikipedia_urls[:url_limit]
output_paths = list(
map(
lambda url: os.path.join(output_path, url.split("/")[-1] + f".{output_type}"),
wikipedia_urls,
)
)
if not raw_download_dir:
raw_download_dir = os.path.join(output_path, "downloads")
expand_outdir_and_mkdir(raw_download_dir)
downloader = WikipediaDownloader(download_dir=raw_download_dir)
iterator = WikipediaIterator(language=language)
extractor = WikipediaExtractor(language=language)
output_format = {
"text": str,
"title": str,
"id": str,
"url": str,
"language": str,
"source_id": str,
"filename": str,
}
dataset = download_and_extract(
wikipedia_urls,
output_paths,
downloader,
iterator,
extractor,
output_format,
output_type=output_type,
keep_raw_download=keep_raw_download,
force_download=force_download,
item_limit=item_limit,
)
return dataset