MNT Add unit test for text encoder (#23)

.gitignore

@@ -6,6 +6,7 @@ bark
 encodec
 main
 tests/test-tokenizer
+tests/test-text-encoder
 
 *.o
 *.plist

Makefile

@@ -1,8 +1,8 @@
 # Define the default target now so that it is always the first target
-BUILD_TARGETS = main 
+BUILD_TARGETS = main
 
 # Binaries only useful for tests
-TEST_TARGETS = tests/test-tokenizer
+TEST_TARGETS = tests/test-tokenizer tests/test-text-encoder
 
 default: $(BUILD_TARGETS)
 
@@ -315,3 +315,6 @@ tests: $(TEST_TARGETS)
 
 tests/test-tokenizer: tests/test-tokenizer.cpp ggml.o bark.o encodec.o $(OBJS)
  $(CXX) $(CXXFLAGS) $(filter-out %.txt,$^) -o $@ $(LDFLAGS)
+
+tests/test-text-encoder: tests/test-text-encoder.cpp ggml.o bark.o encodec.o $(OBJS)
+ $(CXX) $(CXXFLAGS) $(filter-out %.txt,$^) -o $@ $(LDFLAGS)

bark.cpp

@@ -1160,7 +1160,6 @@ bool gpt_eval(
  if (mem_per_token == 0) {
  mem_per_token = ggml_used_mem(ctx0)/N;
  }
- //printf("used_mem = %zu\n", ggml_used_mem(ctx0));
 
  ggml_free(ctx0);
 
@@ -1215,6 +1214,60 @@ bark_vocab::id gpt_sample(
  return logits_id[idx].second;
 }
 
+std::vector<bark_vocab::id> bark_forward_text_encoder(
+ const std::vector<bark_vocab::id> & tokens,
+ const gpt_model model,
+ const int n_threads,
+ const float temp,
+ const bool early_stop,
+ const float min_eos_p) {
+
+ std::vector<bark_vocab::id> out;
+ int n_past = 0;
+ float eos_p = 0;
+
+ std::vector<bark_vocab::id> input = tokens;
+ std::vector<float> logits;
+
+ std::mt19937 rng(0);
+
+ // dry run to estimate mem_per_token
+ size_t mem_per_token = 0;
+ gpt_eval(model, n_threads, 0, false, { 0, 1, 2, 3 }, logits, mem_per_token);
+
+ for (int i = 0; i < 768; i++) {
+ const bool merge_ctx = i == 0;
+ gpt_eval(model, n_threads, n_past, merge_ctx, input, logits, mem_per_token);
+
+ float logits_pad_token = logits[SEMANTIC_PAD_TOKEN];
+ logits.resize(SEMANTIC_VOCAB_SIZE);
+
+ if (early_stop)
+ logits.push_back(logits[logits_pad_token]);
+
+ n_past += input.size();
+ if (i == 0)
+ n_past -= 256; // first step, context are merged
+
+ input.clear();
+
+ bark_vocab::id sampled_id = gpt_sample(logits, temp, rng, &eos_p);
+
+ if (early_stop && ((sampled_id == SEMANTIC_VOCAB_SIZE) || (eos_p > min_eos_p)))
+ break;
+
+ input.push_back(sampled_id);
+ out.push_back(sampled_id);
+
+ printf("%d ", sampled_id);
+ fflush(stdout);
+ }
+
+ printf("\n\nsemantic sequence length: %zu\n\n", out.size());
+
+ return out;
+}
+
 bool bark_generate_audio(
  bark_model model,
  const bark_vocab& vocab,
@@ -1227,17 +1280,15 @@ bool bark_generate_audio(
  const int seed = 0;
 
  // const float top_p = 0.2;
- const float temp = 0.7;
+ const float temp = 1.0;
  const float fine_temp = 0.5;
 
- const int early_stop = true;
+ const bool early_stop = true;
 
  const int sliding_window_size = 60;
  const int max_coarse_history = 630;
 
- // in the original implementation, min_eos_p=0.2, yet for bark.cpp this seems too
- // high and this generates overly long sequence.
- const float min_eos_p = 0.15;
+ const float min_eos_p = 0.2;
 
  std::mt19937 rng(seed);
 
@@ -1279,47 +1330,8 @@ bool bark_generate_audio(
  printf("\n\n");
 
  // encode text (text model)
- std::vector<bark_vocab::id> out_semantic;
- {
- int n_past = 0;
- float eos_p = 0;
-
- std::vector<bark_vocab::id> input = tokens;
- std::vector<float> logits;
-
- // dry run to estimate mem_per_token
- size_t mem_per_token = 0;
- gpt_eval(model.text_model, n_threads, 0, false, { 0, 1, 2, 3 }, logits, mem_per_token);
-
- for (int i = 0; i < 768; i++) {
- const bool merge_ctx = i == 0;
- gpt_eval(model.text_model, n_threads, n_past, merge_ctx, input, logits, mem_per_token);
-
- float logits_pad_token = logits[SEMANTIC_PAD_TOKEN];
- logits.resize(SEMANTIC_VOCAB_SIZE);
-
- if (early_stop)
- logits.push_back(logits[logits_pad_token]);
-
- n_past += input.size();
- if (i == 0)
- n_past -= 256; // first step, context are merged
-
- input.clear();
-
- bark_vocab::id sampled_id = gpt_sample(logits, temp, rng, &eos_p);
- input.push_back(sampled_id);
- out_semantic.push_back(sampled_id);
-
- printf("%d ", sampled_id);
- fflush(stdout);
-
- if (early_stop && ((sampled_id == SEMANTIC_VOCAB_SIZE) || (eos_p > min_eos_p)))
- break;
- }
-
- printf("\n\nsemantic sequence length: %zu\n\n", out_semantic.size());
- }
+ std::vector<bark_vocab::id> out_semantic = bark_forward_text_encoder(
+ tokens, model.text_model, n_threads, temp, early_stop, min_eos_p);
 
  // coarse encoding (coarse model)
  std::vector<std::vector<bark_vocab::id>> out_coarse(N_COARSE_CODEBOOKS);

bark.h

@@ -1,6 +1,7 @@
 #include "encodec.h"
 
 #include <map>
+#include <random>
 #include <vector>
 
 #define CLS_TOKEN_ID 101
@@ -114,6 +115,21 @@ struct bark_model {
 
 bool gpt_model_load(const std::string& fname, gpt_model& model);
 
+bool gpt_eval(
+ const gpt_model & model,
+ const int n_threads,
+ const int n_past,
+ const bool merge_ctx,
+ const std::vector<bark_vocab::id> & embd_inp,
+ std::vector<float> & embd_w,
+ size_t & mem_per_token);
+
+bark_vocab::id gpt_sample(
+ const std::vector<float>& logits,
+ double temp,
+ std::mt19937 & rng,
+ float * eos_p);
+
 bool bark_model_load(const std::string & dirname, bark_model & model);
 
 bool bark_vocab_load(const std::string& fname, bark_vocab& vocab, int32_t expected_size);
@@ -130,3 +146,11 @@ bool bark_generate_audio(
  const bark_vocab& vocab,
  const char * text,
  const int n_threads);
+
+std::vector<bark_vocab::id> bark_forward_text_encoder(
+ const std::vector<bark_vocab::id> & tokens,
+ const gpt_model model,
+ const int n_threads,
+ const float temp,
+ const bool early_stop,
+ const float min_eos_p);

examples/main.cpp

@@ -25,7 +25,7 @@ int main() {
  printf("\n");
 
  // forward pass
- const std::string prompt = "This is an audio";
+ const std::string prompt = "hello world";
  {
  const int64_t t_eval_us_start = ggml_time_us();
 

tests/CMakeLists.txt

@@ -7,3 +7,4 @@ function(bark_add_test source)
 endfunction()
 
 bark_add_test(test-tokenizer.cpp ${CMAKE_CURRENT_SOURCE_DIR}/../ggml_weights/ggml_vocab.bin)
+bark_add_test(test-text-encoder.cpp ${CMAKE_CURRENT_SOURCE_DIR}/../ggml_weights/ggml_weights_text.bin)

tests/test-text-encoder.cpp

@@ -0,0 +1,92 @@
+#include "bark.h"
+
+#include <cstdio>
+#include <string>
+#include <map>
+#include <random>
+#include <vector>
+
+static const std::map<std::vector<bark_vocab::id>, std::vector<bark_vocab::id>> & k_tests()
+{
+ static const std::vector<bark_vocab::id> seq1 = { 71742, 20181, 21404 };
+ static const std::vector<bark_vocab::id> seq2 = { 20579, 20172, 20199, 33733 };
+ static const std::vector<bark_vocab::id> seq3 = { 21113, 35307, 10165, 62571, 10165, 23622, 20236, 20959, 52867 };
+
+ static const std::vector<bark_vocab::id> ans1 = { 3264, 6121, 6414, 7799, 6121, 1907, 1888, 206, 1888, 1888, 6143, 2131, 728, 6328, 3393, 5990, 2992, 8837, 206, 7799, 2533, 7374, 2992, 4059, 5990, 5990, 5236, 5939, 10, 2137, 2137, 6021, 176, 176, 2584, 499, 2382, 499, 8051, 4218, 8051, 8909, 306, 6804, 6804, 292, 292, 675, 3927, 3819, 624, 6583, 8843, 6583, 6583, 9628, 3062, 6652, 6652, 6652, 4792, 6608, 4792, 6608, 9025, 8534, 709, 7431, 7431, 9693, 8858, 8858, 3820, 8858, 682, 4076, 8996, 4909, 5682, 6139, 6139, 9133, 445, 971, 7542, 4564, 9931, 9931, 785, 785, 157, 5897, 5897, 9527, 1233, 138, 131, 10, 266, 266, 1572, 1572, 206, 206, 3533, 206, 4874, 1444, 3533, 206, 206, 7397, 206, 3252, 206, 2314, 91, 206, 7567, 841, 5346, 3252, 206, 841, 3366, 517, 517, 3252, 344, 344, 1278, 3950, 57, 57, 597, 7160, 121, 7334, 631, 292, 41, 41, 8944, 1991, 1408, 1408, 1408, 1462, 3, 166, 8745, 17, 2332, 1574, 7443, 50, 17, 27, 429, 9225, 713, 4099, 4099, 4099, 75, 555, 5932, 8870, 7627, 7627, 5661, 3088, 26, 288, 262 };
+ static const std::vector<bark_vocab::id> ans2 = { 3205, 6179, 7731, 6972, 5722, 602, 441, 125, 147, 991, 1573, 402, 402, 6774, 1913, 8020, 8572, 8572, 1722, 5681, 1133, 4694, 1133, 7517, 9575, 8125, 5905, 6486, 1797, 6486, 5138, 5138, 4150, 2630, 2879, 59, 28, 28, 385, 1741, 4042, 9898, 302, 9600, 7231, 5673, 5475, 321, 171, 321, 164, 1025, 4681, 6202, 6752, 8288, 6747, 7656, 9804, 9804, 2411, 178, 50, 441, 6401, 5899, 79, 6511, 6511, 9629, 6511, 6154, 2224, 2224, 73, 73, 9814, 6303, 1997, 1997, 7396, 8062, 825, 441};
+ static const std::vector<bark_vocab::id> ans3 = { 3946, 8514, 7741, 9262, 5153, 4400, 4509, 512, 1136, 4631, 8486, 4631, 3954, 7234, 993, 4412, 993, 9161, 332, 8209, 5565, 4224, 4344, 6152, 6152, 2704, 2285, 4438, 232, 131, 10, 5038, 2430, 59, 28, 28, 28, 28, 1310, 4449, 5920, 9449, 2002, 9693, 7939, 4049, 4049, 6059, 210, 100, 10, 10, 282, 3968, 988, 9790, 1728, 2587, 4405, 2948, 232, 232, 100, 3621, 8680, 417, 10, 2595, 7352, 9725, 6445, 2428, 41, 41, 10, 41, 9261, 4212, 3963, 6261, 8210, 9588, 934, 441, 1025, 2875, 8558, 6968, 116, 41, 41, 7789, 5721, 5721, 267, 2116, 579, 100, 1133, 3446, 2599, 7503, 3390, 3390, 4485, 657, 1385, 1385, 7691, 7557, 5272, 8887, 10, 6619, 3592, 6394, 5272, 5272, 8887, 1841, 602, 441, 217, 4542, 5861, 5861, 3803, 4542, 4542, 4542, 6675, 7204, 131, 100, 790, 2832, 266, 6115, 4209, 1739, 1739, 1444, 8659, 1739, 1739, 1739, 1133, 1739, 1739, 2556, 2556, 413, 413, 10, 3373, 7966, 2330, 1588, 409, 2942, 59, 28, 28, 28, 10, 28, 3160, 9569, 5920, 5887, 9693, 6290, 3458, 1242, 50, 210, 2977, 1433, 1433, 6150, 6150, 1136, 6413, 9693, 3441, 9598, 9061, 7949, 9137, 5615, 131, 100, 652, 7863, 7344, 8899, 7765, 50, 10, 100, 7399, 9915, 7557, 4509, 8486, 6264, 6133, 6133, 6133, 6619, 5210, 5210, 9629, 2555, 2339, 9486, 1425, 2762, 2466, 1079, 10 };
+
+ static std::map<std::vector<bark_vocab::id>, std::vector<bark_vocab::id>> _k_tests = {
+ { seq1, ans1 }, // hello world
+ { seq2, ans2 }, // this is an audio
+ { seq3, ans3 }, // You cannot, sir, take from me anything
+ };
+ return _k_tests;
+};
+
+std::vector<bark_vocab::id> pad_input(const std::vector<bark_vocab::id> & input) {
+ int original_sz = input.size();
+ std::vector<bark_vocab::id> pad(input.begin(), input.end());
+ pad.resize(513);
+
+ for (int i = original_sz; i < 256; i++)
+ pad[i] = TEXT_PAD_TOKEN;
+ for (int i = 256; i < 512; i++)
+ pad[i] = SEMANTIC_PAD_TOKEN;
+ pad[512] = SEMANTIC_INFER_TOKEN;
+
+ return pad;
+}
+
+int main(int argc, char** argv) {
+ if (argc < 2) {
+ fprintf(stderr, "Usage: %s <model-file>\n", argv[0]);
+ return 1;
+ }
+
+ const std::string fname = argv[1];
+
+ gpt_model model;
+ const int n_threads = 4;
+ const float min_eos_p = 0.2;
+
+ printf("%s: reading bark text model\n", __func__);
+ if(!gpt_model_load(fname, model)) {
+ fprintf(stderr, "%s: invalid model file '%s'\n", __func__, fname.c_str());
+ return 1;
+ }
+
+ for (const auto & test_kv : k_tests()) {
+ std::vector<bark_vocab::id> pad = pad_input(test_kv.first);
+ std::vector<bark_vocab::id> res = bark_forward_text_encoder(
+ pad, model, n_threads, 1.0f, true, min_eos_p);
+
+ bool correct = res.size() == test_kv.second.size();
+
+ for (int i = 0; i < (int) res.size() && correct; ++i) {
+ if (res[i] != test_kv.second[i]) {
+ correct = false;
+ }
+ }
+
+ if (!correct) {
+ fprintf(stderr, "%s : failed test \n", __func__);
+ fprintf(stderr, "%s : expected tokens (n=%zu): ", __func__, test_kv.second.size());
+ for (const auto & t : test_kv.second) {
+ fprintf(stderr, "%d ", t);
+ }
+ fprintf(stderr, "\n");
+ fprintf(stderr, "%s : got tokens (n=%zu): ", __func__, res.size());
+ for (const auto & t : res) {
+ fprintf(stderr, "%d ", t);
+ }
+ fprintf(stderr, "\n");
+
+ return 3;
+ }
+ }
+
+ fprintf(stderr, "%s : tests passed successfully.\n", __func__);
+
+ return 0;
+}