summaryrefslogtreecommitdiff
path: root/rag/generator
diff options
context:
space:
mode:
Diffstat (limited to 'rag/generator')
-rw-r--r--rag/generator/__init__.py1
-rw-r--r--rag/generator/abstract.py4
2 files changed, 3 insertions, 2 deletions
diff --git a/rag/generator/__init__.py b/rag/generator/__init__.py
index 7da603c..541eff8 100644
--- a/rag/generator/__init__.py
+++ b/rag/generator/__init__.py
@@ -4,6 +4,7 @@ from .abstract import AbstractGenerator
from .ollama import Ollama
from .cohere import Cohere
+MODELS = ["ollama", "cohere"]
def get_generator(model: str) -> Type[AbstractGenerator]:
match model:
diff --git a/rag/generator/abstract.py b/rag/generator/abstract.py
index a53b5d8..5b336ea 100644
--- a/rag/generator/abstract.py
+++ b/rag/generator/abstract.py
@@ -1,11 +1,11 @@
-from abc import ABC, abstractmethod
+from abc import abstractmethod
from typing import Any, Generator
from .prompt import Prompt
-class AbstractGenerator(ABC, type):
+class AbstractGenerator(type):
_instances = {}
def __call__(cls, *args, **kwargs):