diff --git a/python/tests/test_translator.py b/python/tests/test_translator.py index f76b78c31..ece251437 100644 --- a/python/tests/test_translator.py +++ b/python/tests/test_translator.py @@ -30,6 +30,60 @@ def _get_transliterator(): return ctranslate2.Translator(_get_model_path()) +def _make_scaled_transformer(tmp_dir): + hidden_size = 4 + ffn_size = 8 + vocabulary = ["", "", "", "", "a", "b", "c", "d"] + rng = np.random.default_rng(0) + model = ctranslate2.specs.TransformerSpec.from_config(1, 1) + + def set_linear(linear, output_size, input_size): + linear.weight = rng.normal(size=(output_size, input_size)).astype(np.float32) + + def set_norm(norm): + norm.gamma = np.ones(hidden_size, dtype=np.float32) + norm.beta = np.zeros(hidden_size, dtype=np.float32) + + def set_attention(attention, self_attention): + set_norm(attention.layer_norm) + if self_attention: + set_linear(attention.linear[0], 3 * hidden_size, hidden_size) + set_linear(attention.linear[1], hidden_size, hidden_size) + else: + set_linear(attention.linear[0], hidden_size, hidden_size) + set_linear(attention.linear[1], 2 * hidden_size, hidden_size) + set_linear(attention.linear[2], hidden_size, hidden_size) + + def set_ffn(ffn): + set_norm(ffn.layer_norm) + set_linear(ffn.linear_0, ffn_size, hidden_size) + set_linear(ffn.linear_1, hidden_size, ffn_size) + + model.encoder.embeddings[0].weight = rng.normal( + size=(len(vocabulary), hidden_size) + ).astype(np.float32) + model.decoder.embeddings.weight = rng.normal( + size=(len(vocabulary), hidden_size) + ).astype(np.float32) + model.decoder.scale_embeddings = True + model.decoder.start_from_zero_embedding = True + set_norm(model.encoder.layer_norm) + set_norm(model.decoder.layer_norm) + set_attention(model.encoder.layer[0].self_attention, self_attention=True) + set_ffn(model.encoder.layer[0].ffn) + set_attention(model.decoder.layer[0].self_attention, self_attention=True) + set_attention(model.decoder.layer[0].attention, self_attention=False) + set_ffn(model.decoder.layer[0].ffn) + set_linear(model.decoder.projection, len(vocabulary), hidden_size) + model.register_source_vocabulary(vocabulary) + model.register_target_vocabulary(vocabulary) + model.validate() + + model_dir = str(tmp_dir.join("scaled_transformer").ensure(dir=1)) + model.save(model_dir) + return model_dir, vocabulary + + def test_invalid_model_path(): with pytest.raises(RuntimeError, match="open file"): ctranslate2.Translator("xxx") @@ -620,6 +674,37 @@ def test_return_alternatives(): assert output[0].hypotheses[1] == ["a", "t", "s", "u", "m", "o", "n"] +def test_return_alternatives_with_scaled_prefix(tmp_dir): + # Regression test for https://github.com/OpenNMT/CTranslate2/issues/2014. + model_dir, vocabulary = _make_scaled_transformer(tmp_dir) + translator = ctranslate2.Translator(model_dir) + source = [["a", "b"]] + prefix = ["c", "d"] + + greedy_result = translator.translate_batch( + source, + target_prefix=[prefix], + beam_size=1, + max_decoding_length=len(prefix) + 1, + return_logits_vocab=True, + )[0] + + alternatives_result = translator.translate_batch( + source, + target_prefix=[prefix], + num_hypotheses=1, + return_alternatives=True, + max_decoding_length=len(prefix) + 1, + return_scores=True, + )[0] + + logits = np.array(greedy_result.logits[0][-1]) + logits[vocabulary.index("")] = -np.inf + max_logit = np.max(logits) + expected_score = max_logit - (max_logit + np.log(np.exp(logits - max_logit).sum())) + assert alternatives_result.scores[0] == pytest.approx(expected_score, abs=1e-5) + + def test_return_alternatives_with_vmap(tmp_dir): model_dir = _get_model_path_with_vmap(tmp_dir, ["z", "s", "u", "m", "o", "n"]) translator = ctranslate2.Translator(model_dir) diff --git a/src/layers/transformer.cc b/src/layers/transformer.cc index 758b2c243..a4336e1e4 100644 --- a/src/layers/transformer.cc +++ b/src/layers/transformer.cc @@ -634,9 +634,11 @@ namespace ctranslate2 { StorageView layer_out(dtype, device); _embeddings(ids, layer_in); + // A sequence starting at step 0 only skips preprocessing on its first embedding. + const dim_t embedding_step = (is_sequence && step == 0 && ids.dim(-1) > 1) ? -1 : step; if (_start_from_zero_embedding) - zero_first_timestep(layer_in, step); - if (_embeddings_scale && (!_start_from_zero_embedding || step != 0)) + zero_first_timestep(layer_in, embedding_step); + if (_embeddings_scale && (!_start_from_zero_embedding || embedding_step != 0)) ops::Mul()(layer_in, *_embeddings_scale, layer_in); if (_project_in) { (*_project_in)(layer_in, layer_out);