Split sentences for model that requires it. Default beams to 1

app.py

@@ -66,9 +66,9 @@ It was a quite young girl, unknown to me, with a hood over her head, and with la
     st.session_state["text"] = st.text_area(
         "Enter text", st.session_state.prompt_box, height=300
     )
-    num_beams = st.sidebar.number_input("Num beams", min_value=1, max_value=10, value=6)
+    num_beams = st.sidebar.number_input("Num beams", min_value=1, max_value=10, value=1)
     num_beam_groups = st.sidebar.number_input(
-        "Num beam groups", min_value=1, max_value=10, value=3
+        "Num beam groups", min_value=1, max_value=10, value=1
     )
     length_penalty = st.sidebar.number_input(
         "Length penalty", min_value=0.0, max_value=2.0, value=1.2, step=0.1
@@ -97,7 +97,7 @@ and the [Huggingface text generation interface doc](https://huggingface.co/trans
             time_end = time.time()
             time_diff = time_end - time_start
 
-            st.write(result[0].replace("\n", "  \n"))
+            st.write(result.replace("\n", "  \n"))
             text_line = ", ".join([f"{k}={v}" for k, v in params_used.items()])
             st.markdown(f"    🕙 *generated in {time_diff:.2f}s, `{text_line}`*")
 

generator.py

@@ -43,6 +43,7 @@ class Generator:
         self.model_name = model_name
         self.task = task
         self.desc = desc
+        self.split_sentences = split_sentences
         self.tokenizer = None
         self.model = None
         self.prefix = ""
@@ -92,8 +93,14 @@ class Generator:
     def generate(self, text: str, **generate_kwargs) -> (str, dict):
         # Replace two or more newlines with a single newline in text
         text = re.sub(r"\n{2,}", "\n", text)
-
         generate_kwargs = {**self.gen_kwargs, **generate_kwargs}
+
+        # if there are newlines in the text, and the model needs line-splitting, split the text
+        if re.search(r"\n", text) and self.split_sentences:
+            lines = text.splitlines()
+            translated = [self.generate(line, **generate_kwargs)[0] for line in lines]
+            return "\n".join(translated), generate_kwargs
+
         batch_encoded = self.tokenizer(
             self.prefix + text,
             max_length=generate_kwargs["max_length"],
@@ -115,7 +122,7 @@ class Generator:
             pred.replace("<pad> ", "").replace("<pad>", "").replace("</s>", "")
             for pred in decoded_preds
         ]
-        return decoded_preds, generate_kwargs
+        return decoded_preds[0], generate_kwargs
 
     def __str__(self):
         return self.desc