From 91ddb3672e514fa9824609ff047d7cab0c65631a Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Tue, 9 Apr 2024 00:14:00 +0200
Subject: Refactor

---
 rag/generator/__init__.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 rag/generator/__init__.py

(limited to 'rag/generator/__init__.py')

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)
-- 
cgit v1.2.3-70-g09d2