add new language

app.py

@@ -5,8 +5,19 @@ import os
 
 auth_token = os.environ.get("TOKEN_FROM_SECRET") or True
 
-tokenizer= T5Tokenizer.from_pretrained("Voicelab/vlt5-base-keywords-v4_3-en", use_auth_token=auth_token)
-model = T5ForConditionalGeneration.from_pretrained("Voicelab/vlt5-base-keywords-v4_3-en", use_auth_token=auth_token)
+tokenizer_en = T5Tokenizer.from_pretrained(
+    "Voicelab/vlt5-base-keywords-v4_3-en", use_auth_token=auth_token
+)
+model_en = T5ForConditionalGeneration.from_pretrained(
+    "Voicelab/vlt5-base-keywords-v4_3-en", use_auth_token=auth_token
+)
+
+tokenizer_pl = T5Tokenizer.from_pretrained(
+    "Voicelab/vlt5-base-keywords-v4_3", use_auth_token=auth_token
+)
+model_pl = T5ForConditionalGeneration.from_pretrained(
+    "Voicelab/vlt5-base-keywords-v4_3", use_auth_token=auth_token
+)
 
 img_full = Image.open("images/vl-logo-nlp-blue.png")
 img_short = Image.open("images/sVL-NLP-short.png")
@@ -15,27 +26,49 @@ max_length: int = 1000
 cache_size: int = 100
 
 st.set_page_config(
-    page_title='DEMO - keywords generation',
+    page_title="DEMO - keywords generation",
     page_icon=img_favicon,
     initial_sidebar_state="expanded",
 )
 
-def get_predictions(text):
-    input_ids = tokenizer(
-        text, return_tensors="pt", truncation=True
-    ).input_ids
-    output = model.generate(input_ids, no_repeat_ngram_size=3, num_beams=4)
-    predicted_kw = tokenizer.decode(output[0], skip_special_tokens=True)
+
+def get_predictions(text, language):
+    if language == "Polish":
+        input_ids = tokenizer_pl(text, return_tensors="pt", truncation=True).input_ids
+        output = model_pl.generate(
+            input_ids,
+            no_repeat_ngram_size=2,
+            num_beams=3,
+            num_beam_groups=3,
+            repetition_penalty=1.5,
+            diversity_penalty=2.0,
+            length_penalty=2.0,
+        )
+        predicted_kw = tokenizer_pl.decode(output[0], skip_special_tokens=True)
+    elif language == "English":
+        input_ids = tokenizer_en(text, return_tensors="pt", truncation=True).input_ids
+        output = model_en.generate(
+            input_ids,
+            no_repeat_ngram_size=2,
+            num_beams=3,
+            num_beam_groups=3,
+            repetition_penalty=1.5,
+            diversity_penalty=2.0,
+            length_penalty=2.0,
+        )
+        predicted_kw = tokenizer_en.decode(output[0], skip_special_tokens=True)
     return predicted_kw
 
+
 def trim_length():
     if len(st.session_state["input"]) > max_length:
         st.session_state["input"] = st.session_state["input"][:max_length]
 
 
 if __name__ == "__main__":
+    st.sidebar.image(img_short)
     st.image(img_full)
-    st.title('VLT5 - keywords generation')
+    st.title("VLT5 - keywords generation")
 
     generated_keywords = ""
     user_input = st.text_area(
@@ -45,9 +78,17 @@ if __name__ == "__main__":
         on_change=trim_length,
         key="input",
     )
-    
+
+    language = st.sidebar.title("Model settings")
+    language = st.sidebar.radio(
+        "Select model to test",
+        [
+            "Polish",
+            "English",
+        ],
+    )
+
     result = st.button("Generate keywords")
     if result:
-        generated_keywords = get_predictions(text=user_input)
+        generated_keywords = get_predictions(text=user_input, language=language)
         st.text_area("Generated keywords", generated_keywords)
-