jax

app.py

@@ -5,8 +5,7 @@ import streamlit as st
 import matplotlib.pyplot as plt
 import seaborn as sns
 
-import torch
-import torch.nn.functional as F
+import jax
 
 from transformers import AlbertTokenizer
 
@@ -50,12 +49,12 @@ if __name__=='__main__':
     tokenizer,model = load_model()
     mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
 
-    input_ids = tokenizer('This is a sample sentence.',return_tensors='pt').input_ids
+    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)
-    logprobs = F.log_softmax(outputs.logits, dim = -1)
+    logprobs = jax.nn.log_softmax(outputs.logits, axis = -1)
     st.write(logprobs.shape)
-    preds = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1).item() for probs in logprobs[0]]
+    preds = [np.choice(np.arange(len(probs)),p=np.exp(probs)) for probs in logprobs[0]]
     st.write([tokenizer.decode([token]) for token in preds])