add m2m model
Browse files
app.py
CHANGED
|
@@ -35,6 +35,9 @@ model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M"
|
|
| 35 |
mbart_m2en_model = EasyNMT("mbart50_m2en")
|
| 36 |
mbart_en2m_model = EasyNMT("mbart50_en2m")
|
| 37 |
|
|
|
|
|
|
|
|
|
|
| 38 |
|
| 39 |
class myTheme(Base):
|
| 40 |
def __init__(
|
|
@@ -169,6 +172,17 @@ def mbart_trans(article, target_language):
|
|
| 169 |
return mbart_m2en_model.translate(article, target_lang="en")
|
| 170 |
else:
|
| 171 |
return mbart_en2m_model.translate(article, target_lang="zh")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 172 |
|
| 173 |
|
| 174 |
def translate(article, toolkit, target_language):
|
|
@@ -178,6 +192,8 @@ def translate(article, toolkit, target_language):
|
|
| 178 |
translated = nllb_trans(article, target_language)
|
| 179 |
elif toolkit == "MBART":
|
| 180 |
translated = mbart_trans(article, target_language)
|
|
|
|
|
|
|
| 181 |
|
| 182 |
return translated
|
| 183 |
|
|
@@ -187,7 +203,7 @@ myTheme = myTheme()
|
|
| 187 |
with gr.Blocks(theme=myTheme) as demo:
|
| 188 |
article = gr.Textbox(label="Article")
|
| 189 |
toolkit_select = gr.Radio(
|
| 190 |
-
["OPUS", "NLLB", "MBART"], label="Select Translation Model", value="OPUS"
|
| 191 |
)
|
| 192 |
lang_select = gr.Radio(["English", "Chinese"], label="Select Desired Language")
|
| 193 |
result = gr.Textbox(label="Translated Result")
|
|
|
|
| 35 |
mbart_m2en_model = EasyNMT("mbart50_m2en")
|
| 36 |
mbart_en2m_model = EasyNMT("mbart50_en2m")
|
| 37 |
|
| 38 |
+
# Initialize m2m_100 models
|
| 39 |
+
m2m_model = EasyNMT("m2m_100_1.2B")
|
| 40 |
+
|
| 41 |
|
| 42 |
class myTheme(Base):
|
| 43 |
def __init__(
|
|
|
|
| 172 |
return mbart_m2en_model.translate(article, target_lang="en")
|
| 173 |
else:
|
| 174 |
return mbart_en2m_model.translate(article, target_lang="zh")
|
| 175 |
+
else:
|
| 176 |
+
return "Error: You chose the same language as the article detected language. Please reselect language and try again."
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def m2m_trans(article, target_language):
|
| 180 |
+
result_lang = detect_lang(article)
|
| 181 |
+
|
| 182 |
+
if result_lang != target_language:
|
| 183 |
+
return m2m_model.translate(article)
|
| 184 |
+
else:
|
| 185 |
+
return "Error: You chose the same language as the article detected language. Please reselect language and try again."
|
| 186 |
|
| 187 |
|
| 188 |
def translate(article, toolkit, target_language):
|
|
|
|
| 192 |
translated = nllb_trans(article, target_language)
|
| 193 |
elif toolkit == "MBART":
|
| 194 |
translated = mbart_trans(article, target_language)
|
| 195 |
+
elif toolkit == "M2M":
|
| 196 |
+
translated = m2m_trans(article, target_language)
|
| 197 |
|
| 198 |
return translated
|
| 199 |
|
|
|
|
| 203 |
with gr.Blocks(theme=myTheme) as demo:
|
| 204 |
article = gr.Textbox(label="Article")
|
| 205 |
toolkit_select = gr.Radio(
|
| 206 |
+
["OPUS", "NLLB", "MBART", "M2M"], label="Select Translation Model", value="OPUS"
|
| 207 |
)
|
| 208 |
lang_select = gr.Radio(["English", "Chinese"], label="Select Desired Language")
|
| 209 |
result = gr.Textbox(label="Translated Result")
|