From c5fd04eff3a89891bddabc235681bb1c0e0083d5 Mon Sep 17 00:00:00 2001 From: Xiao Zhang Date: Mon, 1 Jul 2024 21:05:15 +0800 Subject: [PATCH] fix ollama mode can not work (#154) --- aios/llm_kernel/llm_classes/ollama_llm.py | 4 +++- aios/llm_kernel/llms.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/aios/llm_kernel/llm_classes/ollama_llm.py b/aios/llm_kernel/llm_classes/ollama_llm.py index 6e479313..780b74fb 100755 --- a/aios/llm_kernel/llm_classes/ollama_llm.py +++ b/aios/llm_kernel/llm_classes/ollama_llm.py @@ -72,7 +72,9 @@ def process(self, response = ollama.chat( model=self.model_name.split("/")[-1], messages=messages, - num_predict=self.max_new_tokens + options= ollama.Options( + num_predict=self.max_new_tokens + ) ) agent_process.set_response( Response( diff --git a/aios/llm_kernel/llms.py b/aios/llm_kernel/llms.py index 31dc45ac..dfcd56c4 100644 --- a/aios/llm_kernel/llms.py +++ b/aios/llm_kernel/llms.py @@ -27,7 +27,7 @@ def __init__(self, ) # For locally-deployed LLM else: - if use_backend == "ollama" and llm_name.startswith("ollama"): + if use_backend == "ollama" or llm_name.startswith("ollama"): self.model = OllamaLLM( llm_name=llm_name, max_gpu_memory=max_gpu_memory,