gfun_multimodal/main.py

import pickle
from argparse import ArgumentParser
from os.path import expanduser
from time import time

from dataManager.amazonDataset import AmazonDataset
from dataManager.multilingualDatset import MultilingualDataset
from dataManager.multiNewsDataset import MultiNewsDataset
from evaluation.evaluate import evaluate, log_eval
from gfun.generalizedFunnelling import GeneralizedFunnelling

""" 
TODO: 
    - add documentations sphinx
    - zero-shot setup
"""


def get_dataset(datasetname):
    assert datasetname in ["multinews", "amazon", "rcv1-2"], "dataset not supported"
    RCV_DATAPATH = expanduser(
        "~/datasets/rcv1-2/rcv1-2_doclist_trByLang1000_teByLang1000_processed_run0.pickle"
    )
    MULTINEWS_DATAPATH = expanduser("~/datasets/MultiNews/20110730/")
    if datasetname == "multinews":
        dataset = MultiNewsDataset(
            expanduser(MULTINEWS_DATAPATH),
            excluded_langs=["ar", "pe", "pl", "tr", "ua"],
        )
    elif datasetname == "amazon":
        dataset = AmazonDataset(
            domains=args.domains,
            nrows=args.nrows,
            min_count=args.min_count,
            max_labels=args.max_labels,
        )
    elif datasetname == "rcv1-2":
        dataset = (
            MultilingualDataset(dataset_name="rcv1-2")
            .load(RCV_DATAPATH)
            .reduce_data(langs=["en", "it", "fr"], maxn=100)
        )
    else:
        raise NotImplementedError
    return dataset


def main(args):
    dataset = get_dataset(args.dataset)
    if isinstance(dataset, MultilingualDataset) or isinstance(
        dataset, MultiNewsDataset
    ):
        lX, lY = dataset.training()
        # lX_te, lY_te = dataset.test()
        lX_te, lY_te = dataset.training()
    else:
        _lX = dataset.dX
        _lY = dataset.dY

    tinit = time()

    if args.load_trained is None:
        assert any(
            [
                args.posteriors,
                args.wce,
                args.multilingual,
                args.multilingual,
                args.transformer,
            ]
        ), "At least one of VGF must be True"

    gfun = GeneralizedFunnelling(
        posterior=args.posteriors,
        multilingual=args.multilingual,
        wce=args.wce,
        transformer=args.transformer,
        langs=dataset.langs(),
        embed_dir="~/resources/muse_embeddings",
        n_jobs=args.n_jobs,
        max_length=args.max_length,
        batch_size=args.batch_size,
        epochs=args.epochs,
        lr=args.lr,
        patience=args.patience,
        evaluate_step=args.evaluate_step,
        transformer_name=args.transformer_name,
        device="cuda",
        optimc=args.optimc,
        load_trained=args.load_trained,
    )

    # gfun.get_config()
    gfun.fit(lX, lY)

    if args.load_trained is None:
        print("[NB: FORCE-SKIPPING MODEL SAVE]")
        gfun.save()

    # if not args.load_model:
    #     gfun.save()

    preds = gfun.transform(lX)

    train_eval = evaluate(lY, preds)
    log_eval(train_eval, phase="train")

    timetr = time()
    print(f"- training completed in {timetr - tinit:.2f} seconds")

    test_eval = evaluate(lY_te, gfun.transform(lX_te))
    log_eval(test_eval, phase="test")

    timeval = time()
    print(f"- testing completed in {timeval - timetr:.2f} seconds")


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("-l", "--load_trained", type=str, default=None)
    # Dataset parameters -------------------
    parser.add_argument("-d", "--dataset", type=str, default="multinews")
    parser.add_argument("--domains", type=str, default="all")
    parser.add_argument("--nrows", type=int, default=10000)
    parser.add_argument("--min_count", type=int, default=10)
    parser.add_argument("--max_labels", type=int, default=50)
    # gFUN parameters ----------------------
    parser.add_argument("-p", "--posteriors", action="store_true")
    parser.add_argument("-m", "--multilingual", action="store_true")
    parser.add_argument("-w", "--wce", action="store_true")
    parser.add_argument("-t", "--transformer", action="store_true")
    parser.add_argument("--n_jobs", type=int, default=1)
    parser.add_argument("--optimc", action="store_true")
    # transformer parameters ---------------
    parser.add_argument("--transformer_name", type=str, default="mbert")
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--lr", type=float, default=1e-5)
    parser.add_argument("--max_length", type=int, default=512)
    parser.add_argument("--patience", type=int, default=5)
    parser.add_argument("--evaluate_step", type=int, default=10)

    args = parser.parse_args()

    main(args)