add num_tokens arg

app.py

@@ -146,9 +146,9 @@ def main():
         act_type = "sigmoid"
     else:
         label2id = cls_model_config.label2id
-        filtered_label2id = {k: v  for k, v in label2id.items() if "negative" in k}
-        label_key = st.selectbox(ATTRIBUTE_LABEL[language], label2id.keys())
-        target_label_id = label2id[label_key]
+        filtered_label2id = {k: v for k, v in label2id.items() if "negative" in k.lower()}
+        label_key = st.selectbox(ATTRIBUTE_LABEL[language], filtered_label2id.keys())
+        target_label_id = filtered_label2id[label_key]
         act_type = "softmax"
     st.write(WARNING_TEXT[language])
     show_pos_alpha = st.checkbox("Show positive alphas", value=False)
@@ -183,6 +183,8 @@ def main():
         prompt = st.text_input(TEXT_PROMPT_LABEL[language], "The movie")
     else:
         prompt = st.text_input(TEXT_PROMPT_LABEL[language], PROMPT_EXAMPLE[language])
+    num_tokens = st.slider("# tokens to be generated", min_value=5, max_value=40, step=1, value=20)
+    num_tokens = int(num_tokens)
     st.subheader("Generated text:")
 
     def generate():
@@ -194,7 +196,8 @@ def main():
             target_label_id=target_label_id,
             entropy_threshold=entropy_threshold,
             fp16=fp16,
-            act_type=act_type
+            act_type=act_type,
+            num_tokens=num_tokens
         )
 
     st.button("Generate new", on_click=generate())
@@ -225,7 +228,8 @@ def inference(
         alpha: float = 5,
         target_label_id: int = 0,
         entropy_threshold: float = 0,
-        act_type: str = "sigmoid"
+        act_type: str = "sigmoid",
+        num_tokens=10,
 ) -> str:
     torch.set_grad_enabled(False)
     generator = load_generator(lm_model_name=lm_model_name)
@@ -259,7 +263,7 @@ def inference(
         sequences, tokens = generator.sample_sequences(
             num_samples=1,
             input_prompt=prompt,
-            max_length=20,
+            max_length=num_tokens,
             caif_period=1,
             entropy=entropy_threshold,
             progress_bar=progress_bar,