AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-12 13:19:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make fast embed and default embed mutually exclusive. (#4121)

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


### Type of change

- [x] Performance Improvement
This commit is contained in:
Kevin Hu 2024-12-19 17:27:09 +08:00 committed by GitHub
parent b35e811fe7
commit d8fca43017
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
rag/llm/embedding_model.py
View File

@ -251,11 +251,8 @@ class OllamaEmbed(Base):
return np.array(res["embedding"]), 128
class FastEmbed(Base):
_model = None
_model_name = ""
_model_lock = threading.Lock()
class FastEmbed(DefaultEmbedding):
def __init__(
self,
key: str | None = None,
@ -267,17 +264,17 @@ class FastEmbed(Base):
if not settings.LIGHTEN and not FastEmbed._model:
with FastEmbed._model_lock:
from fastembed import TextEmbedding
if not FastEmbed._model or model_name != FastEmbed._model_name:
if not DefaultEmbedding._model or model_name != DefaultEmbedding._model_name:
try:
FastEmbed._model = TextEmbedding(model_name, cache_dir, threads, **kwargs)
FastEmbed._model_name = model_name
DefaultEmbedding._model = TextEmbedding(model_name, cache_dir, threads, **kwargs)
DefaultEmbedding._model_name = model_name
except Exception:
cache_dir = snapshot_download(repo_id="BAAI/bge-small-en-v1.5",
local_dir=os.path.join(get_home_cache_dir(),
re.sub(r"^[a-zA-Z0-9]+/", "", model_name)),
local_dir_use_symlinks=False)
FastEmbed._model = TextEmbedding(model_name, cache_dir, threads, **kwargs)
self._model = FastEmbed._model
DefaultEmbedding._model = TextEmbedding(model_name, cache_dir, threads, **kwargs)
self._model = DefaultEmbedding._model
self._model_name = model_name
def encode(self, texts: list):