add language detection

app.py

@@ -2,6 +2,8 @@ import streamlit as st
 import os
 import io
 from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from languages import LANGUANGE_MAP
 import time
 import json
 from typing import List
@@ -135,6 +137,17 @@ def load_model(
     return tokenizer, model
 
 
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def load_detection_model(
+    pretrained_model: str = "ivanlau/language-detection-fine-tuned-on-xlm-roberta-base",
+    cache_dir: str = "models/",
+):
+    tokenizer = AutoTokenizer.from_pretrained(pretrained_model, cache_dir=cache_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(pretrained_model, cache_dir=cache_dir).to(device)
+    model.eval()
+    return tokenizer, model
+
+
 st.title("M2M100 Translator")
 st.write("M2M100 is a multilingual encoder-decoder (seq-to-seq) model trained for Many-to-Many multilingual translation. It was introduced in this paper https://arxiv.org/abs/2010.11125 and first released in https://github.com/pytorch/fairseq/tree/master/examples/m2m_100 repository. The model that can directly translate between the 9,900 directions of 100 languages.\n")
 
@@ -147,26 +160,38 @@ user_input: str = st.text_area(
     max_chars=5120,
 )
 
-source_lang = st.selectbox(label="Source language", options=list(lang_id.keys()))
 target_lang = st.selectbox(label="Target language", options=list(lang_id.keys()))
 
 if st.button("Run"):
     time_start = time.time()
     tokenizer, model = load_model()
+    de_tokenizer, de_model = load_detection_model()
 
-    src_lang = lang_id[source_lang]
-    trg_lang = lang_id[target_lang]
-    tokenizer.src_lang = src_lang
     with torch.no_grad():
-        encoded_input = tokenizer(user_input, return_tensors="pt").to(device)
-        generated_tokens = model.generate(
-            **encoded_input, forced_bos_token_id=tokenizer.get_lang_id(trg_lang)
-        )
-        translated_text = tokenizer.batch_decode(
-            generated_tokens, skip_special_tokens=True
-        )[0]
-
-    time_end = time.time()
-    st.success(translated_text)
-
-    st.write(f"Computation time: {round((time_end-time_start),3)} segs")
+
+        tokenized_sentence = de_tokenizer(user_input, return_tensors='pt')
+        output = de_model(**tokenized_sentence)
+        de_predictions = torch.nn.functional.softmax(output.logits, dim=-1)
+        _, preds = torch.max(de_predictions, dim=-1)
+
+        lang_type = LANGUANGE_MAP[preds.item()]
+        
+        if lang_type not in lang_id:
+            st.success('Unsupported Language')
+            st.write(f"Computation time: {round((time_end-time_start),3)} segs")
+        else:
+            src_lang = lang_id[]
+            trg_lang = lang_id[target_lang]
+            tokenizer.src_lang = src_lang
+            encoded_input = tokenizer(user_input, return_tensors="pt").to(device)
+            generated_tokens = model.generate(
+                **encoded_input, forced_bos_token_id=tokenizer.get_lang_id(trg_lang)
+            )
+            translated_text = tokenizer.batch_decode(
+                generated_tokens, skip_special_tokens=True
+            )[0]
+
+            time_end = time.time()
+            st.success(translated_text)
+
+            st.write(f"Computation time: {round((time_end-time_start),3)} segs")

languages.py

@@ -0,0 +1,47 @@
+LANGUANGE_MAP = {
+    0: 'Arabic',
+    1: 'Basque',
+    2: 'Breton',
+    3: 'Catalan',
+    4: 'Chinese',
+    5: 'Chinese',
+    6: 'Chinese',
+    7: 'Chuvash',
+    8: 'Czech',
+    9: 'Dhivehi',
+    10: 'Dutch',
+    11: 'English',
+    12: 'Esperanto',
+    13: 'Estonian',
+    14: 'French',
+    15: 'Frisian',
+    16: 'Georgian',
+    17: 'German',
+    18: 'Greek',
+    19: 'Hakha_Chin',
+    20: 'Indonesian',
+    21: 'Interlingua',
+    22: 'Italian',
+    23: 'Japanese',
+    24: 'Kabyle',
+    25: 'Kinyarwanda',
+    26: 'Kyrgyz',
+    27: 'Latvian',
+    28: 'Maltese',
+    29: 'Mongolian',
+    30: 'Persian',
+    31: 'Polish',
+    32: 'Portuguese',
+    33: 'Romanian',
+    34: 'Romansh_Sursilvan',
+    35: 'Russian',
+    36: 'Sakha',
+    37: 'Slovenian',
+    38: 'Spanish',
+    39: 'Swedish',
+    40: 'Tamil',
+    41: 'Tatar',
+    42: 'Turkish',
+    43: 'Ukranian',
+    44: 'Welsh'
+ }