add app

app.py

@@ -1,3 +1,7 @@
+import base64
+from io import StringIO
+from math import ceil
+
 from utils import get_resources, simplify
 
 import streamlit as st
@@ -7,25 +11,83 @@ st.set_page_config(
     page_icon="🏃"
 )
 
+BATCH_SIZE = 8
+
+if "text_to_simplify" not in st.session_state:
+    st.session_state["text_to_simplify"] = None
+
 st.title("🏃 Text Simplification in Dutch")
 
-with st.form("input data"):
-    text = st.text_area(label="Input text", value="Met het naderen van de zonovergoten middaghemel op deze betoverende dag, waarbij de atmosferische omstandigheden een onbelemmerde convergentie van cumulusbewolking en uitgestrekte stratosferische azuurblauwe wijdheid faciliteren, lijken de geaggregeerde weersverschijnselen van vandaag, die variëren van sporadische plensbuien tot kalme zuchtjes wind en zeldzame opvlammingen van bliksem, de delicate balans tussen meteorologische complexiteit en eenvoud te weerspiegelen, waardoor de gepassioneerde observator met een gevoel van ontzag en verwondering wordt vervuld.")
-    submitted = st.form_submit_button("Submit")
+fupload_check = st.checkbox("File upload?")
 
-error_ct = st.empty()
-if submitted:
-    text = text.strip()
-    if not text:
-        error_ct.error("Text cannot be empty!", icon="⚠️")
+st.markdown(
+    "Make sure that the file or text in the text box contains **one sentence per line**. Empty lines will"
+    " be removed."
+)
+if fupload_check:
+    uploaded_file = st.file_uploader("Text file", label_visibility="collapsed")
+    if uploaded_file is not None:
+        stringio = StringIO(uploaded_file.getvalue().decode("utf-8"))
+        st.session_state["text_to_simplify"] = stringio.read().strip()
+    else:
+        st.session_state["text_to_simplify"] = None
+else:
+    st.session_state["text_to_simplify"] = st.text_area(
+        label="Sentences to translate", label_visibility="collapsed", height=200,
+        value="Met het naderen van de zonovergoten middaghemel op deze betoverende dag, waarbij de atmosferische omstandigheden een onbelemmerde convergentie van cumulusbewolking en uitgestrekte stratosferische azuurblauwe wijdheid faciliteren, lijken de geaggregeerde weersverschijnselen van vandaag, die variëren van sporadische plensbuien tot kalme zuchtjes wind en zeldzame opvlammingen van bliksem, de delicate balans tussen meteorologische complexiteit en eenvoud te weerspiegelen, waardoor de gepassioneerde observator met een gevoel van ontzag en verwondering wordt vervuld."
+    ).strip()
+
+
+def _get_increment_size(num_sents) -> int:
+    if BATCH_SIZE >= num_sents:
+        return 100
     else:
-        error_ct.info("Generating abstract meaning representation (AMR)...", icon="💻")
+        return ceil(100 / (num_sents / BATCH_SIZE))
 
-        model, tokenizer, streamer = get_resources()
+btn_col, results_col = st.columns(2)
+btn_ct = btn_col.empty()
+error_ct = st.empty()
+simpl_ct = st.container()
+if st.session_state["text_to_simplify"]:
+    if btn_ct.button("Simplify text"):
         error_ct.empty()
+        lines = [strip_line for line in st.session_state["text_to_simplify"].splitlines() if (strip_line := line.strip())]
+        num_sentences = len(lines)
+
+        pbar = st.progress(0, text=f"Simplifying sentences in batches of {BATCH_SIZE}...")
+        increment = _get_increment_size(num_sentences)
+        percent_done = 0
+
+        model, tokenizer = get_resources()
+
+        simpl_ct.caption("Simplified text")
+        output_ct = simpl_ct.empty()
+        all_simplifications = []
+        html = "<ol>"
+        for input_batch, simplifications in simplify(lines, model, tokenizer):
+            for input_text, simplification in zip(input_batch, simplifications):
+                output_ct.empty()
+                html += f"""<li>
+                    <ul>
+                        <li><strong>Input text:</strong> {input_text}</li>
+                        <li><strong>Simplification:</strong> {simplification}</li>
+                    </ul>
+                </li>"""
+                output_ct.markdown(html+"</ol>", unsafe_allow_html=True)
+
+            all_simplifications.extend(simplifications)
+
+            percent_done += increment
+            pbar.progress(min(percent_done, 100))
+        pbar.empty()
 
-        for stream_simplification in simplify(text, model, tokenizer, streamer):
-            st.write(stream_simplification)
+        all_simplifications = "\n".join(all_simplifications) + "\n"
+        b64 = base64.b64encode(all_simplifications.encode("utf-8")).decode("utf-8")
+        results_col.markdown(f'<a download="dutch-simplifications.txt" href="data:file/txt;base64,{b64}" title="Download">Download simplifications</a>', unsafe_allow_html=True)
+else:
+    btn_ct.empty()
+    error_ct.error("Text cannot be empty!", icon="⚠️")
+    simpl_ct.container()
 
 
 ########################

utils.py

@@ -1,5 +1,5 @@
 from threading import Thread
-from typing import Tuple, Generator
+from typing import Tuple, Generator, List
 
 from optimum.bettertransformer import BetterTransformer
 import streamlit as st
@@ -10,7 +10,7 @@ from transformers import T5ForConditionalGeneration, T5Tokenizer, TextStreamer,
 
 
 @st.cache_resource(show_spinner=False)
-def get_resources(quantize: bool = True, no_cuda: bool = False) -> Tuple[T5ForConditionalGeneration, T5Tokenizer, TextIteratorStreamer]:
+def get_resources(quantize: bool = True, no_cuda: bool = False) -> Tuple[T5ForConditionalGeneration, T5Tokenizer]:
     """
     """
     tokenizer = T5Tokenizer.from_pretrained("BramVanroy/ul2-base-dutch-simplification-mai-2023", use_fast=False)
@@ -25,34 +25,36 @@ def get_resources(quantize: bool = True, no_cuda: bool = False) -> Tuple[T5ForCo
         model = quantize_dynamic(model, {nn.Linear, nn.Dropout, nn.LayerNorm}, dtype=qint8)
 
     model.eval()
-    streamer = TextIteratorStreamer(tokenizer, decode_kwargs={"skip_special_tokens": True, "clean_up_tokenization_spaces": True})
 
-    return model, tokenizer, streamer
+    return model, tokenizer
+
+
+def batchify(iterable, batch_size=16):
+    num_items = len(iterable)
+    for idx in range(0, num_items, batch_size):
+        yield iterable[idx:min(idx + batch_size, num_items)]
 
 
 def simplify(
-        text: str,
+        texts: List[str],
         model: T5ForConditionalGeneration,
         tokenizer: T5Tokenizer,
-        streamer: TextIteratorStreamer
-) -> Generator:
+        batch_size: int = 16
+) -> List[str]:
     """
     """
-    text = "[NLG] " + text
-
-    encoded = tokenizer(text, return_tensors="pt")
-    encoded = {k: v.to(model.device) for k, v in encoded.items()}
-    gen_kwargs = {
-        **encoded,
-        "max_new_tokens": 128,
-        "streamer": streamer,
-    }
-
-    with torch.no_grad():
-        thread = Thread(target=model.generate, kwargs=gen_kwargs)
-        thread.start()
-
-        generated_text = ""
-        for new_text in streamer:
-            generated_text += new_text
-            yield generated_text
+
+    for batch_texts in batchify(texts, batch_size=batch_size):
+        nlg_batch_texts = ["[NLG] " + text for text in batch_texts]
+        encoded = tokenizer(nlg_batch_texts, return_tensors="pt", padding=True, truncation=True)
+        encoded = {k: v.to(model.device) for k, v in encoded.items()}
+        gen_kwargs = {
+            "max_new_tokens": 128,
+            "num_beams": 3,
+        }
+
+        with torch.no_grad():
+            encoded = {k: v.to(model.device) for k, v in encoded.items()}
+            generated = model.generate(**encoded, **gen_kwargs).cpu()
+
+        yield batch_texts, tokenizer.batch_decode(generated, skip_special_tokens=True)