Add limit on tokens length

translate.py

@@ -34,6 +34,9 @@ def translator_fn(input_text: str, k=10) -> TranslationResult:
     input_tokens = tokenizer.batch_decode(inputs.input_ids[0])
     input_special_mask = torch.tensor([1 if t in tokenizer.all_special_tokens else 0 for t in input_tokens]).to(device)
 
+    if len(input_tokens) > model.config.d_model:
+        raise ValueError("Input text is too long")
+
     # Generate output
     outputs = model.generate(**inputs, return_dict_in_generate=True, output_scores=True, output_attentions=True)
     output_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)