fix

app.py

@@ -56,5 +56,5 @@ if __name__=='__main__':
     outputs = model(input_ids)
     logprobs = jax.nn.log_softmax(outputs.logits, axis = -1)
     st.write(logprobs.shape)
-    preds = [np.random.choice(jnp.arange(len(probs)),p=jnp.exp(probs)/jnp.sum(jnp.exp(probs))) for probs in logprobs[0]]
+    preds = [np.random.choice(np.arange(len(probs)),p=np.exp(probs)/np.sum(np.exp(probs))) for probs in logprobs[0]]
     st.write([tokenizer.decode([token]) for token in preds])