script.py

import os
import re
from pathlib import Path

from src.utils import load_csv


def gen_vocab_files(exp_folder: Path):
    constants = set()
    preds = set()
    data = load_csv(exp_folder / "csvs" / "all.csv")
    print(len(data))
    for entry in data:
        print(entry)
        preds.add(entry[0])
        constants.add(entry[1])
        constants.add(entry[2])

    constants = list(sorted(constants))
    preds = list(sorted(preds))
    print("PREDS", preds)

    (exp_folder / "vocab").mkdir(exist_ok=True)
    with open(exp_folder / "vocab" / "constants.txt", "w+") as f:
        f.write("\n".join(constants))

    with open(exp_folder / "vocab" / "pred.txt", "w+") as f:
        f.write("\n".join(preds))

    with open(exp_folder / "program.dlog", "r") as f:
        program = f.read()
    special = [x for x in re.findall('"(.*?)"', program) if x[0] == "s"]
    with open(exp_folder / "vocab" / "special.txt", "w+") as f:
        f.write("\n".join(special))


def preprocess_data(exp_folder: Path | str):
    exp_folder = Path(exp_folder)
    file_names = ["kg.csv", "train.csv", "test.csv", "val.csv"]
    file_names = [exp_folder / "csvs" / f for f in file_names]
    for f in file_names:
        if f.exists():
            to_proper_csv(f)

    gen_vocab_files(exp_folder)


def prolog_to_csv(line):
    assert ")" in line
    return line.replace(".", "").replace(")", "").replace("(", ",")


def kg_to_csv(line):
    line = line.split("\t")
    assert len(line) == 3
    line = line[1], line[0], line[2]
    return ",".join(line)


def to_proper_csv(file_name):
    with open(file_name, "r") as f:
        data = f.readlines()
    data = [line.strip() for line in data]
    data = [line for line in data if line]
    if any("(" in line for line in data):
        data = [prolog_to_csv(line) for line in data]
    if any("\t" in line for line in data):
        data = [kg_to_csv(line) for line in data]
    with open(file_name, "w") as f:
        f.write("\n".join(data))


def sort_file(filename):
    with open(filename, "r") as f:
        lines = f.read().split("\n")
    lines = [line.strip() for line in lines if line.strip()]
    lines = list(sorted(lines))
    with open(filename, "w") as f:
        f.write("\n".join(lines))


if __name__ == "__main__":
    kg = "umls2"
    to_proper_csv(f"data/{kg}/csvs/val.csv")
    to_proper_csv(f"data/{kg}/csvs/test.csv")
    to_proper_csv(f"data/{kg}/csvs/train.csv")
    to_proper_csv(f"data/{kg}/csvs/kg.csv")
    to_proper_csv(f"data/{kg}/csvs/all.csv")
    gen_vocab_files(Path(f"data/{kg}"))