local env = import "../env.jsonnet";
local base = import "ace.jsonnet";

local fn_path = "data/framenet/full/full.jsonl";
local mapping_path = "data/ace/framenet2ace/";

local debug = false;

# training
local lr = env.json("PT_LR", "5e-5");
local cuda_devices = base.cuda_devices;

# mapping
local min_weight = env.json("MIN_WEIGHT", '0.0');
local max_weight = env.json("MAX_WEIGHT", '5.0');

{
    dataset_reader: {
        type: "semantic_role_labeling",
        debug: debug,
        pretrained_model: base.dataset_reader.pretrained_model,
        ignore_label: false,
        [ if debug then "max_instances" ]: 128,
        event_smoothing_factor: base.dataset_reader.event_smoothing_factor,
        arg_smoothing_factor: base.dataset_reader.arg_smoothing_factor,
        ontology_mapping_path: mapping_path + '/ontology_mapping.json',
        min_weight: min_weight,
        max_weight: max_weight,
    },
    validation_dataset_reader: base.dataset_reader,
    train_data_path: fn_path,
    validation_data_path: base.validation_data_path,
    test_data_path: base.test_data_path,
    vocabulary: {
        type: "extend",
        directory: mapping_path + "/vocabulary"
    },

    datasets_for_vocab_creation: ["train"],

    data_loader: base.data_loader,
    validation_data_loader: base.validation_data_loader,

    model: base.model,

    trainer: {
        num_epochs: base.trainer.num_epochs,
        patience: base.trainer.patience,
        [if std.length(cuda_devices) == 1 then "cuda_device"]: cuda_devices[0],
        validation_metric: "+arg-c_f",
        num_gradient_accumulation_steps: base.trainer.num_gradient_accumulation_steps,
        optimizer: {
            type: "transformer",
            base: {
                type: "adam",
                lr: lr,
            },
            embeddings_lr: 0.0,
            encoder_lr: 1e-5,
            pooler_lr: 1e-5,
            layer_fix: base.trainer.optimizer.layer_fix,
        }
    },

    [if std.length(cuda_devices) > 1 then "distributed"]: {
        "cuda_devices": cuda_devices
    },
    [if std.length(cuda_devices) == 1 then "evaluate_on_test"]: true
}