llm/llm_factory.py

from typing import Dict

from .huggingfacehub.hf_model import HF_Mistaril, HF_TinyLlama
from .llamacpp.lc_model import LC_TinyLlama, LC_Phi3


class LLM_Factory:

    # trigger = {"model_type": "execution_type"} -> {"hf": "small"}
    @staticmethod
    def create_llm(prompt_entity: str, prompt_id: int, trigger: Dict[str, str]):

        print(trigger)

        for key, value in trigger.items():
            if key == "hf" and value == "effective":
                model = HF_Mistaril(prompt_entity=prompt_entity, prompt_id=prompt_id)
            elif key == "hf" and value == "small":
                model = HF_TinyLlama(prompt_entity=prompt_entity, prompt_id=prompt_id)
            elif key == "lc" and value == "effective":
                model = LC_Phi3(prompt_entity=prompt_entity, prompt_id=prompt_id)
            elif key == "lc" and value == "small":
                model = LC_TinyLlama(prompt_entity=prompt_entity, prompt_id=prompt_id)
            else:
                model = None

        return model