Repository: OpenNMT/OpenNMT-py

Issue:
Translation server error with 'wu' coverage penalty and replace unk options
Hi,

Using latest version of OpenNMT-py 3.1.1, the following two options result in different errors:

1. Using a coverage penalty, specifically "wu"
2. Setting replace_unk to True


This is a working config with above mentioned options disabled:
```
{
"models_root": "...",
"models": [
  {"id": 1, "model": "model.pt", "load": true, "opt": {"gpu": 0, "beam_size": 5, "length_penalty": "wu", "coverage_penalty": "none", "alpha": 0.2, "beta": 0.2}, "tokenizer": {"type": "sentencepiece", "model": "spm.model"}}
]
}
```


Failing config 1 (coverage_penalty = wu):
```
{
"models_root": "...",
"models": [
  {"id": 1, "model": "model.pt", "load": true, "opt": {"gpu": 0, "beam_size": 5, "length_penalty": "wu", "coverage_penalty": "wu", "alpha": 0.2, "beta": 0.2}, "tokenizer": {"type": "sentencepiece", "model": "spm.model"}}
]
}
```

Error output:
```
[2023-05-08 14:10:17,051 ERROR] Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translation_server.py", line 587, in run
    self.translator._translate(infer_iter)
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 353, in _translate
    batch_data = self.translate_batch(
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 801, in translate_batch
    return self._translate_batch_with_strategy(
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 891, in _translate_batch_with_strategy
    decode_strategy.advance(log_probs, attn)
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/beam_search.py", line 321, in advance
    self.topk_scores -= cov_penalty.view(_B, self.beam_size).float()
RuntimeError: shape '[1, 5]' is invalid for input of size 25
```


Failing config 2 (replace_unk = true):
```
{
"models_root": "...",
"models": [
  {"id": 1, "model": "model.pt", "load": true, "opt": {"gpu": 0, "beam_size": 5, "length_penalty": "wu", "coverage_penalty": "none", "replace_unk": true, "alpha": 0.2, "beta": 0.2}, "tokenizer": {"type": "sentencepiece", "model": "spm.model"}}
]
}
```

Error output:
```
[2023-05-08 14:13:13,926 ERROR] Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translation_server.py", line 587, in run
    self.translator._translate(infer_iter)
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 353, in _translate
    batch_data = self.translate_batch(
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 801, in translate_batch
    return self._translate_batch_with_strategy(
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 894, in _translate_batch_with_strategy
    decode_strategy.update_finished()
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/beam_search.py", line 166, in update_finished
    self.alive_attn.view(
RuntimeError: shape '[3, 1, 5, 6]' is invalid for input of size 450
```

Thanks.


Discussion:
do you have a model config to share so that I can try ?
Thanks for your response @vince62s.

The models are trained on proprietary data, so unfortunately I can't share them. 
We tried with models trained with V3 and V2 models converted using the convertv2_v3.py tool, in both cases with either coverage penalty "wu" or the replace_unk option, we got the errors I reported above. 


I did not ask you your data but the model config, don't think it would be an issue. Otherwise can't help you further.
OK, just to clarify then, are you asking for the config used to train the model or the actual model? I can provide the config used to train the model, but I cannot provide the actual model itself.  
I am not native :) but "model config" means the config of the model, right ?
Indeed it does, I've heard multiple people refer to models as model configs before, so wasn't sure which definition you were using ;)

Here's the config used to train our models: 
```
src_seq_length: 200
tgt_seq_length: 200

keep_checkpoint: 20
early_stopping: 10
save_checkpoint_steps: 5000
seed: 1234
report_every: 100
train_steps: 1000000
warmup_steps: 8000
valid_steps: 5000

decoder_type: transformer
encoder_type: transformer
word_vec_size: 512
rnn_size: 512
layers: 6
transformer_ff: 2048
heads: 8

valid_batch_size: 16
batch_size_multiple: 1
accum_count: 3
accum_steps: 0

optim: adam
adam_beta1: 0.9
adam_beta2: 0.998
decay_method: noam
learning_rate: 2.0
max_grad_norm: 0.0

batch_size: 4096
batch_type: tokens
normalization: tokens
dropout: 0.1
label_smoothing: 0.1

max_generator_batches: 2

param_init: 0.0
param_init_glorot: 'true'
position_encoding: 'true'
```
so it's a regular vanilla transformer.
replace_unk won't work, searach previous issues / forum for explanation

for the coverage_penalty: what batch_size are you using when you got the error ? (post the command line / inference config you are using)
The configs I posted in the description were using the default batch size and batch type i.e. I didn't specify a batch_size or batch_type. 

Below are outputs with different batch_type and batch_size, the error in all cases is the same. 

Command used to start the server: 
`python3 -m onmt.bin.server --port 8080 --ip 0.0.0.0 --config models.conf `

Curl command: 
`curl <host>:<port>/translator/translate -X POST -d '[{"id": 1, "src": "Where are my cats today"}]' -H 'content-type: application/json'`

Model config 1:
```
{
"models_root": "/data/training/",
"models": [
  {"id": 1, "model": "model_v3.pt", "load": true, "opt": {"gpu": 0, "beam_size": 5, "length_penalty": "wu", "coverage_penalty": "wu", "alpha": 0.2, "beta": 0.2, "batch_size": 32, "batch_type": "tokens"}, "tokenizer": {"type": "sentencepiece", "model": "spm.model"}}
]
}
```

Output with model conf 1:
```
[2023-05-19 10:39:24,741 ERROR] Traceback (most recent call last):
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translation_server.py", line 587, in run
    self.translator._translate(infer_iter)
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 353, in _translate
    batch_data = self.translate_batch(
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 801, in translate_batch
    return self._translate_batch_with_strategy(
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 891, in _translate_batch_with_strategy
    decode_strategy.advance(log_probs, attn)
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/beam_search.py", line 321, in advance
    self.topk_scores -= cov_penalty.view(_B, self.beam_size).float()
RuntimeError: shape '[1, 5]' is invalid for input of size 25
```

Model conf 2 (without wu coverage penalty, same batch size and batch type):
```
{
"models_root": "/data/training/",
"models": [
  {"id": 1, "model": "model_v3.pt", "load": true, "opt": {"gpu": 0, "beam_size": 5, "length_penalty": "wu", "coverage_penalty": "none", "alpha": 0.2, "beta": 0.2, "batch_size": 32, "batch_type": "tokens"}, "tokenizer": {"type": "sentencepiece", "model": "spm.model"}}
]
}
```

Output with model conf 2:
```
[[{"n_best":1,"pred_score":-2.1375200748443604,"src":"Where are my cats today","tgt":"Wo sind meine Katzen heute"}]]
```

Model config 3, using `sents` batch type, and `batch_size` 1 and `wu` coverage penalty:
```
{
"models_root": "/data/training/",
"models": [
  {"id": 1, "model": "model_v3.pt", "load": true, "opt": {"gpu": 0, "beam_size": 5, "length_penalty": "wu", "coverage_penalty": "wu", "alpha": 0.2, "beta": 0.2, "batch_size": 1, "batch_type": "sents"}, "tokenizer": {"type": "sentencepiece", "model": "spm.model"}}
]
}
```

Output model config3:
```
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translation_server.py", line 587, in run
    self.translator._translate(infer_iter)
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 353, in _translate
    batch_data = self.translate_batch(
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 801, in translate_batch
    return self._translate_batch_with_strategy(
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/translator.py", line 891, in _translate_batch_with_strategy
    decode_strategy.advance(log_probs, attn)
  File "/usr/local/lib/python3.8/site-packages/onmt/translate/beam_search.py", line 321, in advance
    self.topk_scores -= cov_penalty.view(_B, self.beam_size).float()
RuntimeError: shape '[1, 5]' is invalid for input of size 25
```


The issue occurs with a `beam_size > 1`, just tried the below config and it worked. 

Model config: 
```
{
"models_root": "/data/training/",
"models": [
  {"id": 1, "model": "model_v3.pt", "load": true, "opt": {"gpu": 0, "beam_size": 1, "length_penalty": "wu", "coverage_penalty": "wu", "alpha": 0.2, "beta": 0.2, "batch_size": 32, "batch_type": "tokens"}, "tokenizer": {"type": "sentencepiece", "model": "spm.model"}}
]
}
```

Output:
```
[[{"n_best":1,"pred_score":-2.816187858581543,"src":"Where are my cats today","tgt":"Wo sind meine Katzen heute"}]]
```

Also, the issue affects coverage penalty in general when `beam_size > 1`; I tried with coverage penalty of `summary` and `beam_size=5` and got the same error. 

Resolving patch:
diff --git a/onmt/translate/beam_search.py b/onmt/translate/beam_search.py
--- a/onmt/translate/beam_search.py
+++ b/onmt/translate/beam_search.py
@@ -202,7 +202,7 @@ def update_finished(self):
         predictions = self.alive_seq.view(_B_old, self.beam_size, step)
         attention = (
             self.alive_attn.view(
-                step - 1, _B_old, self.beam_size, self.alive_attn.size(-1)
+                _B_old, self.beam_size, step - 1, self.alive_attn.size(-1)
             )
             if self.alive_attn is not None
             else None
@@ -221,7 +221,7 @@ def update_finished(self):
                     (
                         self.topk_scores[i, j],
                         predictions[i, j, 1:],  # Ignore start_token.
-                        attention[:, i, j, : self.src_len[i]]
+                        attention[i, j, :, : self.src_len[i]]
                         if attention is not None
                         else None,
                     )
@@ -275,14 +275,14 @@ def remove_finished_batches(
         self.maybe_update_target_prefix(self.select_indices)
         if self.alive_attn is not None:
             inp_seq_len = self.alive_attn.size(-1)
-            self.alive_attn = attention.index_select(1, non_finished).view(
-                step - 1, _B_new * self.beam_size, inp_seq_len
+            self.alive_attn = attention.index_select(0, non_finished).view(
+                _B_new * self.beam_size, step - 1, inp_seq_len
             )
             if self._cov_pen:
                 self._coverage = (
-                    self._coverage.view(1, _B_old, self.beam_size, inp_seq_len)
-                    .index_select(1, non_finished)
-                    .view(1, _B_new * self.beam_size, inp_seq_len)
+                    self._coverage.view(_B_old, self.beam_size, 1, inp_seq_len)
+                    .index_select(0, non_finished)
+                    .view(_B_new * self.beam_size, 1, inp_seq_len)
                 )
                 if self._stepwise_cov_pen:
                     self._prev_penalty = self._prev_penalty.index_select(
@@ -345,7 +345,7 @@ def advance(self, log_probs, attn):
         self.maybe_update_forbidden_tokens()
 
         if self.return_attention or self._cov_pen:
-            current_attn = attn.index_select(1, self.select_indices)
+            current_attn = attn.index_select(0, self.select_indices)
             if step == 1:
                 self.alive_attn = current_attn
                 # update global state (step == 1)
@@ -353,11 +353,11 @@ def advance(self, log_probs, attn):
                     self._prev_penalty = torch.zeros_like(self.topk_log_probs)
                     self._coverage = current_attn
             else:
-                self.alive_attn = self.alive_attn.index_select(1, self.select_indices)
-                self.alive_attn = torch.cat([self.alive_attn, current_attn], 0)
+                self.alive_attn = self.alive_attn.index_select(0, self.select_indices)
+                self.alive_attn = torch.cat([self.alive_attn, current_attn], 1)
                 # update global state (step > 1)
                 if self._cov_pen:
-                    self._coverage = self._coverage.index_select(1, self.select_indices)
+                    self._coverage = self._coverage.index_select(0, self.select_indices)
                     self._coverage += current_attn
                     self._prev_penalty = self.global_scorer.cov_penalty(
                         self._coverage, beta=self.global_scorer.beta


Test patch:
diff --git a/onmt/tests/test_beam_search.py b/onmt/tests/test_beam_search.py
--- a/onmt/tests/test_beam_search.py
+++ b/onmt/tests/test_beam_search.py
@@ -62,7 +62,7 @@ def test_advance_with_all_repeats_gets_blocked(self):
                 word_probs = torch.full((batch_sz * beam_sz, n_words), -float("inf"))
                 word_probs[0::beam_sz, repeat_idx] = 0
 
-                attns = torch.randn(1, batch_sz * beam_sz, 53)
+                attns = torch.randn(batch_sz * beam_sz, 1, 53)
                 beam.advance(word_probs, attns)
 
                 if i < ngram_repeat:
@@ -133,7 +133,7 @@ def test_advance_with_some_repeats_gets_blocked(self):
                     word_probs[0::beam_sz, repeat_idx] = 0
                     # continue pushing around what beam 1 predicts
                     word_probs[1::beam_sz, repeat_idx + i + 1] = 0
-                attns = torch.randn(1, batch_sz * beam_sz, 53)
+                attns = torch.randn(batch_sz * beam_sz, 1, 53)
                 beam.advance(word_probs, attns)
                 if i < ngram_repeat:
                     self.assertFalse(
@@ -207,7 +207,7 @@ def test_repeating_excluded_index_does_not_die(self):
                     word_probs[1::beam_sz, repeat_idx + i + 1] = 0
                     # predict the allowed-repeat again in beam 2
                     word_probs[2::beam_sz, repeat_idx_ignored] = 0
-                attns = torch.randn(1, batch_sz * beam_sz, 53)
+                attns = torch.randn(batch_sz * beam_sz, 1, 53)
                 beam.advance(word_probs, attns)
                 if i < ngram_repeat:
                     self.assertFalse(
@@ -293,7 +293,7 @@ def test_doesnt_predict_eos_if_shorter_than_min_len(self):
                         beam_idx = min(beam_sz - 1, k)
                         word_probs[beam_idx::beam_sz, j] = score
 
-                attns = torch.randn(1, batch_sz * beam_sz, 53)
+                attns = torch.randn(batch_sz * beam_sz, 1, 53)
                 all_attns.append(attns)
                 beam.advance(word_probs, attns)
                 if i < min_length:
@@ -364,7 +364,7 @@ def test_beam_is_done_when_X_beams_eos_using_min_length(self):
                 for j in range(beam_sz):
                     word_probs[j::beam_sz, eos_idx] = valid_score_dist[0]
 
-            attns = torch.randn(1, batch_sz * beam_sz, 53)
+            attns = torch.randn(batch_sz * beam_sz, 1, 53)
             beam.advance(word_probs, attns)
             if i < min_length:
                 self.assertFalse(beam.done)
@@ -435,7 +435,7 @@ def test_beam_returns_attn_with_correct_length(self):
                 for j in range(beam_sz):
                     word_probs[j::beam_sz, eos_idx] = valid_score_dist[0]
 
-            attns = torch.randn(1, batch_sz * beam_sz, 53)
+            attns = torch.randn(batch_sz * beam_sz, 1, 53)
             beam.advance(word_probs, attns)
             if i < min_length:
                 self.assertFalse(beam.done)
@@ -483,7 +483,7 @@ class TestBeamSearchAgainstReferenceCase(unittest.TestCase):
     INP_SEQ_LEN = 53
 
     def random_attn(self):
-        return torch.randn(1, self.BATCH_SZ * self.BEAM_SZ, self.INP_SEQ_LEN)
+        return torch.randn(self.BATCH_SZ * self.BEAM_SZ, 1, self.INP_SEQ_LEN)
 
     def init_step(self, beam, expected_len_pen):
         # init_preds: [4, 3, 5, 6, 7] - no EOS's