Update app.py

app.py

@@ -13,23 +13,26 @@ st.title('KoGPT2 Demo')
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-tokenizer = AutoTokenizer.from_pretrained("skt/kogpt2-base-v2")
-model = GPT2LMHeadModel.from_pretrained('skt/kogpt2-base-v2')
+from transformers import AutoTokenizer, AutoModelWithLMHead
+tokenizer = AutoTokenizer.from_pretrained("BigSalmon/SimplifyText")
+model = AutoModelWithLMHead.from_pretrained("BigSalmon/SimplifyText")
 
 with st.form(key='my_form'):
-    text_input = st.text_input(label='Enter sentence')
+    prompt = st.text_input(label='Enter sentence')
     submit_button = st.form_submit_button(label='Submit')
 
     if submit_button:
       with torch.no_grad():
-        inputs = tokenizer.encode(text_input)
-        gen_ids = model.generate(torch.tensor([inputs]),
-                           max_length=128,
-                           repetition_penalty=2.0,
-                           pad_token_id=tokenizer.pad_token_id,
-                           eos_token_id=tokenizer.eos_token_id,
-                           bos_token_id=tokenizer.bos_token_id,
-                           use_cache=True)
-        generated = tokenizer.decode(gen_ids[0,:].tolist())
-              
-        st.write(generated)
+        text = tokenizer.encode(prompt)
+        myinput, past_key_values = torch.tensor([text]), None
+        myinput = myinput
+        myinput= myinput.to(device)
+        logits, past_key_values = model(myinput, past_key_values = past_key_values, return_dict=False)
+        logits = logits[0,-1]
+        probabilities = torch.nn.functional.softmax(logits)
+        best_logits, best_indices = logits.topk(4)
+        best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]
+        text.append(best_indices[0].item())
+        best_probabilities = probabilities[best_indices].tolist()
+        words = []              
+        st.write(best_words)