Make it fun

app.py

@@ -6,8 +6,6 @@ import torch.nn.functional as F
 import transformers
 import pandas as pd
 
-st.title("Streamlit + Transformers")
-
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 from transformers import MarianMTModel, MarianTokenizer
@@ -21,13 +19,15 @@ model_name = st.radio("Select a model", [
 if model_name == 'other':
     model_name = st.text_input("Enter model name", 'Helsinki-NLP/opus-mt-ROMANCE-en')
 
+if not hasattr(st, "cache_resource"):
+    st.cache_resource = st.experimental_singleton
 
 
-@st.experimental_singleton
+@st.cache_resource
 def get_tokenizer(model_name):
     return MarianTokenizer.from_pretrained(model_name)
 
-@st.experimental_singleton
+@st.cache_resource
 def get_model(model_name):
     model = MarianMTModel.from_pretrained(model_name).to(device)
     print(f"Loaded model, {model.num_parameters():,d} parameters.")
@@ -36,8 +36,10 @@ 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}")
+if tokenizer.supported_language_codes:
+    lang_code = st.selectbox("Select a language", tokenizer.supported_language_codes)
+else:
+    lang_code = None
 
 
 input_text = st.text_input("Enter text to translate", "Hola, mi nombre es Juan")
@@ -45,10 +47,22 @@ input_text = input_text.strip()
 if not input_text:
     st.stop()
 
+# prepend the language code if necessary
+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)
 
+example_generations = model.generate(
+    input_ids,
+    num_beams=4,
+    num_return_sequences=4,
+)
+st.write("Example generations:")
+st.write(tokenizer.batch_decode(example_generations, skip_special_tokens=True))
+
 # tokenize the output so far
 with tokenizer.as_target_tokenizer():
     output_tokens = tokenizer.tokenize(output_so_far)
@@ -62,7 +76,6 @@ with torch.no_grad():
         input_ids = input_ids,
         decoder_input_ids = torch.tensor([decoder_input_ids]).to(device))
 
-
 last_token_logits = model_output.logits[0, -1].cpu()
 assert len(last_token_logits.shape) == 1
 most_likely_tokens = last_token_logits.topk(k=20)
@@ -79,5 +92,11 @@ with tokenizer.as_target_tokenizer():
         'cumulative probability': probs_for_likely_tokens.cumsum(0)
     })
 
-
 st.write(probs_table)
+
+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()
+})
+st.write(loss_table)
+st.write("Total loss so far:", loss_table.loss.sum())