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 used for looking up text embeddings on the CPU. | |
| // | |
| // Ideally text embedding lookups should be a part of the main model but there | |
| // are cases where the embedding lookup needs to be done separately for now. For | |
| // example, large embedding tables may use too much memory on the accelerator | |
| // and so they need to be placed on the CPU. Currently there is no mechanism | |
| // to tell a delegate to move embedding lookups to the CPU. | |
| class EmbeddingLookupText : public EmbeddingLookup { | |
| public: | |
| ~EmbeddingLookupText() override = default; | |
| // Creates a EmbeddingLookupText instance. The reference of |model| is kept | |
| // in the returned instance, so the caller must ensure that |model| outlives | |
| // the returned instance. If the model has more than one signature, and | |
| // signature_key is not provided, the first signature will be used by default. | |
| static absl::StatusOr<std::unique_ptr<EmbeddingLookupText>> Create( | |
| const litert::Model* absl_nonnull model, | |
| std::optional<std::string> signature_key = std::nullopt, | |
| litert::Environment* absl_nullable env = nullptr); | |
| // For a given token, looks up the embedding and stores it in the | |
| // provided vector. The caller is responsible for ensuring that the vector is | |
| // the correct size for the embedding. | |
| // | |
| // 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 LiteRtTensor to store the result in. | |
| absl::Status LookupDecode(int token, | |
| std::vector<float>& decode_output_vector) override; | |
| // For a given token, looks up the embedding and stores it in the | |
| // output tensor. | |
| absl::Status LookupDecode(int token, | |
| litert::TensorBuffer* decode_output) override; | |
| // For a given token, looks up the embedding and stores it in the | |
| // provided vector. The caller is responsible for ensuring that the vector is | |
| // the correct size for the embedding model output. | |
| // | |
| // 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 LiteRtTensor to store the result in. | |
| absl::Status LookupPrefill( | |
| int token, std::vector<float>& prefill_output_vector) override; | |
| // For a given list of tokens, looks up the embeddings, concatenates them and | |
| // returns the result through the output tensor. | |
| // | |
| // Support is only partially implemented right now. This function only | |
| // supports the case where the output tensor's 0th dimension is of size | |
| // 1, its 1st dimension is greater than or equal to tokens.size(), and its | |
| // subsequent dimensions match the dimensions of the embedding model output. | |
| // In other words, if the embedding model outputs [B=1, T=1, N, H], then the | |
| // output tensor must be [1, >=tokens.size(), N, H]. | |
| // | |
| // bytes_offset is used to indicate what byte 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. | |
| absl::Status LookupPrefill(absl::Span<const int> tokens, | |
| litert::TensorBuffer* prefill_output, | |
| size_t byte_offset) override; | |
| // Returns number of floats per token in the output tensor. | |
| size_t GetFloatsPerToken(); | |
| // Returns the default embedding vector to use when a token is not found in | |
| // the lookup table. | |
| const std::vector<float>& GetDefaultEmbeddingVector() const { | |
| return default_embedding_vector_; | |
| } | |
| // Returns the signature key to use for the embedding model. | |
| std::optional<litert::RankedTensorType> GetOutputBufferType() const { | |
| return output_buffer_type_; | |
| } | |
| protected: | |
| EmbeddingLookupText(litert::Environment& env, | |
| const litert::Model* absl_nonnull model, | |
| std::optional<std::string> signature_key) | |
| : env_(env), model_(*model), signature_key_(signature_key) {} | |
| // Loads the provided model. This must be called before Lookup. | |
| absl::Status Initialize(); | |
| // Internal implementation of Lookup for both the single and multiple token | |
| // cases. | |
| absl::Status LookupInternal(int token, absl::Span<uint8_t> buffer); | |
| // The environment for the embedding lookup. | |
| litert::Environment& env_; | |
| // The model for the embedding lookup. The actual model instance is owned by | |
| // the model resources. | |
| const litert::Model& model_; | |
| // The compiled model for the embedding model. | |
| std::optional<litert::CompiledModel> compiled_model_; | |
| // The input buffer for the embedding model. | |
| std::vector<litert::TensorBuffer> input_buffers_; | |
| // The output buffers for the embedding model. | |
| std::vector<litert::TensorBuffer> output_buffers_; | |
| // The output buffer type for the embedding model. | |
| std::optional<litert::RankedTensorType> output_buffer_type_; | |
| // The size of the output tensor needed for a single token. | |
| size_t floats_per_token_output_; | |
| // The default embedding vector to use when a token is not found in the | |
| // lookup table. This is set to the value of token id 0. | |
| std::vector<float> default_embedding_vector_; | |
| // The signature key to use for the embedding model. If not provided, the | |
| // first signature key will be used. | |
| std::optional<std::string> signature_key_; | |
| }; | |
| } // namespace litert::lm | |