Update app.py

app.py

@@ -1,68 +1,74 @@
-import re
 import streamlit as st
+import os
+import io
+from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+import time
+import json
+from typing import List
+import torch
+import random
+import logging
 
-from datetime import datetime
-from transformers import pipeline
-from available_models import MODELS
+if torch.cuda.is_available():
+    device = torch.device("cuda:0")
+else:
+    device = torch.device("cpu")
+    logging.warning("GPU not found, using CPU, translation will be very slow.")
 
+st.cache(suppress_st_warning=True, allow_output_mutation=True)
+st.set_page_config(page_title="M2M100 Translator")
 
-st.set_page_config(page_title="Translator", page_icon="🗣️")
-st.title("🗣️ Translator")
-st.subheader("Translation made fast and easy.")
+lang_id = {
+    
+    "English": "en",
 
-st.write("To add a new model, hit me up! ⬆️")
+    "French": "fr",
 
-with st.expander(label="❓ How does it work", expanded=True):
+}
 
 
-lang1, lang2 = st.columns(2)
-lang1.selectbox(
-    "Source Language", ["🇬🇧 English", "🇫🇷 French", "🇩🇪 German", "🇪🇸 Spanish", "🇷🇺 Russian"], 
-    key="input_lang", index=1,
-)
-lang2.selectbox(
-    "Target Language", ["🇬🇧 English", "🇫🇷 French", "🇩🇪 German", "🇪🇸 Spanish", "🇷🇺 Russian"], 
-    key="output_lang", index=0,
+@st.cache(suppress_st_warning=True, allow_output_mutation=True)
+def load_model(
+    pretrained_model: str = "facebook/m2m100_418M",
+    cache_dir: str = "models/",
+):
+    tokenizer = M2M100Tokenizer.from_pretrained(pretrained_model, cache_dir=cache_dir)
+    model = M2M100ForConditionalGeneration.from_pretrained(
+        pretrained_model, cache_dir=cache_dir
+    ).to(device)
+    model.eval()
+    return tokenizer, model
+
+
+st.title("M2M100 Translator")
+
+
+user_input: str = st.text_area(
+    "Input text",
+    height=200,
+    max_chars=5120,
 )
 
-selected_model = MODELS[f"{st.session_state['input_lang']}->{st.session_state['output_lang']}"]
+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()
 
-if selected_model[0] == None:
-    st.write("No model available for this pair.")
-elif selected_model[0] == 0:
-    st.write("No translation necessary.")
-else:
-    st.markdown(f"""
-    **Selected model:** [{selected_model[0]}]({selected_model[1]})
-    """)
-
-    input_text = st.text_area("Enter text to translate:", height=400, key="input")
-    translate_text = st.button("Translate")
-
-    if translate_text:
-        with st.spinner(text="⚙️ Model loading..."):
-            task = pipeline(
-                "translation",
-                model=selected_model[0],
-                tokenizer=selected_model[0],
-            )
-
-        progress_bar = st.progress(0)
-        with st.spinner(text="🔄 Translating..."):
-            text_to_translate = re.split('(?<=[.!?]) +', input_text)
-            total_progress = len(text_to_translate)
-
-            for i, text in enumerate(text_to_translate):
-                translation = task(text)
-                text_to_translate[i] = translation[0]["translation_text"]
-                progress_bar.progress((i + 1) / total_progress)
-
-        st.success("🗣️ Translated!")
-        st.write(f"**Translation:** {' '.join(text_to_translate)}")
-        st.download_button(
-            label="Download translated text", 
-            data="\n".join(text_to_translate), 
-            file_name=f"{st.session_state['input_lang']}-{st.session_state['output_lang']}-{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}.txt",
-            mime="text/plain"
-        )
+    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")