From 3f447bff69c20109474c455f1ad52bd547ab49e9 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Tue, 9 Apr 2024 00:41:55 +0200 Subject: Update --- rag/generator/__init__.py | 1 + rag/generator/abstract.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) (limited to 'rag/generator') 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): -- cgit v1.2.3-70-g09d2