Generate an embedding vector for each prompt. This class automatically batches the given prompts, considering the memory constraint. For the best performance, put all of your prompts into a single list and pass it to this method. Args: prompts:
(
self,
prompts: Union[PromptType, Sequence[PromptType]],
/,
*,
use_tqdm: bool = True,
pooling_params: Optional[Union[PoolingParams,
Sequence[PoolingParams]]] = None,
lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,
prompt_adapter_request: Optional[PromptAdapterRequest] = None,
)
| 958 | PoolingRequestOutput) |
| 959 | |
| 960 | def embed( |
| 961 | self, |
| 962 | prompts: Union[PromptType, Sequence[PromptType]], |
| 963 | /, |
| 964 | *, |
| 965 | use_tqdm: bool = True, |
| 966 | pooling_params: Optional[Union[PoolingParams, |
| 967 | Sequence[PoolingParams]]] = None, |
| 968 | lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None, |
| 969 | prompt_adapter_request: Optional[PromptAdapterRequest] = None, |
| 970 | ) -> list[EmbeddingRequestOutput]: |
| 971 | """ |
| 972 | Generate an embedding vector for each prompt. |
| 973 | |
| 974 | This class automatically batches the given prompts, considering |
| 975 | the memory constraint. For the best performance, put all of your prompts |
| 976 | into a single list and pass it to this method. |
| 977 | |
| 978 | Args: |
| 979 | prompts: The prompts to the LLM. You may pass a sequence of prompts |
| 980 | for batch inference. See :class:`~vllm.inputs.PromptType` |
| 981 | for more details about the format of each prompts. |
| 982 | pooling_params: The pooling parameters for pooling. If None, we |
| 983 | use the default pooling parameters. |
| 984 | use_tqdm: Whether to use tqdm to display the progress bar. |
| 985 | lora_request: LoRA request to use for generation, if any. |
| 986 | prompt_adapter_request: Prompt Adapter request to use for |
| 987 | generation, if any. |
| 988 | |
| 989 | Returns: |
| 990 | A list of ``EmbeddingRequestOutput`` objects containing the |
| 991 | embedding vectors in the same order as the input prompts. |
| 992 | """ |
| 993 | if self.llm_engine.model_config.task != "embed": |
| 994 | raise ValueError( |
| 995 | "Embedding API is only enabled for `--task embed`") |
| 996 | |
| 997 | items = self.encode(prompts, |
| 998 | use_tqdm=use_tqdm, |
| 999 | pooling_params=pooling_params, |
| 1000 | lora_request=lora_request, |
| 1001 | prompt_adapter_request=prompt_adapter_request) |
| 1002 | |
| 1003 | return [EmbeddingRequestOutput.from_base(item) for item in items] |
| 1004 | |
| 1005 | def classify( |
| 1006 | self, |