import

next_token.py

@@ -0,0 +1,68 @@
+import streamlit as st
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+import pandas as pd
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+from transformers import MarianMTModel, MarianTokenizer
+model_name = 'Helsinki-NLP/opus-mt-ROMANCE-en'
+
+@st.cache_resource
+def get_tokenizer(model_name):
+    return MarianTokenizer.from_pretrained(model_name)
+
+@st.cache_resource
+def get_model(model_name):
+    return MarianMTModel.from_pretrained(model_name).to(device)
+
+tokenizer = get_tokenizer(model_name)
+model = get_model(model_name)
+
+print(f"The model has {model.num_parameters():,d} parameters.")
+
+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()
+
+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)
+
+# tokenize the output so far
+with tokenizer.as_target_tokenizer():
+    output_tokens = tokenizer.tokenize(output_so_far)
+    decoder_input_ids = tokenizer.convert_tokens_to_ids(output_tokens)
+
+# Add the start token
+decoder_input_ids = [model.config.decoder_start_token_id] + decoder_input_ids
+
+with torch.no_grad():
+    model_output = model(
+        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=5)
+
+probs = last_token_logits.softmax(dim=-1)
+probs_for_likely_tokens = probs[most_likely_tokens.indices]
+
+with tokenizer.as_target_tokenizer():
+    probs_table = pd.DataFrame({
+        '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)
+    })
+
+
+st.write(probs_table)
+st.write(model.config.decoder_start_token_id)

requirements.txt

@@ -0,0 +1,2 @@
+transformers
+pandas