Update

app.py

@@ -1,17 +1,14 @@
-import transformers
 import streamlit as st
-
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-st.write("Streamlit version:", st.__version__)
-
 st.set_page_config(
     page_title="Romanian Text Generator",
     page_icon="🇷🇴",
     layout="wide"
 )
 
-st.write("Type your text here and press Ctrl+Enter to generate the next sequence:")
+#############################################
+# Python stuff here
 
 model_list = [
     "dumitrescustefan/gpt-neo-romanian-780m",
@@ -20,68 +17,137 @@ model_list = [
     "readerbench/RoGPT2-large"
 ]
 
-st.sidebar.header("Select model")
-model_checkpoint = st.sidebar.selectbox("Select model", model_list)
-
-
-t1, t2, t3, t4, t5, t6 = st.tabs(["Greedy", "Beam-search", "Sampling", "Top-k Sampling", "Top-p Sampling", "Typical Sampling"])
-t1.write("Greedy")
-t2.write("Greedy2")
-t3.write("Greedy3")
-t4.write("Greedy4")
-
+def greedy_search(model, input_ids, attention_mask, no_repeat_ngram_size, max_length):
+    return model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+        max_length=max_length
+    )
 
+def beam_search(model, input_ids, attention_mask, no_repeat_ngram_size, max_length, num_beams):
+    return model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+        max_length=max_length,
+        num_beams=num_beams,
+        early_stopping=True
+    )
 
-st.sidebar.header("Select generation parameters")
-max_length = t5.slider("Max Length", value=20, min_value=10, max_value=200)
-temperature = st.sidebar.slider("Temperature", value=1.0, min_value=0.0, max_value=1.0, step=0.05)
-top_k = st.sidebar.slider("Top-k", min_value=0, max_value=15, step=1, value=0)
-top_p = st.sidebar.slider("Top-p", min_value=0.0, max_value=1.0, step=0.05, value=0.9)
+def sampling(model, input_ids, attention_mask, no_repeat_ngram_size, max_length, temperature, top_k, top_p):
+    return model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+        max_length=max_length,
+        do_sample=True,
+        temperature=temperature,
+        top_k=top_k,
+        top_p=top_p
+    )
 
-text_element = st.text_input('Text:', 'Acesta este un exemplu,')
+def typical_sampling(model, input_ids, attention_mask, no_repeat_ngram_size, max_length, temperature, typical_p):
+    return model.generate(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+        max_length=max_length,
+        do_sample=True,
+        temperature=temperature,
+        typical_p=typical_p,
+        top_k=0
+    )
 
 @st.cache(allow_output_mutation=True)
 def setModel(model_checkpoint):
     model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
     tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
     return model, tokenizer
+#############################################
 
-def infer(model, tokenizer, text, max_length, temperature, top_k, top_p):
-    encoded_prompt = tokenizer(text, add_special_tokens=False, return_tensors="pt")
+col1, _, col2 = st.columns([10, 1, 16])
 
-    output_sequences = model.generate(
-        input_ids=encoded_prompt.input_ids,
-        max_length=max_length,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
-        do_sample=True,
-        num_return_sequences=1
-    )
+with col1:
+    st.write("Step 1: Select model")
 
-    return output_sequences
+    model_checkpoint = st.selectbox("Select model", model_list)
 
-#model, tokenizer = setModel(model_checkpoint)
-#output_sequences = infer(model, tokenizer, text_element, max_length, temperature, top_k, top_p)
-"""
-for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
-    print(f"=== GENERATED SEQUENCE {generated_sequence_idx + 1} ===")
-    generated_sequences = generated_sequence.tolist()
+    st.write("Step 2: Adjust specific text generation parameters")
 
-    # Decode text
-    text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
+    tab_greedy, tab_beamsearch, tab_sampling, tab_typical = st.tabs(["Greedy", "Beam-search", "Sampling", "Typical Sampling"])
 
-    # Remove all text after the stop token
-    # text = text[: text.find(args.stop_token) if args.stop_token else None]
+    with tab_greedy:
+        st.write("as")
 
-    # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
-    total_sequence = (
-            text_element + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)):]
-    )
+    with tab_beamsearch:
+        num_beams = st.slider("Num beams", min_value=1, max_value=30, step=5, value=5)
+
+    with tab_sampling:
+        top_p = st.slider("Top-p", min_value=0.0, max_value=1.0, step=0.05, value=0.9)
+        top_k = st.slider("Top-k", min_value=0, max_value=100, step=10, value=0)
+
+    with tab_typical:
+        typical_p = st.slider("Typical-p", min_value=0., max_value=1., step=.10, value=1.0)
+
+    st.markdown("""---""")
+
+    st.write("Step 3: Adjust common text generation parameters")
 
-    generated_sequences.append(total_sequence)
-    print(total_sequence)
+    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)
 
-st.write(generated_sequences[-1], text_element)
 
-"""
+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")
+
+
+@st.cache(allow_output_mutation=True)
+def setModel(model_checkpoint):
+    model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
+    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+    return model, tokenizer
+
+#####################################################
+# run-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 = ""
+
+if button_greedy:
+    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)
+    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
+    output = sampling(model, input_ids, attention_mask, no_repeat_ngrams, max_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)
+    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")
+if details != "":
+    col2.write(details)

requirements.txt

@@ -1,4 +1,5 @@
 streamlit>=1.12.0
+st_btn_select
 git+https://github.com/huggingface/transformers
 torch
 protobuf~=3.19.0