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

Refactor embedding batch_size (#3825)

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

Refactor embedding batch_size. Close #3657

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
- [x] Refactoring
This commit is contained in:
Zhichang Yu 2024-12-03 16:22:39 +08:00 committed by GitHub
parent 934dbc2e2b
commit 92ab7ef659
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 160 additions and 109 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
api/db/services/llm_service.py
View File

@ -232,13 +232,13 @@ class LLMBundle(object):
self.max_length = lm.max_tokens
break
def encode(self, texts: list, batch_size=32):
emd, used_tokens = self.mdl.encode(texts, batch_size)
def encode(self, texts: list):
embeddings, used_tokens = self.mdl.encode(texts)
if not TenantLLMService.increase_usage(
self.tenant_id, self.llm_type, used_tokens):
logging.error(
"LLMBundle.encode can't update token usage for {}/EMBEDDING used_tokens: {}".format(self.tenant_id, used_tokens))
return emd, used_tokens
return embeddings, used_tokens
def encode_queries(self, query: str):
emd, used_tokens = self.mdl.encode_queries(query)
@ -280,7 +280,7 @@ class LLMBundle(object):
logging.error(
"LLMBundle.tts can't update token usage for {}/TTS".format(self.tenant_id))
return
yield chunk
yield chunk
def chat(self, system, history, gen_conf):
txt, used_tokens = self.mdl.chat(system, history, gen_conf)

13
rag/benchmark.py
View File

@ -63,16 +63,13 @@ class Benchmark:
run[query][c["chunk_id"]] = c["similarity"]
return run
def embedding(self, docs, batch_size=16):
vects = []
cnts = [d["content_with_weight"] for d in docs]
for i in range(0, len(cnts), batch_size):
vts, c = self.embd_mdl.encode(cnts[i: i + batch_size])
vects.extend(vts.tolist())
assert len(docs) == len(vects)
def embedding(self, docs):
texts = [d["content_with_weight"] for d in docs]
embeddings, _ = self.embd_mdl.encode(texts)
assert len(docs) == len(embeddings)
vector_size = 0
for i, d in enumerate(docs):
v = vects[i]
v = embeddings[i]
vector_size = len(v)
d["q_%d_vec" % len(v)] = v
return docs, vector_size

248
rag/llm/embedding_model.py
View File

@ -38,7 +38,7 @@ class Base(ABC):
def __init__(self, key, model_name):
pass
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
raise NotImplementedError("Please implement encode method!")
def encode_queries(self, text: str):
@ -78,15 +78,16 @@ class DefaultEmbedding(Base):
use_fp16=torch.cuda.is_available())
self._model = DefaultEmbedding._model
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
batch_size = 16
texts = [truncate(t, 2048) for t in texts]
token_count = 0
for t in texts:
token_count += num_tokens_from_string(t)
res = []
ress = []
for i in range(0, len(texts), batch_size):
res.extend(self._model.encode(texts[i:i + batch_size]).tolist())
return np.array(res), token_count
ress.extend(self._model.encode(texts[i:i + batch_size]).tolist())
return np.array(ress), token_count
def encode_queries(self, text: str):
token_count = num_tokens_from_string(text)
@ -101,12 +102,18 @@ class OpenAIEmbed(Base):
self.client = OpenAI(api_key=key, base_url=base_url)
self.model_name = model_name
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
# OpenAI requires batch size <=16
batch_size = 16
texts = [truncate(t, 8191) for t in texts]
res = self.client.embeddings.create(input=texts,
model=self.model_name)
return np.array([d.embedding for d in res.data]
), res.usage.total_tokens
ress = []
total_tokens = 0
for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size],
model=self.model_name)
ress.extend([d.embedding for d in res.data])
total_tokens += res.usage.total_tokens
return np.array(ress), total_tokens
def encode_queries(self, text):
res = self.client.embeddings.create(input=[truncate(text, 8191)],
@ -123,12 +130,14 @@ class LocalAIEmbed(Base):
self.client = OpenAI(api_key="empty", base_url=base_url)
self.model_name = model_name.split("___")[0]
def encode(self, texts: list, batch_size=16):
res = self.client.embeddings.create(input=texts, model=self.model_name)
return (
np.array([d.embedding for d in res.data]),
1024,
) # local embedding for LmStudio donot count tokens
def encode(self, texts: list):
batch_size = 16
ress = []
for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name)
ress.extend([d.embedding for d in res.data])
# local embedding for LmStudio donot count tokens
return np.array(ress), 1024
def encode_queries(self, text):
embds, cnt = self.encode([text])
@ -155,12 +164,12 @@ class BaiChuanEmbed(OpenAIEmbed):
class QWenEmbed(Base):
def __init__(self, key, model_name="text_embedding_v2", **kwargs):
dashscope.api_key = key
self.key = key
self.model_name = model_name
def encode(self, texts: list, batch_size=10):
def encode(self, texts: list):
import dashscope
batch_size = min(batch_size, 4)
batch_size = 4
try:
res = []
token_count = 0
@ -169,6 +178,7 @@ class QWenEmbed(Base):
resp = dashscope.TextEmbedding.call(
model=self.model_name,
input=texts[i:i + batch_size],
api_key=self.key,
text_type="document"
)
embds = [[] for _ in range(len(resp["output"]["embeddings"]))]
@ -186,6 +196,7 @@ class QWenEmbed(Base):
resp = dashscope.TextEmbedding.call(
model=self.model_name,
input=text[:2048],
api_key=self.key,
text_type="query"
)
return np.array(resp["output"]["embeddings"][0]
@ -200,7 +211,7 @@ class ZhipuEmbed(Base):
self.client = ZhipuAI(api_key=key)
self.model_name = model_name
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
arr = []
tks_num = 0
for txt in texts:
@ -221,7 +232,7 @@ class OllamaEmbed(Base):
self.client = Client(host=kwargs["base_url"])
self.model_name = model_name
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
arr = []
tks_num = 0
for txt in texts:
@ -252,13 +263,13 @@ class FastEmbed(Base):
from fastembed import TextEmbedding
self._model = TextEmbedding(model_name, cache_dir, threads, **kwargs)
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
# Using the internal tokenizer to encode the texts and get the total
# number of tokens
encodings = self._model.model.tokenizer.encode_batch(texts)
total_tokens = sum(len(e) for e in encodings)
embeddings = [e.tolist() for e in self._model.embed(texts, batch_size)]
embeddings = [e.tolist() for e in self._model.embed(texts, batch_size=16)]
return np.array(embeddings), total_tokens
@ -278,11 +289,15 @@ class XinferenceEmbed(Base):
self.client = OpenAI(api_key=key, base_url=base_url)
self.model_name = model_name
def encode(self, texts: list, batch_size=16):
res = self.client.embeddings.create(input=texts,
model=self.model_name)
return np.array([d.embedding for d in res.data]
), res.usage.total_tokens
def encode(self, texts: list):
batch_size = 16
ress = []
total_tokens = 0
for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name)
ress.extend([d.embedding for d in res.data])
total_tokens += res.usage.total_tokens
return np.array(ress), total_tokens
def encode_queries(self, text):
res = self.client.embeddings.create(input=[text],
@ -306,7 +321,8 @@ class YoudaoEmbed(Base):
model_name_or_path=model_name.replace(
"maidalun1020", "InfiniFlow"))
def encode(self, texts: list, batch_size=10):
def encode(self, texts: list):
batch_size = 10
res = []
token_count = 0
for t in texts:
@ -332,15 +348,21 @@ class JinaEmbed(Base):
}
self.model_name = model_name
def encode(self, texts: list, batch_size=None):
def encode(self, texts: list):
texts = [truncate(t, 8196) for t in texts]
data = {
"model": self.model_name,
"input": texts,
'encoding_type': 'float'
}
res = requests.post(self.base_url, headers=self.headers, json=data).json()
return np.array([d["embedding"] for d in res["data"]]), res["usage"]["total_tokens"]
batch_size = 16
ress = []
token_count = 0
for i in range(0, len(texts), batch_size):
data = {
"model": self.model_name,
"input": texts[i:i + batch_size],
'encoding_type': 'float'
}
res = requests.post(self.base_url, headers=self.headers, json=data).json()
ress.extend([d["embedding"] for d in res["data"]])
token_count += res["usage"]["total_tokens"]
return np.array(ress), token_count
def encode_queries(self, text):
embds, cnt = self.encode([text])
@ -394,12 +416,17 @@ class MistralEmbed(Base):
self.client = MistralClient(api_key=key)
self.model_name = model_name
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
texts = [truncate(t, 8196) for t in texts]
res = self.client.embeddings(input=texts,
model=self.model_name)
return np.array([d.embedding for d in res.data]
), res.usage.total_tokens
batch_size = 16
ress = []
token_count = 0
for i in range(0, len(texts), batch_size):
res = self.client.embeddings(input=texts[i:i + batch_size],
model=self.model_name)
ress.extend([d.embedding for d in res.data])
token_count += res.usage.total_tokens
return np.array(ress), token_count
def encode_queries(self, text):
res = self.client.embeddings(input=[truncate(text, 8196)],
@ -418,7 +445,7 @@ class BedrockEmbed(Base):
self.client = boto3.client(service_name='bedrock-runtime', region_name=self.bedrock_region,
aws_access_key_id=self.bedrock_ak, aws_secret_access_key=self.bedrock_sk)
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
texts = [truncate(t, 8196) for t in texts]
embeddings = []
token_count = 0
@ -436,7 +463,6 @@ class BedrockEmbed(Base):
return np.array(embeddings), token_count
def encode_queries(self, text):
embeddings = []
token_count = num_tokens_from_string(text)
if self.model_name.split('.')[0] == 'amazon':
@ -453,20 +479,26 @@ class BedrockEmbed(Base):
class GeminiEmbed(Base):
def __init__(self, key, model_name='models/text-embedding-004',
**kwargs):
genai.configure(api_key=key)
self.key = key
self.model_name = 'models/' + model_name
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
texts = [truncate(t, 2048) for t in texts]
token_count = sum(num_tokens_from_string(text) for text in texts)
result = genai.embed_content(
model=self.model_name,
content=texts,
task_type="retrieval_document",
title="Embedding of list of strings")
return np.array(result['embedding']),token_count
genai.configure(api_key=self.key)
batch_size = 16
ress = []
for i in range(0, len(texts), batch_size):
result = genai.embed_content(
model=self.model_name,
content=texts[i, i + batch_size],
task_type="retrieval_document",
title="Embedding of single string")
ress.extend(result['embedding'])
return np.array(ress),token_count
def encode_queries(self, text):
genai.configure(api_key=self.key)
result = genai.embed_content(
model=self.model_name,
content=truncate(text,2048),
@ -495,19 +527,22 @@ class NvidiaEmbed(Base):
if model_name == "snowflake/arctic-embed-l":
self.base_url = "https://ai.api.nvidia.com/v1/retrieval/snowflake/arctic-embed-l/embeddings"
def encode(self, texts: list, batch_size=None):
payload = {
"input": texts,
"input_type": "query",
"model": self.model_name,
"encoding_format": "float",
"truncate": "END",
}
res = requests.post(self.base_url, headers=self.headers, json=payload).json()
return (
np.array([d["embedding"] for d in res["data"]]),
res["usage"]["total_tokens"],
)
def encode(self, texts: list):
batch_size = 16
ress = []
token_count = 0
for i in range(0, len(texts), batch_size):
payload = {
"input": texts[i : i + batch_size],
"input_type": "query",
"model": self.model_name,
"encoding_format": "float",
"truncate": "END",
}
res = requests.post(self.base_url, headers=self.headers, json=payload).json()
ress.extend([d["embedding"] for d in res["data"]])
token_count += res["usage"]["total_tokens"]
return np.array(ress), token_count
def encode_queries(self, text):
embds, cnt = self.encode([text])
@ -541,16 +576,20 @@ class CoHereEmbed(Base):
self.client = Client(api_key=key)
self.model_name = model_name
def encode(self, texts: list, batch_size=16):
res = self.client.embed(
texts=texts,
model=self.model_name,
input_type="search_query",
embedding_types=["float"],
)
return np.array([d for d in res.embeddings.float]), int(
res.meta.billed_units.input_tokens
)
def encode(self, texts: list):
batch_size = 16
ress = []
token_count = 0
for i in range(0, len(texts), batch_size):
res = self.client.embed(
texts=texts[i : i + batch_size],
model=self.model_name,
input_type="search_document",
embedding_types=["float"],
)
ress.extend([d for d in res.embeddings.float])
token_count += res.meta.billed_units.input_tokens
return np.array(ress), token_count
def encode_queries(self, text):
res = self.client.embed(
@ -599,19 +638,23 @@ class SILICONFLOWEmbed(Base):
self.base_url = base_url
self.model_name = model_name
def encode(self, texts: list, batch_size=16):
payload = {
"model": self.model_name,
"input": texts,
"encoding_format": "float",
}
res = requests.post(self.base_url, json=payload, headers=self.headers).json()
if "data" not in res or not isinstance(res["data"], list) or len(res["data"])!= len(texts):
raise ValueError(f"SILICONFLOWEmbed.encode got invalid response from {self.base_url}")
return (
np.array([d["embedding"] for d in res["data"]]),
res["usage"]["total_tokens"],
)
def encode(self, texts: list):
batch_size = 16
ress = []
token_count = 0
for i in range(0, len(texts), batch_size):
texts_batch = texts[i : i + batch_size]
payload = {
"model": self.model_name,
"input": texts_batch,
"encoding_format": "float",
}
res = requests.post(self.base_url, json=payload, headers=self.headers).json()
if "data" not in res or not isinstance(res["data"], list) or len(res["data"]) != len(texts_batch):
raise ValueError(f"SILICONFLOWEmbed.encode got invalid response from {self.base_url}")
ress.extend([d["embedding"] for d in res["data"]])
token_count += res["usage"]["total_tokens"]
return np.array(ress), token_count
def encode_queries(self, text):
payload = {
@ -632,9 +675,14 @@ class ReplicateEmbed(Base):
self.model_name = model_name
self.client = Client(api_token=key)
def encode(self, texts: list, batch_size=16):
res = self.client.run(self.model_name, input={"texts": json.dumps(texts)})
return np.array(res), sum([num_tokens_from_string(text) for text in texts])
def encode(self, texts: list):
batch_size = 16
token_count = sum([num_tokens_from_string(text) for text in texts])
ress = []
for i in range(0, len(texts), batch_size):
res = self.client.run(self.model_name, input={"texts": texts[i : i + batch_size]})
ress.extend(res)
return np.array(ress), token_count
def encode_queries(self, text):
res = self.client.embed(self.model_name, input={"texts": [text]})
@ -673,11 +721,17 @@ class VoyageEmbed(Base):
self.client = voyageai.Client(api_key=key)
self.model_name = model_name
def encode(self, texts: list, batch_size=16):
res = self.client.embed(
texts=texts, model=self.model_name, input_type="document"
)
return np.array(res.embeddings), res.total_tokens
def encode(self, texts: list):
batch_size = 16
ress = []
token_count = 0
for i in range(0, len(texts), batch_size):
res = self.client.embed(
texts=texts[i : i + batch_size], model=self.model_name, input_type="document"
)
ress.extend(res.embeddings)
token_count += res.total_tokens
return np.array(ress), token_count
def encode_queries(self, text):
res = self.client.embed(
@ -694,7 +748,7 @@ class HuggingFaceEmbed(Base):
self.model_name = model_name
self.base_url = base_url or "http://127.0.0.1:8080"
def encode(self, texts: list, batch_size=16):
def encode(self, texts: list):
embeddings = []
for text in texts:
response = requests.post(