StaticLLMPipeline: Handle single element list of prompts (openvinotoo…

…lkit#848)
yatarkan · Sep 10, 2024 · a1fbf86 · a1fbf86
2 parents 05d1d47 + 77ffb2e
commit a1fbf86
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/src/cpp/src/llm_pipeline_static.cpp b/src/cpp/src/llm_pipeline_static.cpp
@@ -277,12 +277,18 @@ DecodedResults StaticLLMPipeline::generate(
     StreamerVariant streamer
 ) {
     GenerationConfig config = (generation_config.has_value()) ? *generation_config : m_generation_config;
-    if (std::holds_alternative<std::vector<std::string>>(inputs)) {
-        OPENVINO_THROW("Currently only batch size=1 is supported");
-    }
 
-    OPENVINO_ASSERT(std::holds_alternative<std::string>(inputs));
-    auto& prompt = std::get<std::string>(inputs);
+    std::string prompt;
+    if (auto input_vector = std::get_if<std::vector<std::string>>(&inputs)) {
+        if (input_vector->size() > 1u) {
+            OPENVINO_THROW("Currently only batch size=1 is supported");
+        }
+        OPENVINO_ASSERT(!input_vector->empty());
+        prompt = std::move(input_vector->front());
+    } else {
+        OPENVINO_ASSERT(std::holds_alternative<std::string>(inputs));
+        prompt = std::get<std::string>(inputs);
+    }
 
     if (m_is_chat_conversation) {
         m_history.push_back({{"role", "user"}, {"content", prompt}});