anzorq commited on
Commit
9634fea
·
1 Parent(s): fe57abb
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -12,6 +12,8 @@ model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
12
  def translate(text, num_beams=4, num_return_sequences=4):
13
  inputs = tokenizer(text, return_tensors="pt")
14
 
 
 
15
  translated_tokens = model.generate(
16
  **inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], num_beams=num_beams, num_return_sequences=num_return_sequences
17
  )
 
12
  def translate(text, num_beams=4, num_return_sequences=4):
13
  inputs = tokenizer(text, return_tensors="pt")
14
 
15
+ num_return_sequences = min(num_return_sequences, num_beams)
16
+
17
  translated_tokens = model.generate(
18
  **inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang], num_beams=num_beams, num_return_sequences=num_return_sequences
19
  )