Skip to content

Commit

Permalink
tests
Browse files Browse the repository at this point in the history
  • Loading branch information
iefode committed Sep 27, 2024
1 parent 5bfa2e7 commit 9f44fe4
Show file tree
Hide file tree
Showing 3 changed files with 257 additions and 11 deletions.
1 change: 1 addition & 0 deletions src/cpp/src/sampler.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ class Sampler {
Tokenizer m_tokenizer;

public:
Sampler() = default;
Sampler(Tokenizer & tokenizer) : m_tokenizer(tokenizer) {};

SamplerOutput sample(std::vector<SequenceGroup::Ptr> & sequence_groups, ov::Tensor logits, bool is_validation_mode_enabled);
Expand Down
3 changes: 2 additions & 1 deletion src/cpp/src/sequence_group.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -428,7 +428,8 @@ class SequenceGroup {
}

void set_num_validated_tokens(size_t k) {
m_num_validated_tokens = k;
// in case of non-prompt we need to take prev tokens + token to validate
m_num_validated_tokens = get_num_processed_tokens() ? k + 1 : k;
}

size_t get_num_available_tokens_for_batching() const {
Expand Down
264 changes: 254 additions & 10 deletions tests/cpp/sampler.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -32,30 +32,274 @@ TEST(SamplerStopTokenIdsTest, multiple_stop_sequence_no_match) {
ASSERT_FALSE(is_stop_token_id_hit(generated_tokens.back(), stop_token_ids));
}

TEST(SamplerValidationMode, check_validation_mode) {
TEST(SamplerValidationMode, gen_phase_to_cut_whole_seq) {
auto sampling_config = ov::genai::greedy();
// create sequence group with prompt [0, 1, 2, 3, 4]
std::vector<int64_t> input_vector{0, 1, 2, 3, 4};
ov::Tensor input_ids(ov::element::i64, ov::Shape{1, 5}, input_vector.data());
ov::Tensor input_tensor(ov::element::i64, ov::Shape{1, 5}, input_vector.data());
std::vector<SequenceGroup::Ptr> sequence_groups{
SequenceGroup::Ptr(new SequenceGroup(0, input_ids, sampling_config, 32, false)),
SequenceGroup::Ptr(new SequenceGroup(0, input_tensor, sampling_config, 32, false)),
};
sequence_groups.front()->get_sequences().front()->append_token(0, 1.f);
sequence_groups.front()->get_sequences().front()->append_token(1, 1.f);
sequence_groups.front()->get_sequences().front()->append_token(2, 1.f);
sequence_groups.front()->get_sequences().front()->append_token(2, 1.f);

// to emulate processed prompt and add next token [ 0 ]
sequence_groups.front()->get_sequences().front()->append_token(0, 1.f);
sequence_groups.front()->update_processed_tokens_num(5);

// append candidates [ 2, 3, 4 ]
size_t num_validated_tokens = 3;
for (size_t i = 1; i <= num_validated_tokens; ++i) {
sequence_groups.front()->get_sequences().front()->append_token(i + 1, 1.f);
}

// generated sequence [0, 1, 2, 3, 4] -> [0, 2, 3, 4]
sequence_groups.front()->set_num_validated_tokens(num_validated_tokens);
const auto num_scheduled_tokens = sequence_groups.front()->get_num_available_tokens_for_batching();
ASSERT_EQ(num_scheduled_tokens, num_validated_tokens + 1);
sequence_groups.front()->schedule_tokens(num_scheduled_tokens);

// create ref tensor : to generate candidates + next token
std::vector<float> logits = {
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
0, 0, 0, 0, 1.f,
};

// shape 4 tokens + 1 batch + 5 vocab
ov::Tensor gen_input_ids(ov::element::f32, ov::Shape{4, 1, 5}, logits.data());

// Tokenizer t("");
Sampler sampler;
sampler.sample(sequence_groups, gen_input_ids, true);

TokenIds actual = sequence_groups.front()->get_sequences().front()->get_generated_ids(),
expected{0, 1};
ASSERT_EQ(sequence_groups.front()->get_sequences().front()->get_generated_ids(), expected);
}

TEST(SamplerValidationMode, gen_phase_to_cut_part_seq) {
auto sampling_config = ov::genai::greedy();
// create sequence group with prompt [0, 1, 2, 3, 4]
std::vector<int64_t> input_vector{0, 1, 2, 3, 4};
ov::Tensor input_tensor(ov::element::i64, ov::Shape{1, 5}, input_vector.data());
std::vector<SequenceGroup::Ptr> sequence_groups{
SequenceGroup::Ptr(new SequenceGroup(0, input_tensor, sampling_config, 32, false)),
};

// to emulate processed prompt and add next token [ 0 ]
sequence_groups.front()->get_sequences().front()->append_token(0, 1.f);
sequence_groups.front()->update_processed_tokens_num(5);
sequence_groups.front()->set_num_validated_tokens(3);

// append candidates [ 1, 2, 2 ]
size_t num_validated_tokens = 3;
for (size_t i = 1; i <= num_validated_tokens; ++i) {
int64_t token_id = i == num_validated_tokens ? i - 1 : i;
sequence_groups.front()->get_sequences().front()->append_token(token_id, 1.f);
}

// generated sequence [0, 1, 2, 3, 4] -> [0, 1, 2, 2]
sequence_groups.front()->set_num_validated_tokens(num_validated_tokens);
const auto num_scheduled_tokens = sequence_groups.front()->get_num_available_tokens_for_batching();
ASSERT_EQ(num_scheduled_tokens, num_validated_tokens + 1);
sequence_groups.front()->schedule_tokens(num_scheduled_tokens);

// create ref tensor : to generate candidates + next token
std::vector<float> logits = {
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
0, 0, 0, 0, 1.f,
};

// shape 4 tokens + 1 batch + 5 vocab
ov::Tensor gen_input_ids(ov::element::f32, ov::Shape{4, 1, 5}, logits.data());

Tokenizer t("");
Sampler sampler(t);
// Tokenizer t("");
Sampler sampler;
sampler.sample(sequence_groups, gen_input_ids, true);

TokenIds actual = sequence_groups.front()->get_sequences().front()->get_generated_ids(),
expected{0, 1, 2, 3};
ASSERT_EQ(sequence_groups.front()->get_sequences().front()->get_generated_ids(), expected);
}

TEST(SamplerValidationMode, gen_phase) {
auto sampling_config = ov::genai::greedy();
// create sequence group with prompt [0, 1, 2, 3, 4]
std::vector<int64_t> input_vector{0, 1, 2, 3, 4};
ov::Tensor input_tensor(ov::element::i64, ov::Shape{1, 5}, input_vector.data());
std::vector<SequenceGroup::Ptr> sequence_groups{
SequenceGroup::Ptr(new SequenceGroup(0, input_tensor, sampling_config, 32, false)),
};

// to emulate processed prompt and add next token [ 0 ]
sequence_groups.front()->get_sequences().front()->append_token(0, 1.f);
sequence_groups.front()->update_processed_tokens_num(5);

// append candidates [ 1, 2, 3 ]
size_t num_validated_tokens = 3;
for (size_t i = 1; i <= num_validated_tokens; ++i) {
sequence_groups.front()->get_sequences().front()->append_token(i, 1.f);
}

// generated sequence [0, 1, 2, 3, 4] -> [0, 1, 2, 3]
sequence_groups.front()->set_num_validated_tokens(num_validated_tokens);
const auto num_scheduled_tokens = sequence_groups.front()->get_num_available_tokens_for_batching();
ASSERT_EQ(num_scheduled_tokens, num_validated_tokens + 1);
sequence_groups.front()->schedule_tokens(num_scheduled_tokens);

// create ref tensor : to generate candidates + next token
std::vector<float> logits = {
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
0, 0, 0, 0, 1.f,
};

// shape 4 tokens + 1 batch + 5 vocab
ov::Tensor gen_input_ids(ov::element::f32, ov::Shape{4, 1, 5}, logits.data());

// Tokenizer t("");
Sampler sampler;
sampler.sample(sequence_groups, gen_input_ids, true);

TokenIds actual = sequence_groups.front()->get_sequences().front()->get_generated_ids(),
expected{0, 1, 2, 3, 4};
ASSERT_EQ(sequence_groups.front()->get_sequences().front()->get_generated_ids(), expected);
}

TEST(SamplerValidationMode, prompt_phase_to_cut_part_seq) {
auto sampling_config = ov::genai::greedy();
// create sequence group with prompt [0, 1, 2, 3, 4]
std::vector<int64_t> input_vector{0, 1, 2, 3, 4};
ov::Tensor input_tensor(ov::element::i64, ov::Shape{1, 5}, input_vector.data());
std::vector<SequenceGroup::Ptr> sequence_groups{
SequenceGroup::Ptr(new SequenceGroup(0, input_tensor, sampling_config, 32, false)),
};

// append candidates [ 0, 1, 1 ]
size_t num_validated_tokens = 3;
for (size_t i = 0; i < num_validated_tokens; ++i) {
int64_t token_id = i + 1 == num_validated_tokens ? i - 1 : i;
sequence_groups.front()->get_sequences().front()->append_token(token_id, 1.f);
}

// generated sequence [0, 1, 2, 3, 4] -> [0, 1, 1]
sequence_groups.front()->set_num_validated_tokens(num_validated_tokens);
const auto num_scheduled_tokens = sequence_groups.front()->get_num_available_tokens_for_batching();
// prompt len + validation
ASSERT_EQ(num_scheduled_tokens, num_validated_tokens + input_vector.size());
sequence_groups.front()->schedule_tokens(num_scheduled_tokens);

// create ref tensor : to generate candidates + next token
std::vector<float> logits = {
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
0, 0, 0, 0, 1.f,
1.f, 0, 0, 0, 0,
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
};

// shape 4 tokens + 1 batch + 5 vocab
ov::Tensor gen_input_ids(ov::element::f32, ov::Shape{8, 1, 5}, logits.data());

// Tokenizer t("");
Sampler sampler;
sampler.sample(sequence_groups, gen_input_ids, true);

TokenIds actual = sequence_groups.front()->get_sequences().front()->get_generated_ids(),
expected{0, 1, 2};
ASSERT_EQ(sequence_groups.front()->get_sequences().front()->get_generated_ids(), expected);
}

TEST(SamplerValidationMode, prompt_phase_to_cut_whole_seq) {
auto sampling_config = ov::genai::greedy();
// create sequence group with prompt [0, 1, 2, 3, 4]
std::vector<int64_t> input_vector{0, 1, 2, 3, 4};
ov::Tensor input_tensor(ov::element::i64, ov::Shape{1, 5}, input_vector.data());
std::vector<SequenceGroup::Ptr> sequence_groups{
SequenceGroup::Ptr(new SequenceGroup(0, input_tensor, sampling_config, 32, false)),
};

// append candidates [ 1, 2, 3 ]
size_t num_validated_tokens = 3;
for (size_t i = 0; i < num_validated_tokens; ++i) {
sequence_groups.front()->get_sequences().front()->append_token(i + 1, 1.f);
}

// generated sequence [0, 1, 2, 3, 4] -> [1, 2, 3]
sequence_groups.front()->set_num_validated_tokens(num_validated_tokens);
const auto num_scheduled_tokens = sequence_groups.front()->get_num_available_tokens_for_batching();
// prompt len + validation
ASSERT_EQ(num_scheduled_tokens, num_validated_tokens + input_vector.size());
sequence_groups.front()->schedule_tokens(num_scheduled_tokens);

// create ref tensor : to generate candidates + next token
std::vector<float> logits = {
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
0, 0, 0, 0, 1.f,
1.f, 0, 0, 0, 0,
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
};

// shape 4 tokens + 1 batch + 5 vocab
ov::Tensor gen_input_ids(ov::element::f32, ov::Shape{8, 1, 5}, logits.data());

// Tokenizer t("");
Sampler sampler;
sampler.sample(sequence_groups, gen_input_ids, true);

TokenIds actual = sequence_groups.front()->get_sequences().front()->get_generated_ids(),
expected{0};
ASSERT_EQ(sequence_groups.front()->get_sequences().front()->get_generated_ids(), expected);
}

TEST(SamplerValidationMode, prompt_phase) {
auto sampling_config = ov::genai::greedy();
// create sequence group with prompt [0, 1, 2, 3, 4]
std::vector<int64_t> input_vector{0, 1, 2, 3, 4};
ov::Tensor input_tensor(ov::element::i64, ov::Shape{1, 5}, input_vector.data());
std::vector<SequenceGroup::Ptr> sequence_groups{
SequenceGroup::Ptr(new SequenceGroup(0, input_tensor, sampling_config, 32, false)),
};

// append candidates [ 0, 1, 2 ]
size_t num_validated_tokens = 3;
for (size_t i = 0; i < num_validated_tokens; ++i) {
sequence_groups.front()->get_sequences().front()->append_token(i, 1.f);
}

// generated sequence [0, 1, 2, 3, 4] -> [0, 1, 2]
sequence_groups.front()->set_num_validated_tokens(num_validated_tokens);
const auto num_scheduled_tokens = sequence_groups.front()->get_num_available_tokens_for_batching();
// prompt len + validation
ASSERT_EQ(num_scheduled_tokens, num_validated_tokens + input_vector.size());
sequence_groups.front()->schedule_tokens(num_scheduled_tokens);

// create ref tensor : to generate candidates + next token
std::vector<float> logits = {
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
0, 0, 0, 0, 1.f,
1.f, 0, 0, 0, 0,
0, 1.f, 0, 0, 0,
0, 0, 1.f, 0, 0,
0, 0, 0, 1.f, 0,
};

// shape 4 tokens + 1 batch + 5 vocab
ov::Tensor gen_input_ids(ov::element::f32, ov::Shape{8, 1, 5}, logits.data());

// Tokenizer t("");
Sampler sampler;
sampler.sample(sequence_groups, gen_input_ids, true);

TokenIds actual = sequence_groups.front()->get_sequences().front()->get_generated_ids(),
Expand Down

0 comments on commit 9f44fe4

Please sign in to comment.