Just playing around

app.py

@@ -1,5 +1,8 @@
 import streamlit as st
 
+if not hasattr(st, "cache_resource"):
+    st.cache_resource = st.experimental_singleton
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -10,18 +13,18 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 from transformers import MarianMTModel, MarianTokenizer
 
-model_name = st.radio("Select a model", [
+model_options = [
     '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')
+col1, col2 = st.columns(2)
 
-if not hasattr(st, "cache_resource"):
-    st.cache_resource = st.experimental_singleton
+with col1:
+    model_name = st.selectbox("Select a model", model_options + ['other'])
 
+    if model_name == 'other':
+        model_name = st.text_input("Enter model name", model_options[0])
 
 @st.cache_resource
 def get_tokenizer(model_name):
@@ -42,7 +45,9 @@ else:
     lang_code = None
 
 
-input_text = st.text_input("Enter text to translate", "Hola, mi nombre es Juan")
+with col2:
+    input_text = st.text_input("Enter text to translate", "Hola, mi nombre es Juan")
+
 input_text = input_text.strip()
 if not input_text:
     st.stop()
@@ -51,7 +56,6 @@ if not input_text:
 if lang_code:
     input_text = f"{lang_code} {input_text}"
 
-output_so_far = st.text_input("Enter text translated so far", "Hello, my")
 
 input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
 
@@ -60,8 +64,18 @@ example_generations = model.generate(
     num_beams=4,
     num_return_sequences=4,
 )
-st.write("Example generations:")
-st.write(tokenizer.batch_decode(example_generations, skip_special_tokens=True))
+
+col1, col2 = st.columns(2)
+with col1:
+    st.write("Example generations:")
+    st.write('\n'.join(
+        '- ' + translation
+        for translation in tokenizer.batch_decode(example_generations, skip_special_tokens=True)))
+
+with col2:
+    example_first_word = tokenizer.decode(example_generations[0, 1])
+    output_so_far = st.text_input("Enter text translated so far", example_first_word)
+
 
 # tokenize the output so far
 with tokenizer.as_target_tokenizer():
@@ -93,10 +107,10 @@ with tokenizer.as_target_tokenizer():
     })
 
 st.subheader("Most likely next tokens")
-st.write(probs_table)
+st.table(probs_table.style.hide(axis='index'))
 
 if len(decoder_input_ids) > 1:
-    st.subheader("Loss by token")
+    st.subheader("Loss by already-generated token")
     loss_table = pd.DataFrame({
         'token': [tokenizer.decode(token_id) for token_id in decoder_input_ids[1:]],
         'loss': F.cross_entropy(model_output.logits[0, :-1], torch.tensor(decoder_input_ids[1:]).to(device), reduction='none').cpu()