Add sampling to vlm pipeline by Sampler

sbalandi · sbalandi · commit 2e3cc746e407 · 2024-10-08T22:55:20.000+01:00
diff --git a/samples/CMakeLists.txt b/samples/CMakeLists.txt
@@ -29,6 +29,8 @@ install(DIRECTORY
             cpp/whisper_speech_recognition
             cpp/stable_diffusion
             cpp/lora_greedy_causal_lm
+        #     cpp/visual_language_chat
+        #     cpp/continuous_batching_accuracy
         DESTINATION samples/cpp COMPONENT cpp_samples_genai)
 
 install(DIRECTORY
diff --git a/samples/cpp/visual_language_chat/visual_language_chat.cpp b/samples/cpp/visual_language_chat/visual_language_chat.cpp
@@ -31,13 +31,20 @@ int main(int argc, char* argv[]) try {
     }
     pipe.generate(
         prompt,
-        ov::genai::image(std::move(image)),
+        // ov::genai::image(std::move(image)),
+        // ov::genai::generation_config(ov::genai::beam_search()),
+        ov::genai::generation_config(ov::genai::greedy()),
+        // ov::genai::generation_config(ov::genai::multinomial()),
         ov::genai::streamer(print_subword)
     );
     std::cout << "\n----------\n"
         "question:\n";
     while (std::getline(std::cin, prompt)) {
-        pipe.generate(prompt, ov::genai::streamer(print_subword));
+        pipe.generate(prompt,
+                    //  ov::genai::generation_config(ov::genai::beam_search()),
+                     ov::genai::generation_config(ov::genai::greedy()),
+                    //  ov::genai::generation_config(ov::genai::multinomial()),
+                      ov::genai::streamer(print_subword));
         std::cout << "\n----------\n"
             "question:\n";
     }
diff --git a/src/cpp/src/sampler.cpp b/src/cpp/src/sampler.cpp
@@ -230,6 +230,20 @@ Sampler::GroupBeamSearcher::GroupBeamSearcher(SequenceGroup::Ptr sequence_group,
     }
 }
 
+std::vector<int32_t> Sampler::GroupBeamSearcher::get_beam_idxs() {
+    std::vector<int32_t> next_beams;
+
+    for (Group& group : m_groups) {
+        if (!group.done) {
+            for (Beam& beam : group.ongoing) {
+                next_beams.push_back(beam.m_global_beam_idx);
+            }
+        }
+    }
+
+    return next_beams;
+}
+
 void Sampler::GroupBeamSearcher::select_next_tokens(const ov::Tensor& logits, SamplerOutput& sampler_output) {
     assert(m_parameters.num_beams % m_parameters.num_beam_groups == 0 &&
         "number of beams should be divisible by number of groups");
@@ -560,6 +574,14 @@ std::vector<int64_t> Sampler::_try_finish_generation(SequenceGroup::Ptr & sequen
     return dropped_seq_ids;
 }
 
+std::vector<int32_t> Sampler::get_beam_idxs(uint64_t request_id) {
+    std::vector<int32_t> beams;
+    if (m_beam_search_info.find(request_id) != m_beam_search_info.end()) {
+        GroupBeamSearcher beam_searcher = m_beam_search_info.at(request_id);
+        std::vector<int32_t> beams = beam_searcher.get_beam_idxs();
+    }
+    return beams;
+}
 
 SamplerOutput Sampler::sample(std::vector<SequenceGroup::Ptr> & sequence_groups, ov::Tensor logits) {
     const float * logits_data = logits.data<float>();
diff --git a/src/cpp/src/sampler.hpp b/src/cpp/src/sampler.hpp
@@ -61,6 +61,8 @@ class Sampler {
     SamplerOutput sample(std::vector<SequenceGroup::Ptr> & sequence_groups, ov::Tensor logits);
     void set_seed(size_t seed) { rng_engine.seed(seed); }
     void clear_beam_search_info(uint64_t request_id);
+
+    std::vector<int32_t> get_beam_idxs(uint64_t request_id);
 };
 
 class Sampler::GroupBeamSearcher {
@@ -105,5 +107,7 @@ class Sampler::GroupBeamSearcher {
 
     void select_next_tokens(const ov::Tensor& logits, SamplerOutput& sampler_output);
     void finalize(SamplerOutput& sampler_output);
+
+    std::vector<int32_t> get_beam_idxs();
 };
 }
diff --git a/src/cpp/src/vlm_pipeline.cpp b/src/cpp/src/vlm_pipeline.cpp
@@ -13,6 +13,10 @@
 #include <optional>
 #include <random>
 
+#include "sampler.hpp"
+
+#include "debug_utils.hpp"
+
 using namespace ov::genai;
 
 namespace {
@@ -416,6 +420,160 @@ ov::Tensor get_image_embedding(const EncodedImage& encoded_image, Tokenizer& tok
 }
 }
 
+void forward_embedings_and_lm(SequenceGroup::CPtr sequence_group, ov::InferRequest& embedding, ov::InferRequest& language, const VLMConfig m_vlm_config, const std::shared_ptr<Sampler> sampler) {
+    // compute aggregated values
+    size_t num_sequences = sequence_group->num_running_seqs();
+    size_t batch_size_in_sequences = num_sequences;
+    size_t total_num_tokens = sequence_group->get_num_scheduled_tokens() * num_sequences;
+    size_t total_num_blocks = sequence_group->get_num_blocks() * num_sequences;
+    size_t max_context_len_val = std::max(max_context_len_val, sequence_group->get_context_len());
+
+    ov::Tensor
+        input_ids(ov::element::i64, {total_num_tokens, 1}),
+        position_ids(ov::element::i64, {total_num_tokens, 1}),
+        beam_idx(ov::element::i32, { total_num_tokens });
+
+    // get raw pointers to copy to
+    int64_t
+        * input_ids_data = input_ids.data<int64_t>(),
+        * position_ids_data = position_ids.data<int64_t>();
+    int32_t
+        * beam_idx_data = beam_idx.data<int32_t>();
+
+    std::vector<Sequence::CPtr> running_sequences = sequence_group->get_running_sequences();
+    size_t num_running_sequences = running_sequences.size();
+    size_t num_scheduled_tokens = sequence_group->get_num_scheduled_tokens();
+    size_t group_position_id = sequence_group->get_num_processed_tokens();
+
+    // spec: In case of multiple input tokens for current sequence (prompt_len > 1),
+    // context_len corresponds to first token within subgroup of scheduled tokens
+    size_t group_context_len = group_position_id;
+
+    for (size_t seq_id = 0; seq_id < num_running_sequences; ++seq_id) {
+        Sequence::CPtr sequence = running_sequences[seq_id];
+
+        for (size_t token_id = 0, position_id = group_position_id; token_id < num_scheduled_tokens; ++token_id, ++position_id) {
+            // compute token for current sequence
+            input_ids_data[token_id] = position_id < sequence_group->get_prompt_len() ?
+                sequence_group->get_prompt_ids()[position_id] :
+                sequence->get_generated_ids()[position_id - sequence_group->get_prompt_len()];
+
+            position_ids_data[token_id] = position_id;
+        }
+
+        // apply strides to shift to a next sequence
+        input_ids_data += num_scheduled_tokens;
+        position_ids_data += num_scheduled_tokens;
+    }
+
+    embedding.set_input_tensor(input_ids);
+
+    embedding.infer();
+    const ov::Tensor& embed_prompt_tensor = embedding.get_output_tensor();
+    float* embed_data = embed_prompt_tensor.data<float>();
+    for (auto idx = 0; idx < embed_prompt_tensor.get_size(); idx++) {
+        embed_data[idx] = embed_data[idx] * m_vlm_config.scale_emb;
+    }
+
+    language.set_tensor("inputs_embeds", embed_prompt_tensor);
+
+    language.get_tensor("attention_mask").set_shape({ total_num_tokens, language.get_tensor("attention_mask").get_shape()[1] + 1 });
+    std::fill_n(language.get_tensor("attention_mask").data<int64_t>(), language.get_tensor("attention_mask").get_size(), 1);
+
+    language.set_tensor("position_ids", position_ids);
+    std::vector<int32_t> beam_idxs = sampler->get_beam_idxs(sequence_group->get_request_id());
+    if (beam_idxs.empty()) {
+        for (size_t i = 0; i < num_sequences; i++) {
+            beam_idx_data[i] = 0;
+        }
+    } else {
+        for (size_t i = 0; i < beam_idxs.size(); i++) {
+            beam_idx_data[i] = beam_idxs.at(i);
+        }
+    }
+    language.set_tensor("beam_idx", beam_idx);
+
+    // print_tensor("input_ids", input_ids);
+    // print_tensor("position_ids", position_ids);
+    // print_tensor("attention_mask", language.get_tensor("attention_mask"));
+    // print_tensor("beam_idx", beam_idx);
+    
+    language.infer();
+}
+
+EncodedGenerationResult get_lm_encoded_results(
+    ov::InferRequest& language,
+    ov::InferRequest& embedding,
+    ov::Tensor inputs_embeds,
+    const VLMConfig m_vlm_config,
+    const std::shared_ptr<StreamerBase> streamer_ptr,
+    const std::shared_ptr<Sampler> sampler,
+    std::vector<SequenceGroup::Ptr> requests
+) {
+    SequenceGroup::Ptr request = requests.back();
+    GenerationHandle generation = std::make_shared<GenerationHandleImpl>(request->get_generation_stream(), request->get_sampling_parameters());
+
+    language.set_tensor("inputs_embeds", inputs_embeds);
+
+    size_t history_len = language.get_tensor("attention_mask").get_shape().at(1);
+    language.get_tensor("attention_mask").set_shape({1, history_len + inputs_embeds.get_shape()[1]});
+    std::fill_n(language.get_tensor("attention_mask").data<int64_t>(), language.get_tensor("attention_mask").get_size(), 1);
+
+    language.get_tensor("position_ids").set_shape({1, inputs_embeds.get_shape().at(1)});
+    std::iota(language.get_tensor("position_ids").data<int64_t>(), language.get_tensor("position_ids").data<int64_t>() + language.get_tensor("position_ids").get_size(), history_len);
+
+    language.get_tensor("beam_idx").set_shape({ BATCH_SIZE });
+    language.get_tensor("beam_idx").data<int32_t>()[0] = 0;
+
+    language.infer();
+
+    int64_t sequence_len = language.get_tensor("logits").get_shape().at(1);
+    request->schedule_tokens(sequence_len);
+
+    SamplerOutput sampler_output = sampler->sample(requests, language.get_tensor("logits"));
+
+    language.get_tensor("inputs_embeds").set_shape({BATCH_SIZE, 1, m_vlm_config.hidden_size});
+    language.get_tensor("position_ids").set_shape({ BATCH_SIZE, 1 });
+
+
+    while (!request->has_finished()) {
+        request->schedule_tokens(1);
+
+        forward_embedings_and_lm(request, embedding, language, m_vlm_config, sampler);
+
+        if (streamer_ptr) {
+            // first sequences
+            int64_t out_token = request.get()->operator[](0)->get_generated_ids().back();
+            if (streamer_ptr->put(out_token)) {
+                break;
+            }
+        }
+
+        sampler_output = sampler->sample(requests, language.get_tensor("logits"));
+    }
+
+    if (streamer_ptr) {
+        streamer_ptr->end();
+    }
+
+    EncodedGenerationResult result;
+    result.m_request_id = 1;
+    std::vector<GenerationOutput> generation_outputs = generation->read_all();
+    std::sort(generation_outputs.begin(), generation_outputs.end(), [=] (GenerationOutput& r1, GenerationOutput& r2) {
+        return r1.score > r2.score;
+    });
+
+    auto num_outputs = std::min(request->get_sampling_parameters().num_return_sequences, generation_outputs.size());
+    for (size_t generation_output_idx = 0; generation_output_idx < num_outputs; ++generation_output_idx) {
+        const auto& generation_output = generation_outputs[generation_output_idx];
+        result.m_generation_ids.push_back(std::move(generation_output.generated_ids));
+        result.m_scores.push_back(generation_output.score);
+    }
+    result.m_status = generation->get_status();
+
+    return result;
+}
+
 VLMPipeline::VLMPipeline(
     const std::filesystem::path& model_dir,
     const Tokenizer& tokenizer,
@@ -473,31 +631,16 @@ DecodedResults VLMPipeline::generate(
         ov::Tensor prompt_tensor = process_prompt(m_embedding, input_ids, m_vlm_config.scale_emb);
         inputs_embeds = concatenate_mid_dim(imgEmbedTensor, prompt_tensor);
     }
-    m_language.set_tensor("inputs_embeds", inputs_embeds);
-    size_t history_len = m_language.get_tensor("attention_mask").get_shape().at(1);
-    m_language.get_tensor("attention_mask").set_shape({1, history_len + inputs_embeds.get_shape()[1]});
-    std::fill_n(m_language.get_tensor("attention_mask").data<int64_t>(), m_language.get_tensor("attention_mask").get_size(), 1);
-    m_language.get_tensor("position_ids").set_shape({1, inputs_embeds.get_shape().at(1)});
-    std::iota(m_language.get_tensor("position_ids").data<int64_t>(), m_language.get_tensor("position_ids").data<int64_t>() + m_language.get_tensor("position_ids").get_size(), history_len);
-    m_language.get_tensor("beam_idx").set_shape({ BATCH_SIZE });
-    m_language.get_tensor("beam_idx").data<int32_t>()[0] = 0;
 
-    m_language.infer();
+    std::shared_ptr<Sampler> sampler = std::make_shared<Sampler>(m_tokenizer);
 
-    ov::Shape logits_shape = m_language.get_tensor("logits").get_shape();
-    auto attention_size = m_language.get_tensor("attention_mask").get_size();
+    std::vector<SequenceGroup::Ptr> requests;
+    // request_id, input_ids, generation_config, block_size, enable_prefix_caching
+    // now we have one prompt as input, so we need one request
+    SequenceGroup::Ptr sequence_group = std::make_shared<SequenceGroup>(0, input_ids, generation_config, 1, false);
+    sequence_group->set_sequence_group_ptr(sequence_group);
+    requests.push_back(sequence_group);
 
-    int64_t sequence_len = m_language.get_tensor("logits").get_shape().at(1) - 1;
-    size_t vocab_size = m_language.get_tensor("logits").get_shape().back();
-    float* logits = m_language.get_tensor("logits").data<float>() + sequence_len * vocab_size;
-    int64_t out_token = std::max_element(logits, logits + vocab_size) - logits;
-
-    m_language.get_tensor("inputs_embeds").set_shape({BATCH_SIZE, 1, m_vlm_config.hidden_size});
-    m_language.get_tensor("position_ids").set_shape({ BATCH_SIZE, 1 });
-
-    m_embedding.get_input_tensor().set_shape({ 1, 1 });
-
-    int64_t eos_token_id = m_tokenizer.get_eos_token_id();
     std::shared_ptr<StreamerBase> streamer_ptr = std::visit(overloaded{
         [&m_tokenizer = m_tokenizer](
             const std::function<bool(std::string)>& callback
@@ -511,47 +654,23 @@ DecodedResults VLMPipeline::generate(
             return std::shared_ptr<StreamerBase>{nullptr};
         },
     }, streamer);
-    std::vector<int64_t> generated;
-    while (true) {  //(out_token != eos_token_id)
-        //out_token embedding
-        m_embedding.get_input_tensor().data<int64_t>()[0] = out_token;
-        m_embedding.infer();
-        const ov::Tensor& embed_prompt_tensor = m_embedding.get_output_tensor();
-        float* embed_data = embed_prompt_tensor.data<float>();
-        for (auto idx = 0; idx < embed_prompt_tensor.get_size(); idx++) {
-            embed_data[idx] = embed_data[idx] * m_vlm_config.scale_emb;
-        }
 
-        m_language.set_tensor("inputs_embeds", embed_prompt_tensor);
-        m_language.get_tensor("attention_mask").set_shape({ BATCH_SIZE, m_language.get_tensor("attention_mask").get_shape()[1] + 1 });
-        std::fill_n(m_language.get_tensor("attention_mask").data<int64_t>(), m_language.get_tensor("attention_mask").get_size(), 1);
-        m_language.get_tensor("position_ids").data<int64_t>()[0] = int64_t(m_language.get_tensor("attention_mask").get_size() - 2);
-
-        m_language.infer();
-
-        generated.push_back(out_token);
-        if (streamer_ptr && streamer_ptr->put(out_token)) {
-            break;
-        }
-        logits = m_language.get_tensor("logits").data<float>();
-
-        out_token = std::max_element(logits, logits + vocab_size) - logits;
-        if (out_token == eos_token_id) {
-            break;
-        }
-    }
-
-    if (streamer_ptr) {
-        streamer_ptr->end();
-    }
+    EncodedGenerationResult encoded_result = get_lm_encoded_results(m_language, m_embedding, inputs_embeds, m_vlm_config, streamer_ptr, sampler, requests);
 
     if (!is_chat_conversation) {
         for (auto& variable : m_language.query_state()) {
             variable.reset();
         }
         m_language.get_tensor("attention_mask").set_shape({1, 0});
     }
-    return {{m_tokenizer.decode(generated)}};
+
+    DecodedResults decoded;
+    for (size_t idx = 0; idx < encoded_result.m_generation_ids.size(); ++idx) {
+        decoded.texts.push_back(m_tokenizer.decode(encoded_result.m_generation_ids.at(idx)));
+        decoded.scores.push_back(encoded_result.m_scores.at(idx));
+    }
+
+    return decoded;
 }
 
 DecodedResults VLMPipeline::generate(
@@ -562,6 +681,15 @@ DecodedResults VLMPipeline::generate(
     ov::genai::OptionalGenerationConfig config_arg = utils::get_config_from_map(config_map);
     GenerationConfig config = (config_arg.has_value()) ? *config_arg : get_generation_config();
     config.update_generation_config(config_map);
+
+    // If eos_token_id was not provided, take value
+    if (config.eos_token_id == -1)
+        config.set_eos_token_id(m_tokenizer.get_eos_token_id());
+
+    // if (is_chat_conversation && config.num_return_sequences > 1) {
+    //     config.num_return_sequences = 1;
+    // }
+
     return generate(
         prompt,
         config_map.end() == image ? std::vector<ov::Tensor>{}
