from dataloader.dsl_dataset import DatasetDSL

from llm_interface.chatgpt import ChatGPT

def load_llm(engine):
    split_point = engine.index("/")
    platform, engine_short = engine[:split_point], engine[split_point+1:]
    if platform == "openai":
        llm = ChatGPT(engine_short)
    else:
        raise NotImplementedError(f"platform {platform} not supported")
    return llm

def load_data(algorithm, dataset, context_len, query_scope):
    if dataset in ["BioEng", "Medical", "Ecology", "Genetics", "Synthesis"]:
        data = DatasetDSL(dataset, algorithm=algorithm, context_len=context_len, query_scope=query_scope)
    else:
        raise "wrong dataset!"

    return data.train_dataset, data.test_examples, data.avg_trainexample_len