增加对Rerank功能支持

2024-08-19 08:59:08 +08:00
parent 8d4382376f
commit 2942730c9a
3 changed files with 18 additions and 3 deletions
@@ -6,9 +6,17 @@ from llama_index.core.settings import Settings
 from llama_index.llms.xinference import Xinference
 from llama_index.llms.xinference.base import DEFAULT_XINFERENCE_TEMP

-from app.xinference.base import XinferenceEmbedding
+from app.xinference.base import XinferenceEmbedding, XinferenceRerank


+def get_node_postprocessors():
+    rerank_model = os.getenv("RERANK_MODEL")
+    rerank_url = os.getenv("RERANK_BASE_URL")
+    postprocess = None
+    if rerank_model is None:
+        postprocess = [XinferenceRerank(rerank_model, rerank_url)]
+    return postprocess
+
 def init_settings():
    model_provider = os.getenv("MODEL_PROVIDER")
    match model_provider: