Update

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 import torch
+from time import perf_counter
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 st.set_page_config(
@@ -104,7 +105,7 @@ with col1:
     temperature = st.slider("Temperature", value=1.0, min_value=0.1, max_value=1.0, step=0.1)
     max_length = st.slider("Number of tokens to generate", value=50, min_value=10, max_value=256)
 
-    st.markdown("**Step 4: Select a prompt or input your own text, and click generate in the left panel**")
+    # st.markdown("**Step 4: Select a prompt or input your own text, and click generate in the left panel**")
 
 
 
@@ -129,6 +130,11 @@ details = ""
 tokenized_text = None
 
 if button_greedy or button_sampling or button_typical:
+    if len(st.session_state['text'].strip()) == 0:
+        col2.warning("Please input some text!")
+        text_element = col2.text_area('Text:', height=400, key="text")
+        st.stop()
+
     model, tokenizer = setModel(model_checkpoint)
 
     tokenized_text = tokenizer(st.session_state['text'], add_special_tokens=False, return_tensors="pt")
@@ -144,7 +150,16 @@ if button_greedy or button_sampling or button_typical:
         previous_ids = None
 
     length = min(512, len(input_ids)+max_length)
-    output = greedy_search(model, input_ids.unsqueeze(dim=0), attention_mask.unsqueeze(dim=0), no_repeat_ngrams, length)
+    timer_mark = perf_counter()
+    if button_greedy:
+        output = greedy_search(model, input_ids.unsqueeze(dim=0), attention_mask.unsqueeze(dim=0), no_repeat_ngrams, length)
+        details = f"Text generated using greedy decoding in {perf_counter()-timer_mark:.2f}s"
+    if button_sampling:
+        output = sampling(model, input_ids.unsqueeze(dim=0), attention_mask.unsqueeze(dim=0), no_repeat_ngrams, length, temperature, top_k, top_p)
+        details = f"Text generated using sampling, top-p={top_p:.2f}, top-k={top_k}, temperature={temperature:.2f} in {perf_counter()-timer_mark:.2f}s"
+    if button_typical:
+        output = typical_sampling(model, input_ids.unsqueeze(dim=0), attention_mask.unsqueeze(dim=0), no_repeat_ngrams, length, temperature, typical_p)
+        details = f"Text generated using typical sampling, typical-p={typical_p:.2f}, temperature={temperature:.2f} in {perf_counter()-timer_mark:.2f}s"
 
     if previous_ids is not None:
         print(f"\nConcat prev id: "+tokenizer.decode(previous_ids, skip_special_tokens=True))
@@ -154,46 +169,8 @@ if button_greedy or button_sampling or button_typical:
         new_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
     st.session_state['text'] = new_text
-    details = "Text generated using greedy decoding"
 
-"""
-if button_greedy:
 
-    tokenized_text = tokenizer(st.session_state['text'], add_special_tokens=False, return_tensors="pt")
-    print(f"len text: {len(tokenized_text.input_ids[0])}")
-    print(f"max_len : {max_length}")
-    if len(tokenized_text.input_ids[0]) + max_length > 512:  # need to keep less words
-        keep_last = 512 - max_length
-        print(f"keep last: {keep_last}")
-        input_ids, attention_mask = tokenized_text.input_ids[0][:-keep_last], tokenized_text.attention_mask[0][:-keep_last]
-        st.warning(f"kept last {keep_last}")
-    else:
-        input_ids, attention_mask = tokenized_text.input_ids[0], tokenized_text.attention_mask[0]
-
-    length = min(512, len(input_ids)+max_length)
-    output = greedy_search(model, input_ids.unsqueeze(dim=0), attention_mask.unsqueeze(dim=0), no_repeat_ngrams, length)
-    st.session_state['text'] = tokenizer.decode(output[0], skip_special_tokens=True)
-    details = "Text generated using greedy decoding"
-
-if button_sampling:
-    model, tokenizer = setModel(model_checkpoint)
-    tokenized_text = tokenizer(st.session_state['text'], add_special_tokens=False, return_tensors="pt")
-    input_ids = tokenized_text.input_ids
-    attention_mask = tokenized_text.attention_mask
-    length = min(512, len(input_ids[0]) + max_length)
-    output = sampling(model, input_ids, attention_mask, no_repeat_ngrams, length, temperature, top_k, top_p)
-    st.session_state['text'] = tokenizer.decode(output[0], skip_special_tokens=True)
-    details = f"Text generated using sampling, top-p={top_p:.2f}, top-k={top_k:.2f}, temperature={temperature:.2f}"
-
-if button_typical:
-    model, tokenizer = setModel(model_checkpoint)
-    tokenized_text = tokenizer(st.session_state['text'], add_special_tokens=False, return_tensors="pt")
-    input_ids, attention_mask = tokenized_text.input_ids, tokenized_text.attention_mask
-    length = min(512, len(input_ids[0]) + max_length)
-    output = typical_sampling(model, input_ids, attention_mask, no_repeat_ngrams, length, temperature, typical_p)
-    st.session_state['text'] = tokenizer.decode(output[0], skip_special_tokens=True)
-    details = f"Text generated using typical sampling, typical-p={typical_p:.2f}, temperature={temperature:.2f}"
-"""
 
 text_element = col2.text_area('Text:', height=400, key="text")
 col2.markdown("""---""")