|
|
|
@@ -2,6 +2,7 @@ import time |
|
|
|
from typing import Optional |
|
|
|
|
|
|
|
import dashscope |
|
|
|
import numpy as np |
|
|
|
|
|
|
|
from core.model_runtime.entities.model_entities import PriceType |
|
|
|
from core.model_runtime.entities.text_embedding_entities import ( |
|
|
|
@@ -21,11 +22,11 @@ class TongyiTextEmbeddingModel(_CommonTongyi, TextEmbeddingModel): |
|
|
|
""" |
|
|
|
|
|
|
|
def _invoke( |
|
|
|
self, |
|
|
|
model: str, |
|
|
|
credentials: dict, |
|
|
|
texts: list[str], |
|
|
|
user: Optional[str] = None, |
|
|
|
self, |
|
|
|
model: str, |
|
|
|
credentials: dict, |
|
|
|
texts: list[str], |
|
|
|
user: Optional[str] = None, |
|
|
|
) -> TextEmbeddingResult: |
|
|
|
""" |
|
|
|
Invoke text embedding model |
|
|
|
@@ -37,16 +38,44 @@ class TongyiTextEmbeddingModel(_CommonTongyi, TextEmbeddingModel): |
|
|
|
:return: embeddings result |
|
|
|
""" |
|
|
|
credentials_kwargs = self._to_credential_kwargs(credentials) |
|
|
|
embeddings, embedding_used_tokens = self.embed_documents( |
|
|
|
credentials_kwargs=credentials_kwargs, |
|
|
|
model=model, |
|
|
|
texts=texts |
|
|
|
) |
|
|
|
|
|
|
|
context_size = self._get_context_size(model, credentials) |
|
|
|
max_chunks = self._get_max_chunks(model, credentials) |
|
|
|
inputs = [] |
|
|
|
indices = [] |
|
|
|
used_tokens = 0 |
|
|
|
|
|
|
|
for i, text in enumerate(texts): |
|
|
|
|
|
|
|
# Here token count is only an approximation based on the GPT2 tokenizer |
|
|
|
num_tokens = self._get_num_tokens_by_gpt2(text) |
|
|
|
|
|
|
|
if num_tokens >= context_size: |
|
|
|
cutoff = int(np.floor(len(text) * (context_size / num_tokens))) |
|
|
|
# if num tokens is larger than context length, only use the start |
|
|
|
inputs.append(text[0:cutoff]) |
|
|
|
else: |
|
|
|
inputs.append(text) |
|
|
|
indices += [i] |
|
|
|
|
|
|
|
batched_embeddings = [] |
|
|
|
_iter = range(0, len(inputs), max_chunks) |
|
|
|
|
|
|
|
for i in _iter: |
|
|
|
embeddings_batch, embedding_used_tokens = self.embed_documents( |
|
|
|
credentials_kwargs=credentials_kwargs, |
|
|
|
model=model, |
|
|
|
texts=inputs[i : i + max_chunks], |
|
|
|
) |
|
|
|
used_tokens += embedding_used_tokens |
|
|
|
batched_embeddings += embeddings_batch |
|
|
|
|
|
|
|
# calc usage |
|
|
|
usage = self._calc_response_usage( |
|
|
|
model=model, credentials=credentials, tokens=used_tokens |
|
|
|
) |
|
|
|
return TextEmbeddingResult( |
|
|
|
embeddings=embeddings, |
|
|
|
usage=self._calc_response_usage(model, credentials_kwargs, embedding_used_tokens), |
|
|
|
model=model |
|
|
|
embeddings=batched_embeddings, usage=usage, model=model |
|
|
|
) |
|
|
|
|
|
|
|
def get_num_tokens(self, model: str, credentials: dict, texts: list[str]) -> int: |
|
|
|
@@ -79,12 +108,16 @@ class TongyiTextEmbeddingModel(_CommonTongyi, TextEmbeddingModel): |
|
|
|
credentials_kwargs = self._to_credential_kwargs(credentials) |
|
|
|
|
|
|
|
# call embedding model |
|
|
|
self.embed_documents(credentials_kwargs=credentials_kwargs, model=model, texts=["ping"]) |
|
|
|
self.embed_documents( |
|
|
|
credentials_kwargs=credentials_kwargs, model=model, texts=["ping"] |
|
|
|
) |
|
|
|
except Exception as ex: |
|
|
|
raise CredentialsValidateFailedError(str(ex)) |
|
|
|
|
|
|
|
@staticmethod |
|
|
|
def embed_documents(credentials_kwargs: dict, model: str, texts: list[str]) -> tuple[list[list[float]], int]: |
|
|
|
def embed_documents( |
|
|
|
credentials_kwargs: dict, model: str, texts: list[str] |
|
|
|
) -> tuple[list[list[float]], int]: |
|
|
|
"""Call out to Tongyi's embedding endpoint. |
|
|
|
|
|
|
|
Args: |
|
|
|
@@ -102,7 +135,7 @@ class TongyiTextEmbeddingModel(_CommonTongyi, TextEmbeddingModel): |
|
|
|
api_key=credentials_kwargs["dashscope_api_key"], |
|
|
|
model=model, |
|
|
|
input=text, |
|
|
|
text_type="document" |
|
|
|
text_type="document", |
|
|
|
) |
|
|
|
data = response.output["embeddings"][0] |
|
|
|
embeddings.append(data["embedding"]) |
|
|
|
@@ -111,7 +144,7 @@ class TongyiTextEmbeddingModel(_CommonTongyi, TextEmbeddingModel): |
|
|
|
return [list(map(float, e)) for e in embeddings], embedding_used_tokens |
|
|
|
|
|
|
|
def _calc_response_usage( |
|
|
|
self, model: str, credentials: dict, tokens: int |
|
|
|
self, model: str, credentials: dict, tokens: int |
|
|
|
) -> EmbeddingUsage: |
|
|
|
""" |
|
|
|
Calculate response usage |
|
|
|
@@ -125,7 +158,7 @@ class TongyiTextEmbeddingModel(_CommonTongyi, TextEmbeddingModel): |
|
|
|
model=model, |
|
|
|
credentials=credentials, |
|
|
|
price_type=PriceType.INPUT, |
|
|
|
tokens=tokens |
|
|
|
tokens=tokens, |
|
|
|
) |
|
|
|
|
|
|
|
# transform usage |
|
|
|
@@ -136,7 +169,7 @@ class TongyiTextEmbeddingModel(_CommonTongyi, TextEmbeddingModel): |
|
|
|
price_unit=input_price_info.unit, |
|
|
|
total_price=input_price_info.total_amount, |
|
|
|
currency=input_price_info.currency, |
|
|
|
latency=time.perf_counter() - self.started_at |
|
|
|
latency=time.perf_counter() - self.started_at, |
|
|
|
) |
|
|
|
|
|
|
|
return usage |