Update app.py

app.py

@@ -21,8 +21,8 @@ model, tokenizer = get_model()
 #g = 
 c = 5
 with st.form(key='my_form'):
-    prompt = st.text_input('Enter sentence', '')
-    #prompt = st.text_area(label='Enter Number of words', value=c)
+    prompt = st.text_input('Enter sentence:', '')
+    c = st.number_input(label='Enter Number of words: ', 1)
     submit_button = st.form_submit_button(label='Submit')
     if submit_button:
       with torch.no_grad():
@@ -39,6 +39,6 @@ with st.form(key='my_form'):
         best_probabilities = probabilities[best_indices].tolist()
         words = []
         
-        best_words = ' '.join(best_words[0:5])
+        best_words = ' '.join(best_words[0:c])
         final_string = prompt + best_words             
         st.write(final_string)