Merge pull request #184 from weaviate/modules/add-support-for-reranke…

…r-jinaai Add support for JinaAI reranker
weaviate · Aug 6, 2024 · c45948a · c45948a
2 parents a5355e2 + 7a68c22
commit c45948a
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 2 deletions.
diff --git a/src/collections/config/types/reranker.ts b/src/collections/config/types/reranker.ts
@@ -8,17 +8,35 @@ export type RerankerVoyageAIConfig = {
   model?: 'rerank-lite-1' | string;
 };
 
+export type RerankerJinaAIConfig = {
+  model?:
+    | 'jina-reranker-v2-base-multilingual'
+    | 'jina-reranker-v1-base-en'
+    | 'jina-reranker-v1-turbo-en'
+    | 'jina-reranker-v1-tiny-en'
+    | 'jina-colbert-v1-en'
+    | string;
+};
+
 export type RerankerConfig =
   | RerankerCohereConfig
   | RerankerTransformersConfig
   | RerankerVoyageAIConfig
   | Record<string, any>
   | undefined;
 
-export type Reranker = 'reranker-cohere' | 'reranker-transformers' | 'reranker-voyageai' | 'none' | string;
+export type Reranker =
+  | 'reranker-cohere'
+  | 'reranker-jinaai'
+  | 'reranker-transformers'
+  | 'reranker-voyageai'
+  | 'none'
+  | string;
 
 export type RerankerConfigType<R> = R extends 'reranker-cohere'
   ? RerankerCohereConfig
+  : R extends 'reranker-jinaai'
+  ? RerankerJinaAIConfig
   : R extends 'reranker-transformers'
   ? RerankerTransformersConfig
   : R extends 'reranker-voyageai'

diff --git a/src/collections/configure/reranker.ts b/src/collections/configure/reranker.ts
@@ -1,4 +1,9 @@
-import { ModuleConfig, RerankerCohereConfig, RerankerVoyageAIConfig } from '../config/types/index.js';
+import {
+  ModuleConfig,
+  RerankerCohereConfig,
+  RerankerJinaAIConfig,
+  RerankerVoyageAIConfig,
+} from '../config/types/index.js';
 
 export default {
   /**
@@ -17,6 +22,22 @@ export default {
       config: config,
     };
   },
+  /**
+   * Create a `ModuleConfig<'reranker-jinaai', RerankerJinaAIConfig>` object for use when reranking using the `reranker-jinaai` module.
+   *
+   * See the [documentation](https://weaviate.io/developers/weaviate/modules/retriever-vectorizer-modules/reranker-jinaai) for detailed usage.
+   *
+   * @param {RerankerJinaAIConfig} [config] The configuration for the `reranker-jinaai` module.
+   * @returns {ModuleConfig<'reranker-jinaai', RerankerJinaAIConfig | undefined>} The configuration object.
+   */
+  jinaai: (
+    config?: RerankerJinaAIConfig
+  ): ModuleConfig<'reranker-jinaai', RerankerJinaAIConfig | undefined> => {
+    return {
+      name: 'reranker-jinaai',
+      config: config,
+    };
+  },
   /**
    * Create a `ModuleConfig<'reranker-transformers', Record<string, never>>` object for use when reranking using the `reranker-transformers` module.
    *