File size: 1,006 Bytes
d5436e0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 |
from typing import Dict, List
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]):
if trigger.keys() == "hf" and trigger.values() == "effective":
model = HF_Mistaril(prompt_entity=prompt_entity, prompt_id=prompt_id)
elif trigger.keys() == "hf" and trigger.values() == "small":
model = HF_TinyLlama(prompt_entity=prompt_entity, prompt_id=prompt_id)
elif trigger.keys() == "lc" and trigger.values() == "effective":
model = LC_Phi3(prompt_entity=prompt_entity, prompt_id=prompt_id)
elif trigger.keys() == "lc" and trigger.values() == "small":
model = LC_TinyLlama(prompt_entity=prompt_entity, prompt_id=prompt_id)
else:
model = None
return model
|