jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor OllamaLLM constructor parameters

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-25 19:53:51 -03:00
commit d258e0faea
1 changed files with 19 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

41
src/backend/langflow/components/llms/OllamaLLM.py
View file

@ -2,6 +2,7 @@ from typing import List, Optional
from langchain.llms.base import BaseLLM
from langchain_community.llms.ollama import Ollama
from langflow import CustomComponent
@ -132,29 +133,25 @@ class OllamaLLM(CustomComponent):
mirostat_eta = None
mirostat_tau = None
llm_params = {
"base_url": base_url,
"model": model,
"mirostat": mirostat_value,
"mirostat_eta": mirostat_eta,
"mirostat_tau": mirostat_tau,
"num_ctx": num_ctx,
"num_gpu": num_gpu,
"num_thread": num_thread,
"repeat_last_n": repeat_last_n,
"repeat_penalty": repeat_penalty,
"temperature": temperature,
"stop": stop,
"tfs_z": tfs_z,
"top_k": top_k,
"top_p": top_p,
}
# None Value remove
llm_params = {k: v for k, v in llm_params.items() if v is not None}
try:
llm = Ollama(**llm_params)
llm = Ollama(
base_url=base_url,
model=model,
mirostat=mirostat_value,
mirostat_eta=mirostat_eta,
mirostat_tau=mirostat_tau,
num_ctx=num_ctx,
num_gpu=num_gpu,
num_thread=num_thread,
repeat_last_n=repeat_last_n,
repeat_penalty=repeat_penalty,
temperature=temperature,
stop=stop,
tfs_z=tfs_z,
top_k=top_k,
top_p=top_p,
)
except Exception as e:
raise ValueError("Could not connect to Ollama.") from e