run_kg.py

import argparse
from pathlib import Path

import torch
import wandb

from src.embeddings import EmbeddingStore
from src.utils import set_seed, load_config, ground_program
from src.dataset import get_dataset
from src.train import train_epoch, test


def get_engine_class(name):
    if name == "glog":
        from src.glog_engine import GlogEngine

        return GlogEngine


def main(experiment_folder: str, eval_split: str):
    config = load_config(experiment_folder)
    name = str(experiment_folder)[5:]
    wandb.init(project="eq_reasoning", config=config, name=name)
    config = wandb.config
    set_seed(config["seed"])

    emb = EmbeddingStore(
        experiment_folder,
        emb_scale=config["emb_scale"],
        special_init=config["special_init"],
    )
    optimizer = torch.optim.AdamW(
        params=emb.parameters(),
        lr=config["embedding_lr"],
        weight_decay=config["weight_decay"],
    )

    engine = get_engine_class("glog")(
        experiment_folder, iter_limit=config.get("iter_limit")
    )

    for epoch in range(config["nb_epochs"]):
        print(f"## EPOCH {epoch+1} ##")
        train_dataset = get_dataset(experiment_folder, "train")
        train_epoch(
            engine,
            train_dataset,
            emb,
            optimizer,
            nb_samples=config["nb_samples"],
            grad_clip=config["grad_clip"],
        )

        map_prob, map_program = emb.MAP()
        wandb.log({"train/map_prob": map_prob.item()})
        ground_program(experiment_folder, map_program)

    test_dataset = get_dataset(experiment_folder, eval_split)
    test(engine, test_dataset, emb, ranks=True)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_folder")
    parser.add_argument("-e", "--eval_split", type=str, default="val")
    args = parser.parse_args()

    main(experiment_folder=Path(args.experiment_folder), eval_split=args.eval_split)