Jesus Carrasco commited on
Commit
bf83ac1
1 Parent(s): 6908aad

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +12 -1
app.py CHANGED
@@ -2,7 +2,18 @@ import gradio as gr
2
  from transformers import MarianMTModel, MarianTokenizer
3
 
4
  def translate(text, target_language):
5
- model_name = f'Helsinki-NLP/opus-mt-en-{target_language}'
 
 
 
 
 
 
 
 
 
 
 
6
  tokenizer = MarianTokenizer.from_pretrained(model_name)
7
  model = MarianMTModel.from_pretrained(model_name)
8
  inputs = tokenizer(text, return_tensors="pt")
 
2
  from transformers import MarianMTModel, MarianTokenizer
3
 
4
  def translate(text, target_language):
5
+ language_codes = {
6
+ "Spanish": "es",
7
+ "French": "fr",
8
+ "Italian": "it",
9
+ "Ukrainian": "uk",
10
+ "Brazilian Portuguese": "pt",
11
+ "Japanese": "ja",
12
+ "Russian": "ru",
13
+ "Chinese": "zh"
14
+ }
15
+ target_language_code = language_codes[target_language]
16
+ model_name = f'helsinki-nlp/opus-mt-en-{target_language_code}'
17
  tokenizer = MarianTokenizer.from_pretrained(model_name)
18
  model = MarianMTModel.from_pretrained(model_name)
19
  inputs = tokenizer(text, return_tensors="pt")