Spaces:
Running
Running
| // Copyright 2025 The ODML Authors. | |
| // | |
| // Licensed under the Apache License, Version 2.0 (the "License"); | |
| // you may not use this file except in compliance with the License. | |
| // You may obtain a copy of the License at | |
| // | |
| // http://www.apache.org/licenses/LICENSE-2.0 | |
| // | |
| // Unless required by applicable law or agreed to in writing, software | |
| // distributed under the License is distributed on an "AS IS" BASIS, | |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| // See the License for the specific language governing permissions and | |
| // limitations under the License. | |
| namespace litert::lm { | |
| class EmbeddingLookupManager { | |
| public: | |
| // Creates an EmbeddingLookupManager. | |
| // | |
| // The end_of_multi_modal_embedding_models is a map of special tokens to the | |
| // corresponding embedding models. The special tokens are used to indicate | |
| // that the corresponding embedding model should be used. | |
| // | |
| // If fully_supports_multi_modal is true, the EmbeddingLookupManager will | |
| // handle multimodal tokens via the multimodal embedding lookup. Otherwise, it | |
| // default any multi-modal tokens to the text embedding value of entry 0. | |
| // If fully_supports_multi_modal is false, the | |
| // end_of_multi_modal_embedding_models must be empty. | |
| // | |
| // If the provide text_embedding_model has more than one signature, the | |
| // signature_key must be provided. | |
| static absl::StatusOr<std::unique_ptr<EmbeddingLookupManager>> Create( | |
| const litert::Model* absl_nonnull text_embedding_model, | |
| absl::flat_hash_map<int, const litert::Model*>& | |
| end_of_multi_modal_embedding_models, | |
| bool fully_supports_multi_modal = true, | |
| std::optional<std::string> signature_key = std::nullopt, | |
| litert::Environment* env = nullptr); | |
| static absl::StatusOr<std::unique_ptr<EmbeddingLookupManager>> Create( | |
| const litert::Model* absl_nonnull text_embedding_model, | |
| bool fully_supports_multi_modal = true, | |
| std::optional<std::string> signature_key = std::nullopt, | |
| litert::Environment* env = nullptr); | |
| // Updates the multimodal embeddings for the given ExecutorInputs. | |
| // Intended to be called at the beginning of the prefill pass. | |
| // | |
| // If fully_supports_multi_modal_ is false, this function will return an error | |
| // if the ExecutorInputs contain any multimodal embeddings. | |
| absl::Status UpdateMultiModalEmbeddings( | |
| const ::litert::lm::ExecutorInputs& inputs); | |
| // Cleans up the multimodal embeddings and verifies that all the embeddings | |
| // have been used. | |
| // Intended to be called at the end of the prefill pass. | |
| absl::Status CleanupMultiModalEmbeddings(); | |
| // For a given token, looks up the embedding and stores it in the output | |
| // vector. | |
| // | |
| // This is used for the case where the llm_litert_executor needs to look up | |
| // embeddings for the current step and then use the result for the next step. | |
| // At that point, it does not have a TfLiteTensor to store the result in. | |
| absl::Status LookupDecode(int token, std::vector<float>& output_vector); | |
| // For a given token, looks up the embedding and stores it in the | |
| // output tensor. | |
| absl::Status LookupDecode(int token, litert::TensorBuffer* output_tensor); | |
| // For a given token, looks up the embedding and stores it in the provided | |
| // vector. This function is responsible for setting the size of the vector to | |
| // the correct size and filling it with the embedding. Any data that was | |
| // previously in the vector will be overwritten. | |
| // | |
| // This is used for the case where the llm_litert_executor needs to look up | |
| // embeddings for the current step and then use the result for the next step. | |
| // At that point, it does not have a TfLiteTensor to store the result in. | |
| absl::Status LookupPrefill(int token, std::vector<float>& output_vector); | |
| // For a given list of tokens, looks up the embeddings, concatenates them and | |
| // returns the result through the output tensor. | |
| // | |
| // token_offset is used to indicate where to start writing to in the | |
| // output_tensor. This is used in cases where the output_tensor has already | |
| // had some embeddings written to it. If this is the first time embeddings are | |
| // being written to the output_tensor, token_offset should be 0. | |
| absl::Status LookupPrefill(absl::Span<const int> tokens, | |
| litert::TensorBuffer* output_tensor, | |
| size_t token_offset); | |
| EmbeddingLookupText* GetTextEmbeddingLookup() const { | |
| return text_embedding_lookup_.get(); | |
| } | |
| protected: | |
| absl::Status Initialize( | |
| const litert::Model* absl_nonnull text_embedding_model, | |
| absl::flat_hash_map<int, const litert::Model*>& | |
| end_of_multi_modal_embedding_models, | |
| bool fully_supports_multi_modal, std::optional<std::string> signature_key, | |
| litert::Environment* env = nullptr); | |
| std::unique_ptr<EmbeddingLookupText> text_embedding_lookup_; | |
| std::vector<std::unique_ptr<EmbeddingLookupMultiModal>> | |
| multi_modal_embedding_lookups_; | |
| std::vector<std::unique_ptr<EndOfMultiModalEmbedding>> | |
| end_of_multi_modal_embedding_lookups_; | |
| // If true, the EmbeddingLookupManager will support multimodal embeddings. | |
| // Otherwise, it will default any multimodal tokens to the text embedding | |
| // value of entry 0. | |
| bool fully_supports_multi_modal_; | |
| }; | |
| } // namespace litert::lm | |