-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_prelude_forcing.cpp
More file actions
85 lines (77 loc) · 2.42 KB
/
Copy pathtest_prelude_forcing.cpp
File metadata and controls
85 lines (77 loc) · 2.42 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
#include <cstdio>
#include <memory>
#include "encode_helpers.h"
#include "events.h"
#include "lm.h"
#include "mt3.h"
#include "pipeline.h"
static int g_fails = 0;
#define CHECK(cond) do { if (!(cond)) { std::fprintf(stderr, "FAIL %s:%d %s\n", __FILE__, __LINE__, #cond); ++g_fails; } } while(0)
int main() {
using namespace muscriptor;
MT3Tokenizer tok;
// Tie section encoding
{
auto ids = tok.tie_section_token_ids({{0, 60}, {0, 64}});
CHECK(ids.size() == 4); // program, pitch, pitch, tie
CHECK(tok.vocab()[ids[0]].type == "program");
CHECK(tok.vocab()[ids.back()].type == "tie");
}
// OpenNoteTracker open_keys after note on
{
auto idx = encode_index_map(1001);
OpenNoteTracker tracker(tok.vocab());
tracker.feed(ChunkBoundary{0.f, std::nullopt});
tracker.feed(idx[{"tie", 0}]);
tracker.feed(idx[{"program", 0}]);
tracker.feed(idx[{"velocity", 1}]);
tracker.feed(idx[{"pitch", 60}]);
auto keys = tracker.open_keys();
CHECK(keys.size() == 1);
CHECK(keys[0].first == 0 && keys[0].second == 60);
}
// batch_size + prelude_forcing rule
{
LMConfig cfg;
cfg.card = 1393;
cfg.dim = 32;
cfg.num_heads = 4;
cfg.num_layers = 1;
auto lm = std::make_shared<LMModel>(LMModel::create_random(cfg, 1));
TranscriptionModel tm(lm, tok);
CHECK(tm.resolve_batch_size(0, true) == 1);
bool threw = false;
try {
tm.resolve_batch_size(4, true);
} catch (...) {
threw = true;
}
CHECK(threw);
CHECK(tm.resolve_batch_size(4, false) == 4);
}
// Tiny LM generate with prompt (prelude)
{
LMConfig cfg;
cfg.card = 1393;
cfg.dim = 32;
cfg.num_heads = 4;
cfg.num_layers = 1;
auto lm = LMModel::create_random(cfg, 7);
auto prompt = tok.tie_section_token_ids({{0, 60}});
GenerateConfig g;
g.max_gen_len = 16;
g.prompt = prompt;
g.early_stop_on_token = tok.eos_id();
auto out = lm.generate({}, g);
CHECK(out.size() >= prompt.size());
for (size_t i = 0; i < prompt.size(); ++i) {
CHECK(out[i] == prompt[i]);
}
}
if (g_fails) {
std::fprintf(stderr, "%d failures\n", g_fails);
return 1;
}
std::printf("ok prelude_forcing\n");
return 0;
}