cleanups

app.py

@@ -63,6 +63,7 @@ example_generations = model.generate(
     input_ids,
     num_beams=4,
     num_return_sequences=4,
+    max_length=100,
 )
 
 col1, col2 = st.columns(2)
@@ -90,21 +91,34 @@ with torch.no_grad():
         input_ids = input_ids,
         decoder_input_ids = torch.tensor([decoder_input_ids]).to(device))
 
+with st.expander("Configuration"):
+    top_k = st.slider("Number of tokens to show", min_value=1, max_value=100, value=5)
+    temperature = st.slider("Temperature", min_value=0.0, max_value=2.0, value=1.0, step=0.01)
+    show_token_ids = st.checkbox("Show token IDs", value=False)
+    show_logprobs = st.checkbox("Show log probabilities", value=False)
+    show_cumulative_probs = st.checkbox("Show cumulative probabilities", value=False)
+
 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)
+# apply temperature
+last_token_logits_with_temperature = last_token_logits / temperature
+most_likely_tokens = last_token_logits.topk(k=top_k)
 
-probs = last_token_logits.softmax(dim=-1)
+probs = last_token_logits_with_temperature.softmax(dim=-1)
 probs_for_likely_tokens = probs[most_likely_tokens.indices]
 
 with tokenizer.as_target_tokenizer():
-    probs_table = pd.DataFrame({
+    prob_table_data = {
         'token': [tokenizer.decode(token_id) for token_id in most_likely_tokens.indices],
-        'id': most_likely_tokens.indices,
-        'probability': probs_for_likely_tokens,
-        'logprob': probs_for_likely_tokens.log(),
-        'cumulative probability': probs_for_likely_tokens.cumsum(0)
-    })
+    }
+    if show_token_ids:
+        prob_table_data['id'] = most_likely_tokens.indices
+    prob_table_data['probability'] = probs_for_likely_tokens
+    if show_logprobs:
+        prob_table_data['logprob'] = last_token_logits.log_softmax(dim=-1)[most_likely_tokens.indices]
+    if show_cumulative_probs:
+        prob_table_data['cumulative probability'] = probs_for_likely_tokens.cumsum(0)
+    probs_table = pd.DataFrame(prob_table_data)
 
 st.subheader("Most likely next tokens")
 st.table(probs_table.style.hide(axis='index'))