summaryrefslogtreecommitdiff
path: root/rag/generator/__init__.py
blob: 770db16e4efa89f76753d7b30ea38846f1da7e84 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from typing import Type

from rag.generator.abstract import AbstractGenerator
from rag.generator.cohere import Cohere
from rag.generator.ollama import Ollama

MODELS = ["local", "cohere"]

def get_generator(model: str) -> Type[AbstractGenerator]:
    match model:
        case "local":
            return Ollama()
        case "cohere":
            return Cohere()
        case _:
            exit(1)