TiberiuCristianLeon commited on
Commit
2045817
·
verified ·
1 Parent(s): a103fac

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -29,7 +29,7 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
29
  tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=nllb_language_codes[sselected_language])
30
  model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
31
  translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=nllb_language_codes[sselected_language], tgt_lang=nllb_language_codes[tselected_language])
32
- translated_text = translator(input_text, max_length=360)
33
  return translated_text[0]['translation_text'], message_text
34
  else:
35
  tokenizer = T5Tokenizer.from_pretrained(model_name)
@@ -39,13 +39,15 @@ def translate_text(input_text, sselected_language, tselected_language, model_nam
39
  if torch.cuda.is_available():
40
  model = model.to('cuda')
41
  print("CUDA is available! Using GPU.")
 
 
42
  if model_name.startswith("Helsinki-NLP"):
43
  prompt = input_text
44
  else:
45
  prompt = f"translate {sselected_language} to {tselected_language}: {input_text}"
46
 
47
  input_ids = tokenizer.encode(prompt, return_tensors="pt")
48
- output_ids = model.generate(input_ids, max_length=360)
49
  translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
50
 
51
  print(f'Translating from {sselected_language} to {tselected_language} with {model_name}:', f'{input_text} = {translated_text}', sep='\n')
 
29
  tokenizer = AutoTokenizer.from_pretrained(model_name, src_lang=nllb_language_codes[sselected_language])
30
  model = AutoModelForSeq2SeqLM.from_pretrained(model_name, device_map="auto")
31
  translator = pipeline('translation', model=model, tokenizer=tokenizer, src_lang=nllb_language_codes[sselected_language], tgt_lang=nllb_language_codes[tselected_language])
32
+ translated_text = translator(input_text, max_length=512)
33
  return translated_text[0]['translation_text'], message_text
34
  else:
35
  tokenizer = T5Tokenizer.from_pretrained(model_name)
 
39
  if torch.cuda.is_available():
40
  model = model.to('cuda')
41
  print("CUDA is available! Using GPU.")
42
+ else:
43
+ print("CUDA not available! Using CPU.")
44
  if model_name.startswith("Helsinki-NLP"):
45
  prompt = input_text
46
  else:
47
  prompt = f"translate {sselected_language} to {tselected_language}: {input_text}"
48
 
49
  input_ids = tokenizer.encode(prompt, return_tensors="pt")
50
+ output_ids = model.generate(input_ids, max_length=512)
51
  translated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
52
 
53
  print(f'Translating from {sselected_language} to {tselected_language} with {model_name}:', f'{input_text} = {translated_text}', sep='\n')