### What problem does this PR solve?


### Type of change

- [x] Refactoring
This commit is contained in:
KevinHuSh
2024-06-12 11:02:15 +08:00
committed by GitHub
parent 2cc89211f6
commit abcd3d2469
7 changed files with 73 additions and 83 deletions

View File

@ -26,9 +26,8 @@ import dashscope
from openai import OpenAI
from FlagEmbedding import FlagModel
import torch
import asyncio
import numpy as np
import asyncio
from api.utils.file_utils import get_home_cache_dir
from rag.utils import num_tokens_from_string, truncate
@ -317,12 +316,12 @@ class InfinityEmbed(Base):
engine_kwargs: dict = {},
key = None,
):
from infinity_emb import EngineArgs
from infinity_emb.engine import AsyncEngineArray
self._default_model = model_names[0]
self.engine_array = AsyncEngineArray.from_args([EngineArgs(model_name_or_path = model_name, **engine_kwargs) for model_name in model_names])
self.engine_array = AsyncEngineArray.from_args([EngineArgs(model_name_or_path = model_name, **engine_kwargs) for model_name in model_names])
async def _embed(self, sentences: list[str], model_name: str = ""):
if not model_name: