CrossEncoderModule with rerank API

This module is closely related to EmbeddingModule. Cross-encoder models use Q and A pairs and are trained return a relevance score for rank(). The existing rerank APIs in EmbeddingModule had to encode Q and A separately and use cosine similarity as a score. So the API is the same, but the results are supposed to be better (and slower). Cross-encoder models do not support returning embedding vectors or sentence-similarity. Support for the existing tokenization and model_info endpoints was also added. Signed-off-by: Mark Sturdevant <mark.sturdevant@ibm.com>
caikit · Sep 11, 2024 · c118db3 · c118db3
1 parent 4a027d3
commit c118db3
Show file tree

Hide file tree

Showing 3 changed files with 1,250 additions and 0 deletions.
diff --git a/caikit_nlp/modules/text_embedding/__init__.py b/caikit_nlp/modules/text_embedding/__init__.py
@@ -29,4 +29,5 @@
 """
 
 # Local
+from .crossencoder import CrossEncoderModule
 from .embedding import EmbeddingModule
-Original file line number
+Diff line change
@@ Expand Up / @@ -29,4 +29,5 @@ @@
     """
     # Local
+    from .crossencoder import CrossEncoderModule
     from .embedding import EmbeddingModule