From 9e0cbcb4e7f1f3f95f304046d3190c6ebc4d3901 Mon Sep 17 00:00:00 2001
From: Gustaf Rydholm <gustaf.rydholm@gmail.com>
Date: Wed, 24 Apr 2024 09:09:24 +0200
Subject: Reformat and fix typo

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

(limited to 'rag/retriever/rerank/__init__.py')

diff --git a/rag/retriever/rerank/__init__.py b/rag/retriever/rerank/__init__.py
new file mode 100644
index 0000000..16b2fac
--- /dev/null
+++ b/rag/retriever/rerank/__init__.py
@@ -0,0 +1,15 @@
+from typing import Type
+
+from rag.retriever.rerank.abstract import AbstractReranker
+from rag.retriever.rerank.cohere import CohereReranker
+from rag.retriever.rerank.local import Reranker
+
+
+def get_reranker(model: str) -> Type[AbstractReranker]:
+    match model:
+        case "local":
+            return Reranker()
+        case "cohere":
+            return CohereReranker()
+        case _:
+            exit(1)
-- 
cgit v1.2.3-70-g09d2