Spaces:
Running
Running
File size: 5,869 Bytes
5f923cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | // 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.
#include "runtime/components/constrained_decoding/constrained_decoder.h"
#include <cstdint>
#include <limits>
#include "absl/status/status.h" // from @com_google_absl
#include "absl/types/span.h" // from @com_google_absl
#include "litert/cc/litert_element_type.h" // from @litert
#include "litert/cc/litert_layout.h" // from @litert
#include "litert/cc/litert_macros.h" // from @litert
#include "litert/cc/litert_model.h" // from @litert
#include "litert/cc/litert_tensor_buffer.h" // from @litert
#include "runtime/util/convert_tensor_buffer.h"
#include "runtime/util/status_macros.h" //NOLINT
#include "tflite/types/half.h" // from @litert
namespace litert::lm {
absl::Status ConstrainedDecoder::UpdateConstraintState(
const ::litert::TensorBuffer& next_token_ids) {
LITERT_ASSIGN_OR_RETURN(auto next_token_ids_span,
ReferTensorBufferAsSpan<int>(next_token_ids));
return UpdateConstraintState(next_token_ids_span);
}
absl::Status ConstrainedDecoder::UpdateConstraintState(
absl::Span<int> next_token_ids) {
RET_CHECK_EQ(next_token_ids.size(), batch_size_)
<< "Batch size [" << next_token_ids.size()
<< "] does not match the expected batch size [" << batch_size_ << "].";
for (int i = 0; i < batch_size_; ++i) {
auto& constraint_state = constraint_states_[i];
ASSIGN_OR_RETURN(
constraint_state,
constraint_->ComputeNext(*constraint_state, next_token_ids[i]));
if (constraint_->IsEnded(*constraint_state)) {
constraint_state = constraint_->Start();
}
}
return absl::OkStatus();
}
absl::Status ConstrainedDecoder::MaskLogits(::litert::TensorBuffer& logits) {
// Compute the allowed tokens bitmap for the current constraint state.
LITERT_ASSIGN_OR_RETURN(auto logits_tensor_type, logits.TensorType());
if (logits_tensor_type.ElementType() == ::litert::ElementType::Float32) {
LITERT_ASSIGN_OR_RETURN(auto logits_span,
ReferTensorBufferAsSpan<float>(logits));
return MaskLogits(logits_span, logits_tensor_type.Layout().Dimensions());
} else if (logits_tensor_type.ElementType() ==
::litert::ElementType::Float16) {
LITERT_ASSIGN_OR_RETURN(auto logits_span,
ReferTensorBufferAsSpan<tflite::half>(logits));
return MaskLogits(logits_span, logits_tensor_type.Layout().Dimensions());
}
return absl::InvalidArgumentError("Unsupported logits type for MaskLogits.");
}
absl::Status ConstrainedDecoder::MaskLogits(
absl::Span<float> logits,
absl::Span<const ::litert::Layout::Dim> logits_dims) {
RET_CHECK_EQ(logits_dims.size(), 3)
<< "Only support logits with dimensions [batch_size, 1, vocab_size].";
int batch_size = logits_dims[0];
int sequence_length = logits_dims[1];
int vocab_size = logits_dims[2];
RET_CHECK_EQ(sequence_length, 1) << "Only support sequence length 1.";
// It is possible that the constraint vocabulary size is larger than the model
// vocabulary size. The remaining tokens in the constraint vocabulary are
// treated as unused tokens.
RET_CHECK_LE(vocab_size, constraint_->GetVocabularySize())
<< "Vocabulary size [" << vocab_size
<< "] does not match the expected vocabulary size ["
<< constraint_->GetVocabularySize() << "].";
RET_CHECK_EQ(batch_size, batch_size_)
<< "Batch size [" << batch_size
<< "] does not match the expected batch size [" << batch_size_ << "].";
for (int b = 0; b < batch_size; ++b) {
auto& constraint_state = constraint_states_[b];
ASSIGN_OR_RETURN(auto bitmap,
constraint_->ComputeBitmap(*constraint_state));
for (int i = 0; i < vocab_size; ++i) {
if (!bitmap->Get(i)) {
logits.data()[b * vocab_size + i] =
std::numeric_limits<float>::lowest();
}
}
}
return absl::OkStatus();
}
absl::Status ConstrainedDecoder::MaskLogits(
absl::Span<tflite::half> logits,
absl::Span<const ::litert::Layout::Dim> logits_dims) {
RET_CHECK_EQ(logits_dims.size(), 3)
<< "Only support logits with dimensions [batch_size, 1, vocab_size].";
int batch_size = logits_dims[0];
int sequence_length = logits_dims[1];
int vocab_size = logits_dims[2];
RET_CHECK_EQ(sequence_length, 1) << "Only support sequence length 1.";
// It is possible that the constraint vocabulary size is larger than the model
// vocabulary size. The remaining tokens in the constraint vocabulary are
// treated as unused tokens.
RET_CHECK_LE(vocab_size, constraint_->GetVocabularySize())
<< "Vocabulary size [" << vocab_size
<< "] does not match the expected vocabulary size ["
<< constraint_->GetVocabularySize() << "].";
RET_CHECK_EQ(batch_size, batch_size_)
<< "Batch size [" << batch_size
<< "] does not match the expected batch size [" << batch_size_ << "].";
for (int b = 0; b < batch_size; ++b) {
auto& constraint_state = constraint_states_[b];
ASSIGN_OR_RETURN(auto bitmap,
constraint_->ComputeBitmap(*constraint_state));
for (int i = 0; i < vocab_size; ++i) {
if (!bitmap->Get(i)) {
logits.data()[b * vocab_size + i] = tflite::half::min();
}
}
}
return absl::OkStatus();
}
} // namespace litert::lm
|