Update app.py

app.py

@@ -19,8 +19,10 @@ def get_model():
  
 model, tokenizer = get_model()
 g = "life is a"
+c = 5
 with st.form(key='my_form'):
     prompt = st.text_area(label='Enter sentence', value=g)
+    prompt = st.text_area(label='Enter Number of words', value=c)
     submit_button = st.form_submit_button(label='Submit')
     if submit_button:
       with torch.no_grad():
@@ -35,5 +37,6 @@ with st.form(key='my_form'):
         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(text)
+        words = []
+        final_string = ' '.join(best_words[c])              
+        st.write(final_string)