Files
multifit/prepare_xnli.py
2018-11-19 06:50:50 +00:00

74 lines
2.2 KiB
Python

import zipfile
from pathlib import Path
from typing import Optional, Union
import fire
from tqdm import tqdm
from fastai.core import *
from fastai.datasets import *
ROOT = Path("data").resolve()
XNLI_DIR = ROOT / "xnli"
if not ROOT.exists():
ROOT.mkdir()
XNLI_DIR.mkdir(exist_ok=True)
print(f"Saving data in {ROOT}")
MT_FILE = "XNLI-MT-1.0.zip"
XNLI_FILE = "XNLI-1.0.zip"
MT_PATH = XNLI_DIR / MT_FILE
XNLI_PATH = XNLI_DIR / XNLI_FILE
MT_URL = "https://s3.amazonaws.com/xnli/XNLI-MT-1.0.zip"
XNLI_URL = "https://s3.amazonaws.com/xnli/XNLI-1.0.zip"
class TqdmUpTo(tqdm):
def update_to(self, b=1, bsize=1, tsize=None):
if tsize is not None:
self.total = tsize
self.update(b * bsize - self.n)
def download_data(url: str, fname: Union[str, Path], dest: Optional[Union[str, Path]]):
"""
Download data if the filename does not exist already
Uses Tqdm to show download progress
"""
from urllib.request import urlretrieve
filepath = (Path(dest) / fname).resolve()
if not filepath.exists():
dirname = Path(filepath.parents[0])
print(f"Creating directory {dirname} from {filepath}")
dirname.mkdir(exist_ok=True)
with TqdmUpTo(unit="B", unit_scale=True, miniters=1, desc=url.split("/")[-1]) as t:
urlretrieve(url, filepath, reporthook=t.update_to)
return str(filepath.resolve().absolute())
def get_and_unzip_data(url: str, fname: Union[str, Path] = None, dest: Union[str, Path] = None):
"""Download `url` if it doesn't exist to `fname` and un-tgz to folder `dest`"""
if dest is None:
dest = url.split("/")[-1]
dest = Path(dest)
fname = dest / fname
if not fname.exists():
download_data(url=url, fname=fname, dest=dest)
print(f"Extracting {fname.resolve().absolute()} \n to {dest}")
zipfile.ZipFile(fname, "r").extractall(dest)
return dest
def get_xnli_and_MT(dest: Union[str, Path] = XNLI_DIR):
get_and_unzip_data(url=XNLI_URL, fname=XNLI_FILE, dest=dest)
get_and_unzip_data(url=MT_URL, fname=MT_FILE, dest=dest)
if __name__ == "__main__":
fire.Fire(get_xnli_and_MT)