File size: 1,019 Bytes
b32ff0a d5436e0 385f711 48edb58 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 26 27 28 |
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
|