AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-14 02:35:55 +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 return np.array(res["embedding"]), 128
class FastEmbed(Base): class FastEmbed(DefaultEmbedding):
_model = None
_model_name = ""
_model_lock = threading.Lock()
def __init__( def __init__(
self, self,
key: str | None = None, key: str | None = None,
@ -267,17 +264,17 @@ class FastEmbed(Base):
if not settings.LIGHTEN and not FastEmbed._model: if not settings.LIGHTEN and not FastEmbed._model:
with FastEmbed._model_lock: with FastEmbed._model_lock:
from fastembed import TextEmbedding 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: try:
FastEmbed._model = TextEmbedding(model_name, cache_dir, threads, **kwargs) DefaultEmbedding._model = TextEmbedding(model_name, cache_dir, threads, **kwargs)
FastEmbed._model_name = model_name DefaultEmbedding._model_name = model_name
except Exception: except Exception:
cache_dir = snapshot_download(repo_id="BAAI/bge-small-en-v1.5", cache_dir = snapshot_download(repo_id="BAAI/bge-small-en-v1.5",
local_dir=os.path.join(get_home_cache_dir(), local_dir=os.path.join(get_home_cache_dir(),
re.sub(r"^[a-zA-Z0-9]+/", "", model_name)), re.sub(r"^[a-zA-Z0-9]+/", "", model_name)),
local_dir_use_symlinks=False) local_dir_use_symlinks=False)
FastEmbed._model = TextEmbedding(model_name, cache_dir, threads, **kwargs) DefaultEmbedding._model = TextEmbedding(model_name, cache_dir, threads, **kwargs)
self._model = FastEmbed._model self._model = DefaultEmbedding._model
self._model_name = model_name self._model_name = model_name
def encode(self, texts: list): def encode(self, texts: list):