Model selector

app.py

@@ -11,7 +11,17 @@ st.title("Streamlit + Transformers")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 from transformers import MarianMTModel, MarianTokenizer
-model_name = 'Helsinki-NLP/opus-mt-ROMANCE-en'
+
+model_name = st.radio("Select a model", [
+    'Helsinki-NLP/opus-mt-roa-en',
+    'Helsinki-NLP/opus-mt-en-roa',
+    'other'
+])
+
+if model_name == 'other':
+    model_name = st.text_input("Enter model name", 'Helsinki-NLP/opus-mt-ROMANCE-en')
+
+
 
 @st.experimental_singleton
 def get_tokenizer(model_name):
@@ -26,6 +36,9 @@ def get_model(model_name):
 tokenizer = get_tokenizer(model_name)
 model = get_model(model_name)
 
+if tokenizer.supported_language_codes is not None:
+    st.write(f"Supported languages: {tokenizer.supported_language_codes}")
+
 
 input_text = st.text_input("Enter text to translate", "Hola, mi nombre es Juan")
 input_text = input_text.strip()
@@ -52,7 +65,7 @@ with torch.no_grad():
 
 last_token_logits = model_output.logits[0, -1].cpu()
 assert len(last_token_logits.shape) == 1
-most_likely_tokens = last_token_logits.topk(k=5)
+most_likely_tokens = last_token_logits.topk(k=20)
 
 probs = last_token_logits.softmax(dim=-1)
 probs_for_likely_tokens = probs[most_likely_tokens.indices]