pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

get keep alive from env (#9039)

Browse Source 
### What problem does this PR solve?

get keepalive from env

### Type of change

- [x] Refactoring
This commit is contained in:
Stephen Hu
2025-07-25 12:16:33 +08:00
committed by GitHub
parent 6aaad85cc6
commit 53b0b0e583
2 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
rag/llm/cv_model.py
View File

@ -469,6 +469,7 @@ class OllamaCV(Base):
self.client = Client(host=kwargs["base_url"]) self.client = Client(host=kwargs["base_url"])
self.model_name = model_name self.model_name = model_name
self.lang = lang self.lang = lang
self.keep_alive = kwargs.get("ollama_keep_alive", int(os.environ.get("OLLAMA_KEEP_ALIVE", -1)))
def describe(self, image): def describe(self, image):
prompt = self.prompt("") prompt = self.prompt("")
@ -517,7 +518,7 @@ class OllamaCV(Base):
model=self.model_name, model=self.model_name,
messages=history, messages=history,
options=options, options=options,
keep_alive=-1, keep_alive=self.keep_alive,
) )
ans = response["message"]["content"].strip() ans = response["message"]["content"].strip()
@ -548,7 +549,7 @@ class OllamaCV(Base):
messages=history, messages=history,
stream=True, stream=True,
options=options, options=options,
keep_alive=-1, keep_alive=self.keep_alive,
) )
for resp in response: for resp in response:
if resp["done"]: if resp["done"]:

5
rag/llm/embedding_model.py
View File

@ -285,6 +285,7 @@ class OllamaEmbed(Base):
def __init__(self, key, model_name, **kwargs): def __init__(self, key, model_name, **kwargs):
self.client = Client(host=kwargs["base_url"]) if not key or key == "x" else Client(host=kwargs["base_url"], headers={"Authorization": f"Bearer {key}"}) self.client = Client(host=kwargs["base_url"]) if not key or key == "x" else Client(host=kwargs["base_url"], headers={"Authorization": f"Bearer {key}"})
self.model_name = model_name self.model_name = model_name
self.keep_alive = kwargs.get("ollama_keep_alive", int(os.environ.get("OLLAMA_KEEP_ALIVE", -1)))
def encode(self, texts: list): def encode(self, texts: list):
arr = [] arr = []
@ -293,7 +294,7 @@ class OllamaEmbed(Base):
# remove special tokens if they exist # remove special tokens if they exist
for token in OllamaEmbed._special_tokens: for token in OllamaEmbed._special_tokens:
txt = txt.replace(token, "") txt = txt.replace(token, "")
res = self.client.embeddings(prompt=txt, model=self.model_name, options={"use_mmap": True}, keep_alive=-1) res = self.client.embeddings(prompt=txt, model=self.model_name, options={"use_mmap": True}, keep_alive=self.keep_alive)
try: try:
arr.append(res["embedding"]) arr.append(res["embedding"])
except Exception as _e: except Exception as _e:
@ -305,7 +306,7 @@ class OllamaEmbed(Base):
# remove special tokens if they exist # remove special tokens if they exist
for token in OllamaEmbed._special_tokens: for token in OllamaEmbed._special_tokens:
text = text.replace(token, "") text = text.replace(token, "")
res = self.client.embeddings(prompt=text, model=self.model_name, options={"use_mmap": True}, keep_alive=-1) res = self.client.embeddings(prompt=text, model=self.model_name, options={"use_mmap": True}, keep_alive=self.keep_alive)
try: try:
return np.array(res["embedding"]), 128 return np.array(res["embedding"]), 128
except Exception as _e: except Exception as _e: