Lenylvt commited on
Commit
2189552
·
verified ·
1 Parent(s): 4b6b297

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -34
app.py CHANGED
@@ -1,47 +1,55 @@
1
  import gradio as gr
2
  from transformers import MarianMTModel, MarianTokenizer
3
 
4
- # Specify the model name from the Hugging Face Hub, for example, an English to French model by the University of Helsinki
5
- model_name = "Helsinki-NLP/opus-mt-en-fr"
6
-
7
- # Load the tokenizer and model
8
- tokenizer = MarianTokenizer.from_pretrained(model_name)
9
- model = MarianMTModel.from_pretrained(model_name)
10
-
11
- # Function to handle translation
12
- def translate_text(text, target_language):
13
- # Adjust the model_name based on the target language
14
- # Note: You'd need to find the exact model names for each language pair you want to support
15
- model_name_map = {
16
- "French": "Helsinki-NLP/opus-mt-en-fr",
17
- "German": "Helsinki-NLP/opus-mt-en-de",
18
- "Spanish": "Helsinki-NLP/opus-mt-en-es",
19
- }
20
-
21
- selected_model_name = model_name_map.get(target_language, "Helsinki-NLP/opus-mt-en-fr")
22
-
23
- # Load the selected model and tokenizer
24
- tokenizer = MarianTokenizer.from_pretrained(selected_model_name)
25
- model = MarianMTModel.from_pretrained(selected_model_name)
26
-
27
- # Prepare the text for translation
28
- encoded_text = tokenizer.prepare_seq2seq_batch([text], return_tensors="pt")
29
-
30
- # Perform the translation
31
- translated = model.generate(**encoded_text)
32
-
33
- # Decode the translated text
34
  translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
35
 
36
  return translated_text
37
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38
  # Define the interface
39
  iface = gr.Interface(
40
  fn=translate_text,
41
- inputs=[gr.Textbox(lines=2, placeholder="Enter text to translate..."), gr.Dropdown(["French", "German", "Spanish"], label="Select Language")],
42
- outputs=[gr.Textbox()],
43
- title="Text Translator with Helsinki NLP Models",
44
- description="Select a language to translate English text into using University of Helsinki models."
45
  )
 
46
  # Launch the app
47
  iface.launch()
 
1
  import gradio as gr
2
  from transformers import MarianMTModel, MarianTokenizer
3
 
4
+ # Function to dynamically load the model and tokenizer based on selected languages
5
+ def translate_text(text, source_language, target_language):
6
+ # Construct model name based on selected languages
7
+ model_name = f"Helsinki-NLP/opus-mt-{source_language}-{target_language}"
8
+
9
+ # Load tokenizer and model
10
+ try:
11
+ tokenizer = MarianTokenizer.from_pretrained(model_name)
12
+ model = MarianMTModel.from_pretrained(model_name)
13
+ except Exception as e:
14
+ return f"Failed to load model for {source_language} to {target_language}: {str(e)}"
15
+
16
+ # Translate text
17
+ translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
  translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
19
 
20
  return translated_text
21
 
22
+ # Define language options (ISO 639-1 codes and names)
23
+ # Note: This is a simplified subset for demonstration. Expand based on available models.
24
+ language_options = [
25
+ ('en', 'English'),
26
+ ('es', 'Spanish'),
27
+ ('fr', 'French'),
28
+ ('de', 'German'),
29
+ ('zh', 'Chinese'),
30
+ ('ru', 'Russian'),
31
+ ('ar', 'Arabic'),
32
+ ('it', 'Italian'),
33
+ ('pt', 'Portuguese'),
34
+ ('nl', 'Dutch'),
35
+ # Add more languages as needed
36
+ ]
37
+
38
+ # Convert language options to the format expected by the dropdown
39
+ language_dropdown_options = [(code, f"{name} ({code})") for code, name in language_options]
40
+
41
+ # Create dropdowns for source and target languages
42
+ source_language_dropdown = gr.inputs.Dropdown(choices=language_dropdown_options, label="Source Language")
43
+ target_language_dropdown = gr.inputs.Dropdown(choices=language_dropdown_options, label="Target Language")
44
+
45
  # Define the interface
46
  iface = gr.Interface(
47
  fn=translate_text,
48
+ inputs=[gr.inputs.Textbox(lines=2, placeholder="Enter text to translate..."), source_language_dropdown, target_language_dropdown],
49
+ outputs=gr.outputs.Textbox(),
50
+ title="Text Translator with Dynamic Helsinki NLP Models",
51
+ description="Select source and target languages to translate text using Helsinki NLP models."
52
  )
53
+
54
  # Launch the app
55
  iface.launch()