Update app.py

app.py

@@ -1,36 +1,40 @@
 import streamlit as st
+#import numpy as np
+#import pandas as pd
+#import os
+#import torch
+#import torch.nn as nn
+#from transformers.activations import get_activation
+from transformers import TFGPT2LMHeadModel, GPT2Tokenizer
+import tensorflow as tf
 
-words= st.text_input('Enter some words')
-num_words= st.slider('How long should the output be?', 0, 100, 5)
-button = st.button('Submit')
+st.title('DeepWords')
+st.text('Still under Construction.')
+st.text('Tip: Try writing a sentence and making the model predict final word.')
+#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-@st.cache # only run the function once
-def download_transformer():
-  #for reproducability
-  #SEED = 12
-  
-  from transformers import TFGPT2LMHeadModel, GPT2Tokenizer
+@st.cache(allow_output_mutation=True) # will only run once
+def get_model():
   tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
   GPT2 = TFGPT2LMHeadModel.from_pretrained("gpt2-medium", pad_token_id=tokenizer.eos_token_id)
-  
-  return tokenizer, GPT2
-  
-   
-tokenizer, GPT2 = download_transformer()
+  return GPT2, tokenizer
 
-def input_seq(input_words):
-  import tensorflow as tf
-  return tokenizer.encode(input_words, return_tensors='tf')
-
-if button:
-
-  sample_output = GPT2.generate(
-                             input_seq(words), 
+c = 5
+with st.form(key='my_form'):
+    prompt = st.text_input('Enter sentence:', '')
+    c = st.number_input('Enter Number of words: ', 1)
+    submit_button = st.form_submit_button(label='Submit')
+    if submit_button:
+      tf.random.set_seed(SEED)
+      input_ids = tokenizer.encode(prompt, return_tensors='tf')
+      
+      sample_output = GPT2.generate(
+                             input_ids, 
                              do_sample = True, 
-                             max_length = num_words, 
+                             max_length = c, 
                              top_p = 0.8, 
                              top_k = 0)
+                             
+      st.write(tokenizer.decode(sample_output[0], skip_special_tokens = True), '...')
 
-
-  st.write('Clicked!')
-  st.write(words, num_words)
+