diff options
Diffstat (limited to 'rag/generator/__init__.py')
-rw-r--r-- | rag/generator/__init__.py | 15 |
1 files changed, 15 insertions, 0 deletions
diff --git a/rag/generator/__init__.py b/rag/generator/__init__.py new file mode 100644 index 0000000..7da603c --- /dev/null +++ b/rag/generator/__init__.py @@ -0,0 +1,15 @@ +from typing import Type + +from .abstract import AbstractGenerator +from .ollama import Ollama +from .cohere import Cohere + + +def get_generator(model: str) -> Type[AbstractGenerator]: + match model: + case "ollama": + return Ollama() + case "cohere": + return Cohere() + case _: + exit(1) |