update app

app.py

@@ -1,13 +1,16 @@
+import numpy as np
 import pandas as pd
+import time
 import streamlit as st
-import numpy as np
 import matplotlib.pyplot as plt
 import seaborn as sns
+
 import torch
 import torch.nn.functional as F
+
 from transformers import AlbertTokenizer
-import time
 
+from custom_modeling_albert_flax import CustomFlaxAlbertForMaskedLM
 
 if __name__=='__main__':
 
@@ -38,5 +41,16 @@ if __name__=='__main__':
     st.markdown(define_margins, unsafe_allow_html=True)
     st.markdown(hide_table_row_index, unsafe_allow_html=True)
 
-    from custom_modeling_albert_flax import CustomFlaxAlbertForMaskedLM
-    model = CustomFlaxAlbertForMaskedLM.from_pretrained('albert-base-v2')
+    model = CustomFlaxAlbertForMaskedLM.from_pretrained('albert-xxlarge-v2',from_pt=True)
+    tokenizer = AlbertTokenizer.from_pretrained('albert-xxlarge-v2')
+    mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
+
+    input_ids = tokenizer('This is a sample sentence.',return_tensors='pt')
+    input_ids[0][4] = mask_id
+
+    with torch.no_grad():
+        outputs = model(input_ids)
+    logprobs = F.log_softmax(outputs.logits, dim = -1)
+    st.write(logprobs.shape)
+    preds = [torch.multinomial(torch.exp(probs), num_samples=1).squeeze(dim=-1).item() for probs in logprobs[0]]
+    st.write([tokenizer.decode([token]) for token in preds])