set ollama keep_alive (#985)

### What problem does this PR solve? #980 ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue)
infiniflow · May 30, 2024 · f077b57 · f077b57
1 parent c62834f
commit f077b57
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/api/db/init_data.py b/api/db/init_data.py
@@ -532,8 +532,8 @@ def init_llm_factory():
         {
             "fid": factory_infos[12]["name"],
             "llm_name": "BAAI/bge-reranker-v2-m3",
-            "tags": "LLM,CHAT,",
-            "max_tokens": 16385,
+            "tags": "RE-RANK,2k",
+            "max_tokens": 2048,
             "model_type": LLMType.RERANK.value
         },
     ]

diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
@@ -303,7 +303,8 @@ def chat(self, system, history, gen_conf):
             response = self.client.chat(
                 model=self.model_name,
                 messages=history,
-                options=options
+                options=options,
+                keep_alive=-1
             )
             ans = response["message"]["content"].strip()
             return ans, response["eval_count"] + response.get("prompt_eval_count", 0)
@@ -325,7 +326,8 @@ def chat_streamly(self, system, history, gen_conf):
                 model=self.model_name,
                 messages=history,
                 stream=True,
-                options=options
+                options=options,
+                keep_alive=-1
             )
             for resp in response:
                 if resp["done"]: