from completion_algorithm.llm_rag_algorithm import LLM_RAG

def load_algorithm(algorithm, dataset, train_dataset, engine, embedding_model, context_num, compiler, chunk_size):
    if algorithm == "Ours" or algorithm == "RAG-LLM" or algorithm == "Pure-LLM":
        RAG = None
        if algorithm == "Ours":
            RAG = "DSL"
        elif algorithm == "RAG-LLM":
            RAG = "NL"
        return LLM_RAG(dataset, train_dataset, engine, RAG, embedding_model, context_num, compiler, chunk_size)
    else:
        raise "wrong algorithm"