backend/inference.py

from backend.section_infer_helper.local_llm_helper import local_llm_helper
from backend.section_infer_helper.online_llm_helper import online_llm_helper
from backend.cwe_infer_helper import cwe_infer_helper

MODEL_TYPE_HELPER_MAP = {
    "Local LLM": local_llm_helper,
    "Online LLM": online_llm_helper
}

PREDEF_MODEL_MAP = {
    "Local LLM": local_llm_helper.PREDEF_MODEL,
    "Online LLM": online_llm_helper.PREDEF_MODEL
}

LOCAL_MODEL_PEFT_MAP = local_llm_helper.MODEL_PEFT_MAP

PREDEF_CWE_MODEL = cwe_infer_helper.PREDEF_MODEL

def section_infer(diff_code, patch_message, model_type, *model_config):
    helper = MODEL_TYPE_HELPER_MAP.get(model_type, None)
    if helper is None:
        raise ValueError(f"Model {model_type} is not supported")
    helper.load_model(*model_config)
    results = helper.infer(diff_code, patch_message)
    return results


def cwe_infer(diff_code, patch_message, *model_config):
    helper = cwe_infer_helper
    helper.load_model(*model_config)
    results = helper.infer(diff_code, patch_message)
    return results