# pylint: disable=redefined-outer-name,unused-import import torch from api import ComposerAPI, GenerationRequest from logitbias_json import json2logitbias from sample import SearchException from test_fixtures import ( untrained_model, untrained_cross_model, untrained_embed_nocross_model, test_vocab, ) import pytest from pprint import pprint import json from data import wavtool2midi, rec_sum from util import configure_logging configure_logging() @pytest.fixture() def api(test_vocab, untrained_model): a = ComposerAPI( test_vocab, untrained_model, ) a.inference_batch_size = 6 return a @pytest.fixture() def api_with_cross(test_vocab, untrained_cross_model): a = ComposerAPI( test_vocab, untrained_cross_model, ) a.inference_batch_size = 6 return a @pytest.fixture() def api_with_embed_nocross(test_vocab, untrained_embed_nocross_model): a = ComposerAPI( test_vocab, untrained_embed_nocross_model, ) a.inference_batch_size = 4 return a def mock_critic(embeddings): return torch.rand((embeddings.shape[0],)) def test_api_nocrash(api): assert ( len( api.generate([], [GenerationRequest(0.7, "steps == 3", None, None, None)])[ 0 ] ) > 0 ) def test_api_multiple_requests(api): assert ( len( api.generate( [], [GenerationRequest(0.7, "steps == 3", None, None, None)] * 3 )[0] ) > 0 ) def check_max_polyphony(notes, limit): if len(notes) == 0: return note_ends = [] for note in notes: note_ends = [t for t in note_ends if t > note["onBeat"]] note_ends.append(note["offBeat"]) if len(note_ends) > limit: pprint(notes) raise Exception(f"Too many notes at {note['onBeat']}: {note_ends}") # @pytest.mark.parametrize("critic", [None, mock_critic]) @pytest.mark.parametrize("critic", [None]) def test_api_polyphony_limits(api, critic): api.critic = critic api.inference_batch_size = 8 bias = json2logitbias({"length > 0.5": -1000}) assert ( len(api.generate([], [GenerationRequest(0.7, "steps == 3", bias, None, 1)])[0]) == 3 ) assert ( len(api.generate([], [GenerationRequest(0.7, "steps == 10", bias, None, 1)])[0]) == 10 ) check_max_polyphony( api.generate([], [GenerationRequest(0.7, "steps == 10", bias, 2, 2)])[0], 2 ) for res in api.generate( [], [GenerationRequest(0.7, "steps == 10", bias, 2, 4)] * 2 ): check_max_polyphony(res, 4) check_max_polyphony( api.generate( [], [GenerationRequest(0.7, "steps == 10", bias, None, 4)], )[0], 4, ) assert ( len( api.generate( [], [GenerationRequest(0.7, "steps == 10", bias, 4, None)], )[0] ) > 0 ) midis = [ wavtool2midi(json.loads(x)) for x in [ '{"notes":[{"pitch":40,"start":1,"end":1.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":40,"start":1.75,"end":1.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.25,"end":0.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":0.75,"end":0.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":41,"start":1.25,"end":1.375,"velocity":0.7874015748031497,"lifted":true},{"pitch":42,"start":1.5,"end":1.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0,"end":0.125,"velocity":0.7874015748031497,"lifted":true},{"pitch":43,"start":0.5,"end":0.625,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":48,"start":1,"end":1.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":0.875,"end":1.875,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":0.125,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":52,"start":1.125,"end":2,"velocity":0.7874015748031497,"lifted":true},{"pitch":51,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":47,"start":2,"end":2.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":2,"end":3,"velocity":0.7874015748031497,"lifted":true},{"pitch":54,"start":1.875,"end":2.75,"velocity":0.7874015748031497,"lifted":true},{"pitch":39,"start":2,"end":2.125,"velocity":0.7874015748031497,"lifted":true}]}', '{"notes":[{"pitch":48,"start":0,"end":1,"velocity":0.7874015748031497,"lifted":true},{"pitch":49,"start":1,"end":2,"velocity":0.7874015748031497,"lifted":true},{"pitch":50,"start":2,"end":3,"velocity":0.7874015748031497,"lifted":true}]}', ] ] @pytest.mark.parametrize("prefix", [[], midis[0]]) @pytest.mark.parametrize("accompany", [None, [], midis[1]]) def test_api_accompany(api, prefix, accompany): # can't have these tests generating off the end bias = json2logitbias({"length > 1.0": -100}) res = api.generate( prefix, [ GenerationRequest(0.7, "steps == 2", bias, None, 2), GenerationRequest(0.7, "steps == 4", bias, None, 2), ], accompany=accompany, ) assert len(res[0]) >= 2 and len(res[0]) <= 4 assert len(res[1]) >= 4 and len(res[1]) <= 8 @pytest.mark.parametrize("prefix", [[], midis[0]]) @pytest.mark.parametrize("accompany", [None, [], midis[1]]) def test_api_wrap_around(api, prefix, accompany): with pytest.raises(SearchException): bias = json2logitbias({"length < 1.0 || length > 2.0": -100}) api.generate( prefix, [ GenerationRequest(0.7, "steps == 32", bias, None, 1), ], accompany=accompany, ) # @pytest.mark.parametrize("critic", [None, mock_critic]) @pytest.mark.parametrize("critic", [None]) def test_api_requests_independent(api, critic): api.critic = critic api.inference_batch_size = 8 # can't have these tests generating off the end bias = json2logitbias({"length > 1.0": -100}) res = api.generate( [], [ GenerationRequest(0.0, "steps == 16", bias, 1, 4), GenerationRequest(0.0, "steps == 16", bias, 1, 4), ], ) assert res[0] == res[1] @pytest.mark.parametrize("prefix", [[], midis[0]]) @pytest.mark.parametrize("text", [None, "hello world"]) # @pytest.mark.parametrize("critic", [None, mock_critic]) @pytest.mark.parametrize("critic", [None]) def test_api_with_cross(api_with_cross, prefix, text, critic): api_with_cross.critic = critic api_with_cross.inference_batch_size = 8 # can't have these tests generating off the end bias = json2logitbias({"length > 1.0": -100}) res = api_with_cross.generate( prefix, [ GenerationRequest(0.7, "steps == 2", bias, None, 2), GenerationRequest(0.7, "steps == 4", bias, None, 2), ], text_prompt=text, ) assert len(res[0]) >= 2 and len(res[0]) <= 4 assert len(res[1]) >= 4 and len(res[1]) <= 8 @pytest.mark.parametrize("prefix", [[], midis[0]]) @pytest.mark.parametrize("text", [None, "hello world"]) # @pytest.mark.parametrize("critic", [None, mock_critic]) @pytest.mark.parametrize("critic", [None]) def test_api_with_embed_nocross(api_with_embed_nocross, prefix, text, critic): api_with_embed_nocross.critic = critic api_with_embed_nocross.inference_batch_size = 8 # can't have these tests generating off the end bias = json2logitbias({"length > 1.0": -100}) res = api_with_embed_nocross.generate( prefix, [ GenerationRequest(0.7, "steps == 2", bias, None, 2), GenerationRequest(0.7, "steps == 4", bias, None, 2), ], text_prompt=text, ) assert len(res[0]) >= 2 and len(res[0]) <= 4 assert len(res[1]) >= 4 and len(res[1]) <= 8 @pytest.mark.parametrize("prefix", [[], midis[0]]) @pytest.mark.parametrize("accompany", [None, midis[1]]) def test_api_dists(api, prefix, accompany): # can't have these tests generating off the end bias = json2logitbias({"length > 1.0": -100}) res = api.generate( prefix, [GenerationRequest(0, "steps == 4", bias, None, 2, dists=True)], accompany=accompany, ) notes = res[0] time = notes[0]["onBeat"] for i, note in enumerate(notes): dists = note["dists"] assert abs(rec_sum(dists["pitch"]) - 1.0) < 1e-2 if note["onBeat"] > time: time = note["onBeat"] arg_max_rest_dur = max( notes[i - 1]["dists"]["restDuration"].items(), key=lambda dp: dp[1] )[0] assert abs(arg_max_rest_dur - (time - notes[i - 1]["onBeat"])) < 1e-8 arg_max_pitch = max( dists["pitch"].items(), key=lambda pp: max(pp[1].values()) if isinstance(pp[1], dict) else pp[1], )[0] arg_max_duration = max(dists["duration"].items(), key=lambda dp: dp[1])[0] assert arg_max_pitch == note["note"] assert abs(arg_max_duration - (note["offBeat"] - note["onBeat"])) < 1e-8 @pytest.mark.parametrize("prefix", [[], midis[0]]) @pytest.mark.parametrize("accompany", [None, midis[1]]) def test_api_bias_lengths(api, prefix, accompany): bias = json2logitbias({"length > 0.1": -1000}) ress = api.generate( prefix, [ GenerationRequest(0.7, "steps == 10", bias, None, None), GenerationRequest(0.7, "steps == 15", bias, None, None), ], accompany=accompany, ) assert len(ress) == 2 for res in ress: assert len(res) > 0 time = res[0]["onBeat"] for note in res: assert note["offBeat"] - note["onBeat"] <= 0.1 assert note["onBeat"] - time <= 0.1 time = note["offBeat"] @pytest.mark.parametrize("prefix", [[], midis[0]]) @pytest.mark.parametrize("accompany", [None, midis[1]]) def test_api_start_offset(api, prefix, accompany): ress = api.generate( prefix, [ GenerationRequest( 0.7, "steps == 4", json2logitbias( { "start && !rest": -1000, "index == 0 && length != 4.0 && rest": -1000, # necessary to prevent running off the end w/ untrained model "index == 0 && length > 1.0 && !rest": -100, "index > 0 && length > 1.0": -100, } ), 1, 1, ), ], accompany=accompany, ) last_start = 0 if not prefix else max(note["onBeat"] for note in prefix) assert len(ress) == 1 res = ress[0] assert res[0]["onBeat"] == last_start + 4.0 assert res[0]["offBeat"] - res[0]["onBeat"] <= 1.0 assert res[1]["onBeat"] - res[0]["onBeat"] <= 1.0 @pytest.mark.parametrize("prefix", [[], midis[0]]) @pytest.mark.parametrize("use_pitch", [True, False]) # @pytest.mark.parametrize("critic", [None, mock_critic]) @pytest.mark.parametrize("critic", [None]) def test_api_distinct(api, prefix, use_pitch, critic): api.critic = critic api.inference_batch_size = 8 rarg = { "redundant_pitch_penalty" if use_pitch else "redundant_duration_penalty": 100.0 } bias = json2logitbias({"length > 2.0": -100}) ress = api.generate( prefix, [ GenerationRequest(0.0, "steps == 3", bias, None, 1), GenerationRequest(0.0, "steps == 3", bias, None, 1), GenerationRequest(0.0, "steps == 3", bias, None, 1), ], ) assert ress[0] == ress[1] ress = api.generate( prefix, [ GenerationRequest(0.0, "steps == 3", bias, None, 1), GenerationRequest(0.0, "steps == 3", bias, None, 1), GenerationRequest(0.0, "steps == 3", bias, None, 1), ], **rarg, ) assert ress[0] != ress[1] and ress[0] != ress[2] and ress[1] != ress[2] rc = [] for _ in range(5): ress = api.generate( prefix, [ GenerationRequest(0.0, "steps == 3", bias, None, 1), ], redundant_clips=rc, **rarg, ) assert all(ress[0] != c for c in rc) rc.append(ress[0]) @pytest.mark.parametrize("accompany", [None, midis[1]]) def test_api_bias_lengths_rests(api, accompany): bias = json2logitbias( {"rest && length > 1.2": -100, "!rest && (length < 1 || length > 1.2)": -100} ) ress = api.generate( [], [ GenerationRequest(0.7, "steps == 5", bias, None, 2), GenerationRequest(0.7, "steps == 7", bias, None, 3), ], accompany=accompany, ) assert len(ress) == 2 for res in ress: assert len(res) > 0 time = res[0]["onBeat"] for note in res: assert 0.99 <= note["offBeat"] - note["onBeat"] <= 1.201 assert note["onBeat"] - time <= 1.2 time = note["offBeat"] def test_api_bias_start_rest(api): bias_start_rest = json2logitbias({"length > 1": -100, "start && !rest": -1000}) bias_start_note = json2logitbias({"length > 1": -100, "start && rest": -1000}) ress = api.generate( [], [ GenerationRequest(0.7, "steps == 2", bias_start_rest, None, 2), GenerationRequest(0.7, "steps == 2", bias_start_note, None, 2), GenerationRequest(0.7, "steps == 2", bias_start_rest, 2, 4), GenerationRequest(0.7, "steps == 2", bias_start_note, 2, 4), ], ) assert len(ress) == 4 for res in (ress[0], ress[2]): assert all(note["onBeat"] > 0 for note in res) for res in (ress[1], ress[3]): assert not all(note["onBeat"] > 0 for note in res) @pytest.mark.parametrize("accompany", [None, midis[1]]) def test_api_bias_pitch(api, accompany): bias = json2logitbias( { "length > 1": -10, "index == 0 && pitch > 1": -10, "index == 1 && (pitch <= pitches[-1] || pitch > 20)": -10, "index == 2 && (pitch <= pitches[-1] || pitch > 30)": -10, } ) ress = api.generate( [], 1 * [GenerationRequest(0.7, "steps == 3", bias, None, 1)], accompany=accompany, ) assert len(ress) == 1 for res in ress: assert len(res) == 3 assert res[0]["note"] == 1 assert res[1]["note"] > 1 and res[1]["note"] <= 20 assert res[2]["note"] <= 30 assert res[1]["note"] < res[2]["note"] @pytest.mark.parametrize("prefix", [[], midis[1]]) def test_api_max_len(api, prefix): bias = json2logitbias({"length > 0.1": -1000}) api.generate(prefix, [GenerationRequest(0, "false", bias, None, None)]) @pytest.mark.parametrize("accompany", [None, midis[1]]) def test_api_stop_at_beat(api, accompany): bias = json2logitbias({"length > 0.5": -1000}) one, two, zero = api.generate( [], [ GenerationRequest(0.7, "beat > 1", bias, None, None), GenerationRequest(0.7, "beat > 2", bias, None, None), GenerationRequest(0.7, "beat < 2", bias, None, None), ], accompany=accompany, ) assert one[-1]["onBeat"] <= 1.0 assert two[-1]["onBeat"] <= 2.0 assert len(zero) == 0 def test_api_stop_on_error(api): ress = api.generate( [], [ GenerationRequest( 0.7, "steps == 3 || error", json2logitbias({"index > 0": -1000}), None, 1, ), GenerationRequest( 0.7, "steps == 3 || error", json2logitbias({"index >= 0": -1000}), None, 1, ), ], ) assert len(ress[0]) == 1 assert len(ress[1]) == 0 with pytest.raises(SearchException): api.generate( [], [ GenerationRequest( 0.7, "steps == 3", json2logitbias({"index > 0": -1000}), None, 1 ), ], ) def test_api_backtracking(api): api.backtrack_limit = 100 ress = api.generate( [], [ GenerationRequest( 0.7, "steps == 6", json2logitbias({"index > 0 && pitches[-1] % 2 == 0": -1000}), None, 1, ), GenerationRequest( 0.7, "steps == 3", None, None, 1, ), ], ) assert len(ress[0]) == 6 assert len(ress[1]) == 3 def test_api_backtracking_distinct(api): api.backtrack_limit = 100 for i in range(20): ress = api.generate( [], [ GenerationRequest( 0.9, "steps == 3", json2logitbias({"index > 0 && pitches[-1] % 2 == 0": -1000}), None, 1, ), GenerationRequest( 0.9, "steps == 3", json2logitbias({"pitch % 2 == 0": -1000}), None, 1, ), ], redundant_pitch_penalty=100.0, ) # this doesn't reveal much aside from not crashing assert ress[0] != ress[1]