|  | @@ -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
 |