implement fix for set inputs

microsoft · Dec 16, 2024 · 3fc4ba1 · 3fc4ba1
1 parent 7735e10
commit 3fc4ba1
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/src/generators.cpp b/src/generators.cpp
@@ -262,7 +262,7 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_
 
   // Temporary solution for multimodal and whisper models
   if (!params.aux_input_ids.empty() && params.aux_input_ids.data() != nullptr) {
-    AppendTokens(params.aux_input_ids);
+    AuxAppendTokens(params.aux_input_ids);
   }
 }
 
@@ -274,6 +274,20 @@ DeviceSpan<int32_t> Generator::AllocateInputIdsOnDevice(const cpu_span<int32_t>
   return input_ids_device;
 }
 
+// TODO(aciddelgado): Remove this function once SetInputs is moved to generator
+void Generator::AuxAppendTokens(const cpu_span<int32_t> input_ids) {
+  ThrowErrorIfSessionTerminated(state_->session_terminated_);
+  if (input_ids.size() == 0)
+    throw std::runtime_error("input_ids is empty");
+  if (search_->GetSequenceLength() != 0 && state_->params_->search.batch_size > 1)
+    throw std::runtime_error("AppendTokens can only be called once for batch_size > 1. To call AppendTokens again, use RewindToLength(0)");
+
+  auto input_ids_device = AllocateInputIdsOnDevice(input_ids);
+  search_->AppendTokens(input_ids_device);
+  computed_logits_ = false;
+  ComputeLogits(input_ids_device);
+}
+
 void Generator::AppendTokens(const cpu_span<int32_t> input_ids) {
   ThrowErrorIfSessionTerminated(state_->session_terminated_);
   if (input_ids.size() == 0)

diff --git a/src/generators.h b/src/generators.h
@@ -128,6 +128,7 @@ struct Generator : LeakChecked<Generator> {
 
  private:
   DeviceSpan<int32_t> AllocateInputIdsOnDevice(const cpu_span<int32_t> input_ids);
+  void AuxAppendTokens(const cpu_span<int32_t> input_ids);
   void ComputeLogits(DeviceSpan<int32_t> next_tokens);
   enum Action { standard,   // Default, set in any other case
                 generated,  // Set after GenerateNextToken