clip-spanish / prepare_wit.py
edugp's picture
Add all necessary files to replicate training run
2daf3c7
raw
history blame
No virus
4.27 kB
import argparse
import json
import logging
import os
import time
from typing import List
import urllib.request
import urllib.error
import pandas as pd
from tqdm import tqdm
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
level=logging.INFO,
)
logger = logging.getLogger(__name__)
def split_and_save_datasets(lines: List[str], output_dir: str, train_proportion: float, valid_proportion: float):
total_lines = len(lines)
train_lines = lines[:int(total_lines * train_proportion)]
valid_lines = lines[int(total_lines * train_proportion):int(total_lines * (train_proportion + valid_proportion))]
test_lines = lines[int(total_lines * (train_proportion + valid_proportion)):]
with open(f"{output_dir}/train_dataset.json", "w") as f:
f.write("\n".join(train_lines))
with open(f"{output_dir}/valid_dataset.json", "w") as f:
f.write("\n".join(valid_lines))
with open(f"{output_dir}/test_dataset.json", "w") as f:
f.write("\n".join(test_lines))
def prepare_wit(
tsv: str, language: str, output_dir: str, seed: int, train_proportion: float, valid_proportion: float, backup_period: int, language_col: str="language", caption_col: str="caption_reference_description", url_col: str="image_url", pause=0.875, retries: int=10):
os.makedirs(output_dir, exist_ok=True)
logger.info("Loading dataset")
df = pd.read_csv(tsv, sep="\t", engine="python")
existing_files = set(os.listdir(output_dir))
not_exists_condition = (~(df[url_col].map(lambda x: x.split("/")[-1][-100:]).isin(existing_files)))
df = df[(df["language"] == language) & (~df["caption_reference_description"].isnull()) & not_exists_condition]
# Shuffle
df = df.sample(frac=1.0, random_state=seed)
logger.info(f"Trying to downloading {df.shape[0]} files")
lines = []
count = 0
try:
with tqdm(total=len(df)) as pbar:
for i, row in tqdm(df.iterrows()):
url = row[url_col]
caption = row[caption_col]
# Trim image file names so that they are no longer than 100 characters
image_filename = url.split("/")[-1][-100:]
image_path = f"{output_dir}/{image_filename}"
for retry in range(retries):
try:
# Download file
urllib.request.urlretrieve(url, image_path)
lines.append(json.dumps({"image_path": image_path, "captions": [caption]}, ensure_ascii=False))
count += 1
break
except urllib.error.HTTPError as e:
time.sleep(pause * 10)
if count % backup_period == 0:
logger.info(f"Saving dataset backup: Number of lines {len(lines)}")
split_and_save_datasets(lines, output_dir, train_proportion, valid_proportion)
if retry == retries - 1:
logger.info(f"Skipping {image_filename}")
pbar.update(1)
# Save existing dataset, even upon failure
finally:
split_and_save_datasets(lines, output_dir, train_proportion, valid_proportion)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description = "Download and prepare the WIT dataset")
parser.add_argument("--tsv", type=str, default=f"/home/{os.environ['USER']}/data/wit/wit_v1.train.all-1percent_sample.tsv")
parser.add_argument("--language", type=str, default="es")
parser.add_argument("--output_dir", type=str, default=f"/home/{os.environ['USER']}/data/wit/prepared_dataset")
parser.add_argument("--random_seed", type=int, default=0)
parser.add_argument("--train_proportion", type=float, default=0.8)
parser.add_argument("--valid_proportion", type=float, default=0.1)
parser.add_argument("--backup_period", type=int, default=1000)
args = parser.parse_args()
assert args.train_proportion + args.valid_proportion < 1.0, "The sum of train_proportion and valid_proportion has to be < 1.0"
prepare_wit(args.tsv, args.language, args.output_dir, args.random_seed, args.train_proportion, args.valid_proportion, args.backup_period)