from fairseq.models.roberta import RobertaModel

import sys

PATH_TO_COSMIC = "../Model/COSMIC"
EXTRACTORS_PATH = PATH_TO_COSMIC + "/feature_extraction"
EPIK_MODEL_DIR = PATH_TO_COSMIC + "/erc_training"

sys.path.append(PATH_TO_COSMIC)
sys.path.append(EXTRACTORS_PATH)
sys.path.append(EPIK_MODEL_DIR)
from Model.COSMIC.feature_extraction.comet.csk_feature_extract import (
    CSKFeatureExtractor,
)

from Model.COSMIC.erc_training.predict_epik import parse_cosmic_args, load_model

roberta = RobertaModel.from_pretrained(
    EXTRACTORS_PATH + "/checkpoints/epik/",
    checkpoint_file="checkpoint_best.pt",
    data_name_or_path="../../epik-bin",
)
roberta.eval()

comet = CSKFeatureExtractor(dir=EXTRACTORS_PATH)

cosmic_args = parse_cosmic_args()

COSMIC_MODEL = load_model(EPIK_MODEL_DIR + "/epik/best_model.pt", cosmic_args)