update

app.py

@@ -6,6 +6,7 @@ import matplotlib.pyplot as plt
 import seaborn as sns
 
 import jax
+import jax.numpy as jnp
 
 from transformers import AlbertTokenizer
 
@@ -52,9 +53,8 @@ if __name__=='__main__':
     input_ids = tokenizer('This is a sample sentence.',return_tensors='np').input_ids
     input_ids[0][4] = mask_id
 
-    with torch.no_grad():
-        outputs = model(input_ids)
+    outputs = model(input_ids)
     logprobs = jax.nn.log_softmax(outputs.logits, axis = -1)
     st.write(logprobs.shape)
-    preds = [np.choice(np.arange(len(probs)),p=np.exp(probs)) for probs in logprobs[0]]
+    preds = [jnp.random.choice(jnp.arange(len(probs)),p=jnp.exp(probs)) for probs in logprobs[0]]
     st.write([tokenizer.decode([token]) for token in preds])