Update app.py

app.py

@@ -1,31 +1,39 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import tensorflow as tf
+import numpy as np
+import pandas as pd
+import os
 import torch
-#maximum number of words in output text
-# MAX_LEN = 30
+import torch.nn as nn
+from transformers.activations import get_activation
+from transformers import AutoTokenizer, AutoModelWithLMHead, AutoModelForCausalLM
 
-title = st.text_input('Enter the seed words', ' ')
-input_sequence = title
+st.title('GPT2:')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-number = st.number_input('Insert how many words', 1)
-MAX_LEN = number
-
-if st.button('Submit'):
-
-    tokenizer = AutoTokenizer.from_pretrained("ml6team/gpt-2-medium-conditional-quote-generator")
+@st.cache(allow_output_mutation=True)
+def get_model():
+tokenizer = AutoTokenizer.from_pretrained("ml6team/gpt-2-medium-conditional-quote-generator")
     model = AutoModelForCausalLM.from_pretrained("ml6team/gpt-2-medium-conditional-quote-generator")
-    inputs = tokenizer.encode(input_sequence, return_tensors="pt")
-# generate text until the output length (which includes the context length) reaches 50
-    #greedy_output = GPT2.generate(input_ids, max_length = MAX_LEN)
-    outputs = model(inputs)
-    outputs = tf.strings.as_string(outputs)
-
-    print("Output:\n" + 100 * '-')
-    st.write(outputs)
-else:
-    st.write(' ')
-
-
-# print("Output:\n" + 100 * '-')
-# print(tokenizer.decode(sample_output[0], skip_special_tokens = True), '...')
+    return model, tokenizer
+    
+ 
+model, tokenizer = get_model()
+g = "life is a"
+with st.form(key='my_form'):
+    prompt = st.text_area(label='Enter sentence', value=g)
+    submit_button = st.form_submit_button(label='Submit')
+    if submit_button:
+      with torch.no_grad():
+        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(350)
+        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)