Update

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 st.set_page_config(
@@ -64,16 +65,24 @@ def setModel(model_checkpoint):
     model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
     tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
     return model, tokenizer
+
+
 #############################################
+col_title, _, col_b1, col_b2, col_b3, _ = st.columns([18, 1, 8, 8, 8, 1])
+col_title.markdown("**Playground for text generation with Romanian models**")
+button_greedy = col_b1.button("Greedy generation")
+button_sampling = col_b2.button("Sampling generation")
+button_typical = col_b3.button("Typical sampling generation")
+
 
 col1, _, col2 = st.columns([10, 1, 16])
 
 with col1:
-    st.write("Step 1: Select model")
+    st.markdown("**Step 1: Select model**")
 
     model_checkpoint = st.selectbox("Select model", model_list)
 
-    st.write("Step 2: Adjust specific text generation parameters")
+    st.markdown("**Step 2: Adjust specific text generation parameters**")
 
     tab_greedy, tab_beamsearch, tab_sampling, tab_typical = st.tabs(["Greedy", "Beam-search", "Sampling", "Typical Sampling"])
 
@@ -92,20 +101,19 @@ with col1:
 
     st.markdown("""---""")
 
-    st.write("Step 3: Adjust common text generation parameters")
+    st.markdown("**Step 3: Adjust common text generation parameters**")
 
     no_repeat_ngrams = st.slider("No repeat n-grams", value=2, min_value=0, max_value=3)
     temperature = st.slider("Temperature", value=1.0, min_value=0.0, max_value=1.0, step=0.05)
-    max_length = st.slider("Max Length", value=20, min_value=10, max_value=200)
+    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**")
 
-with col2:
-    with st.container():
-        button_greedy = st.button("Greedy")
-        button_beam_search = st.button("Beam-search")
-        button_sampling = st.button("Sampling")
-        button_typical = st.button("Typical sampling")
 
+    def update_prompt():
+        st.session_state['text'] = prompt
+
+    prompt = st.selectbox("Select prompt", model_list, on_change=update_prompt)
 
 @st.cache(allow_output_mutation=True)
 def setModel(model_checkpoint):
@@ -114,19 +122,56 @@ def setModel(model_checkpoint):
     return model, tokenizer
 
 #####################################################
-# run-time
+# show-time
 
 if 'text' not in st.session_state:
     st.session_state['text'] = 'Acesta este un exemplu de text generat de un model de limbă.'
 
 details = ""
+tokenized_text = None
 
-if button_greedy:
+if button_greedy or button_sampling or button_typical:
     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
-    output = greedy_search(model, input_ids, attention_mask, no_repeat_ngrams, 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]
+        previous_ids = tokenized_text.input_ids[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]
+        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)
+
+    if previous_ids is not None:
+        new_text = tokenizer.decode(torch.cat([previous_ids, output[0]], dim=1), skip_special_tokens=True)
+    else:
+        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"
 
@@ -135,19 +180,27 @@ if button_sampling:
     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
-    output = sampling(model, input_ids, attention_mask, no_repeat_ngrams, max_length, temperature, top_k, top_p)
+    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 = tokenized_text.input_ids
-    attention_mask = tokenized_text.attention_mask
-    output = typical_sampling(model, input_ids, attention_mask, no_repeat_ngrams, max_length, temperature, typical_p)
+    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("""---""")
+col2.text("Statistics and details:")
 if details != "":
-    col2.write(details)
+    col2.caption("\tGeneration details: " + details)
+if tokenized_text is None:
+    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+tt = tokenizer(text_element, add_special_tokens=False, return_tensors="pt")
+col2.caption(f"\tText length is {len(text_element)} characters, {len(tt.input_ids[0])} tokens.")