Update

app.py

@@ -1,14 +1,7 @@
 import transformers
 import streamlit as st
 
-from transformers import AutoTokenizer, AutoModelWithLMHead
-
-###################
-# global variables
-
-
-###################
-# page configs and functions
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 st.set_page_config(
     page_title="Romanian Text Generator",
@@ -16,29 +9,33 @@ st.set_page_config(
     layout="wide"
 )
 
-model_list = ["dumitrescustefan/gpt-neo-romanian-780m"]
-st.sidebar.header("Select Model")
-model_checkpoint = st.sidebar.radio("", model_list)
-text_element = st.text_input('Text:', 'Acesta este un exemplu,')
-
-
+st.write("Type your text here and press Ctrl+Enter to generate the next sequence:")
 
+model_list = [
+    "dumitrescustefan/gpt-neo-romanian-780m"
+    "readerbench/RoGPT2-base",
+    "readerbench/RoGPT2-medium",
+    "readerbench/RoGPT2-large"
+]
 
+st.sidebar.header("Select model")
+model_checkpoint = st.sidebar.radio("", model_list)
 
-st.sidebar.header("Select type of PERSON detection")
+st.sidebar.header("Select generation parameters")
 max_length = st.sidebar.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)
 
+text_element = st.text_input('Text:', 'Acesta este un exemplu,')
 
 @st.cache(allow_output_mutation=True)
 def setModel(model_checkpoint):
-    model = AutoModelWithLMHead.from_pretrained(model_checkpoint)
+    model = AutoModelForCausalLM.from_pretrained(model_checkpoint)
     tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
     return model, tokenizer
 
-def infer(model, tokenizer, text, input_ids, max_length, temperature, top_k, top_p):
+def infer(model, tokenizer, text, max_length, temperature, top_k, top_p):
     encoded_prompt = tokenizer.encode(text, add_special_tokens=False, return_tensors="pt")
     output_sequences = model.generate(
         input_ids=encoded_prompt.input_ids,
@@ -53,7 +50,7 @@ def infer(model, tokenizer, text, input_ids, max_length, temperature, top_k, top
     return output_sequences
 
 model, tokenizer = setModel(model_checkpoint)
-output_sequences = infer(model, tokenizer, text_element, input_ids, max_length, temperature, top_k, top_p)
+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} ===")