Add language prediction, translation

app.py

@@ -50,7 +50,7 @@ st.set_page_config(page_title="Indonesian GPT-2 Demo")
 
 st.title("Indonesian GPT-2")
 
-# ft_model = fasttext.load_model('lid.176.ftz')
+ft_model = fasttext.load_model('lid.176.ftz')
 
 # Sidebar
 st.sidebar.image(LOGO)
@@ -119,6 +119,12 @@ text = st.text_area("Enter text", session_state.prompt_box)
 
 if st.button("Run"):
     with st.spinner(text="Getting results..."):
+        lang_predictions, lang_probability = ft_model.predict(text.replace("\n", " "), k=3)
+        if "__label__id" in lang_predictions:
+            lang = "id"
+        else:
+            lang = lang_predictions[0].replace("__label__", "")
+            text = translate(text, "id", lang)
 
         st.subheader("Result")
         model = load_gpt(model_name)
@@ -138,3 +144,19 @@ if st.button("Run"):
 
         st.text("Translation")
         translation = translate(text, "en", "id")
+
+        if lang == "id":
+            st.write(translation.replace("\n", "  \n"))
+
+        else:
+            st.write(translate(result, lang, "id").replace("\n", "  \n"))
+
+        image_cat = "https://media.giphy.com/media/vFKqnCdLPNOKc/giphy.gif"
+        image = get_image(translation.replace("\"", "'"))
+
+        if image is not "":
+            st.image(image, width=400)
+
+        else:
+            # display cat image if no image found
+            st.image(image_cat, width=400)