- added choice of several NN, added choice of several summarisation modes

app.py

@@ -1,32 +1,47 @@
 import gradio as gr
-from summarizer import TransformerSummarizer
+from summarizer import TransformerSummarizer, Summarizer
 
 title = "Summarizer"
 description = """
-This demo is GPT-2 based Summarizer, 
-works with English, Ukrainian, and Russian (and a few other languages too, it`s GPT-2 after all).
+This is a demo of a text summarization NN - based on GPT-2, XLNet, BERT, 
+works with English, Ukrainian, and Russian (and a few other languages too, these are SOTA NN after all).
 """
 
+NN_OPTIONS_LIST = ["mean", "max", "min", "median"]
+NN_LIST = ["GPT-2", "XLNet", "BERT"]
 
-def start_fn(article_input: str) -> str:
+
+def start_fn(article_input: str, reduce_option="mean", model_type='GPT-2') -> str:
     """
     GPT-2 based solution, input full text, output summarized text
+    :param model_type:
+    :param reduce_option:
     :param article_input:
     :return summarized article_output:
     """
-    GPT2_model = TransformerSummarizer(transformer_type="GPT2", transformer_model_key="gpt2-medium")
-    full = ''.join(GPT2_model(article_input, min_length=60))
-    return full
+    if model_type == "GPT-2":
+        GPT2_model = TransformerSummarizer(transformer_type="GPT2", transformer_model_key="gpt2-medium",
+                                           reduce_option=reduce_option)
+        full = ''.join(GPT2_model(article_input, min_length=60))
+        return full
+    elif model_type == "XLNet":
+        XLNet_model = TransformerSummarizer(transformer_type="XLNet", transformer_model_key="xlnet-base-cased",
+                                            reduce_option=reduce_option)
+        full = ''.join(XLNet_model(article_input, min_length=60))
+        return full
+
+    elif model_type == "BERT":
+        BERT_model = Summarizer(reduce_option=reduce_option)
+        full = ''.join(BERT_model(article_input, min_length=60))
+        return full
 
 
 face = gr.Interface(fn=start_fn,
-                    inputs=gr.inputs.Textbox(lines=2, placeholder="Paste article here.", label='Input Article'),
+                    inputs=[gr.inputs.Textbox(lines=2, placeholder="Paste article here.", label='Input Article'),
+                            gr.inputs.Dropdown(NN_OPTIONS_LIST, label="Summarize mode"),
+                            gr.inputs.Dropdown(NN_LIST, label="Selected NN")],
                     outputs=gr.inputs.Textbox(lines=2, placeholder="Summarized article here.", label='Summarized '
                                                                                                      'Article'),
                     title=title,
-                    description=description,)
-#face.launch(server_name="0.0.0.0", share=True)
-
-#launching the app
-if __name__ == "__main__":
-    face.launch(debug=True)
+                    description=description, )
+face.launch(server_name="0.0.0.0", share=True)