新增ollama重排类

2024-09-06 10:07:22 +08:00
parent 21fdc16259
commit 60b0f11ca2
3 changed files with 83 additions and 1 deletions
@@ -88,7 +88,18 @@ class OllamaPlatform(ModelPlatform):
        pass

    def rerank(self):
-        pass
+        from app.engine.rerank.ollamRerank import OllamaRerank
+        modelpath = os.getcwd() + os.getenv('RERANK_MODEL')
+        top_n = os.getenv('RERANK_TOP_N',5)
+        threshold =  float(os.getenv('RERANK_THRESHOLD',0.3))
+        rerank = OllamaRerank(
+            model=modelpath,
+            top_n=top_n, 
+            device="cpu",
+            score_threshold= threshold
+        )
+        return [rerank]
+

@register(ModelPlateCategory,'xinference')
 class XinferencePlatform(ModelPlatform):