Update app.py

app.py

@@ -20,8 +20,8 @@ from pyabsa import (
 from pyabsa import ABSAInstruction
 from pyabsa.utils.data_utils.dataset_manager import detect_infer_dataset
 
-download_all_available_datasets()
 
+download_all_available_datasets()
 
 def get_atepc_example(dataset):
     task = TaskCodeOption.Aspect_Polarity_Classification
@@ -46,7 +46,6 @@ def get_atepc_example(dataset):
             )
         return sorted(set(lines), key=lines.index)
 
-
 def get_aste_example(dataset):
     task = TaskCodeOption.Aspect_Sentiment_Triplet_Extraction
     dataset_file = detect_infer_dataset(aste_dataset_items[dataset], task)
@@ -63,9 +62,8 @@ def get_aste_example(dataset):
             fin.close()
         return sorted(set(lines), key=lines.index)
 
-
 def get_acos_example(dataset):
-    task = 'ACOS'
+    task = "ACOS"
     dataset_file = detect_infer_dataset(acos_dataset_items[dataset], task)
 
     for fname in dataset_file:
@@ -78,14 +76,15 @@ def get_acos_example(dataset):
             fin = open(f, "r", encoding="utf-8")
             lines.extend(fin.readlines())
             fin.close()
-        lines = [line.split('####')[0] for line in lines]
+        lines = [line.split("####")[0] for line in lines]
         return sorted(set(lines), key=lines.index)
 
-
 try:
     from pyabsa import AspectTermExtraction as ATEPC
 
-    atepc_dataset_items = {dataset.name: dataset for dataset in ATEPC.ATEPCDatasetList()}
+    atepc_dataset_items = {
+        dataset.name: dataset for dataset in ATEPC.ATEPCDatasetList()
+    }
     atepc_dataset_dict = {
         dataset.name: get_atepc_example(dataset.name)
         for dataset in ATEPC.ATEPCDatasetList()
@@ -100,11 +99,16 @@ except Exception as e:
 try:
     from pyabsa import AspectSentimentTripletExtraction as ASTE
 
-    aste_dataset_items = {dataset.name: dataset for dataset in ASTE.ASTEDatasetList()}
+    aste_dataset_items = {
+        dataset.name: dataset for dataset in ASTE.ASTEDatasetList()
+    }
     aste_dataset_dict = {
-        dataset.name: get_aste_example(dataset.name) for dataset in ASTE.ASTEDatasetList()
+        dataset.name: get_aste_example(dataset.name)
+        for dataset in ASTE.ASTEDatasetList()[:-1]
     }
-    triplet_extractor = ASTE.AspectSentimentTripletExtractor(checkpoint="multilingual")
+    triplet_extractor = ASTE.AspectSentimentTripletExtractor(
+        checkpoint="multilingual"
+    )
 except Exception as e:
     print(e)
     aste_dataset_items = {}
@@ -114,9 +118,12 @@ except Exception as e:
 try:
     from pyabsa import ABSAInstruction
 
-    acos_dataset_items = {dataset.name: dataset for dataset in ABSAInstruction.ACOSDatasetList()}
+    acos_dataset_items = {
+        dataset.name: dataset for dataset in ABSAInstruction.ACOSDatasetList()
+    }
     acos_dataset_dict = {
-        dataset.name: get_acos_example(dataset.name) for dataset in ABSAInstruction.ACOSDatasetList()
+        dataset.name: get_acos_example(dataset.name)
+        for dataset in ABSAInstruction.ACOSDatasetList()
     }
     quadruple_extractor = ABSAInstruction.ABSAGenerator("multilingual")
 except Exception as e:
@@ -156,7 +163,7 @@ def perform_aste_inference(text, dataset):
 
     pred_triplets = pd.DataFrame(result["Triplets"])
     true_triplets = pd.DataFrame(result["True Triplets"])
-    return pred_triplets, true_triplets, "{}".format(text.split('####')[0])
+    return pred_triplets, true_triplets, "{}".format(text.split("####")[0])
 
 
 def perform_acos_inference(text, dataset):
@@ -165,108 +172,144 @@ def perform_acos_inference(text, dataset):
             random.randint(0, len(acos_dataset_dict[dataset]) - 1)
         ]
 
-    raw_output = quadruple_extractor.predict(text.split('####')[0], max_length=128)
+    raw_output = quadruple_extractor.predict(text.split("####")[0], max_length=128)
 
-    result = raw_output['Quadruples']
+    result = raw_output["Quadruples"]
     result = pd.DataFrame(result)
     return result, text
 
 
-demo = gr.Blocks()
+def inference(text, dataset, task):
+    if task == "ATEPC":
+        return perform_atepc_inference(text, dataset)
+    elif task == "ASTE":
+        return perform_aste_inference(text, dataset)
+    elif task == "ACOS":
+        return perform_acos_inference(text, dataset)
+    else:
+        raise Exception("No such task: {}".format(task))
 
-with demo:
-    with gr.Row():
 
-        if quadruple_extractor:
-            with gr.Row():
-                with gr.Column():
-                    gr.Markdown("# <p align='center'> ABSA Quadruple Extraction (Experimental) </p>")
-
-                    acos_input_sentence = gr.Textbox(
-                        placeholder="Leave this box blank and choose a dataset will give you a random example...",
-                        label="Example:",
-                    )
-                    acos_dataset_ids = gr.Radio(
-                        choices=[dataset.name for dataset in ABSAInstruction.ACOSDatasetList()],
-                        value="Laptop14",
-                        label="Datasets",
-                    )
-                    acos_inference_button = gr.Button("Let's go!")
-
-                    acos_output_text = gr.TextArea(label="Example:")
-                    acos_output_pred_df = gr.DataFrame(label="Predicted Triplets:")
-
-                    acos_inference_button.click(
-                        fn=perform_acos_inference,
-                        inputs=[acos_input_sentence, acos_dataset_ids],
-                        outputs=[acos_output_pred_df, acos_output_text],
-                    )
-    with gr.Row():
-        if triplet_extractor:
-            with gr.Column():
-                gr.Markdown("# <p align='center'>Aspect Sentiment Triplet Extraction !</p>")
+if __name__ == "__main__":
+    demo = gr.Blocks()
 
+    with demo:
+        with gr.Row():
+            if quadruple_extractor:
                 with gr.Row():
                     with gr.Column():
-                        aste_input_sentence = gr.Textbox(
-                            placeholder="Leave this box blank and choose a dataset will give you a random example...",
-                            label="Example:",
-                        )
                         gr.Markdown(
-                            "You can find code and dataset at [ASTE examples](https://github.com/yangheng95/PyABSA/tree/v2/examples-v2/aspect_sentiment_triplet_extration)"
-                        )
-                        aste_dataset_ids = gr.Radio(
-                            choices=[dataset.name for dataset in ASTE.ASTEDatasetList()[:-1]],
-                            value="Restaurant14",
-                            label="Datasets",
+                            "# <p align='center'> ABSA Quadruple Extraction (Experimental) </p>"
                         )
-                        aste_inference_button = gr.Button("Let's go!")
 
-                        aste_output_text = gr.TextArea(label="Example:")
-                        aste_output_pred_df = gr.DataFrame(label="Predicted Triplets:")
-                        aste_output_true_df = gr.DataFrame(label="Original Triplets:")
-
-                        aste_inference_button.click(
-                            fn=perform_aste_inference,
-                            inputs=[aste_input_sentence, aste_dataset_ids],
-                            outputs=[aste_output_pred_df, aste_output_true_df, aste_output_text],
-                        )
-        if aspect_extractor:
-            with gr.Column():
-                gr.Markdown(
-                    "# <p align='center'>Multilingual Aspect-based Sentiment Analysis !</p>"
-                )
-                with gr.Row():
-                    with gr.Column():
-                        atepc_input_sentence = gr.Textbox(
+                        acos_input_sentence = gr.Textbox(
                             placeholder="Leave this box blank and choose a dataset will give you a random example...",
                             label="Example:",
                         )
-                        gr.Markdown(
-                            "You can find the datasets at [github.com/yangheng95/ABSADatasets](https://github.com/yangheng95/ABSADatasets/tree/v1.2/datasets/text_classification)"
-                        )
-                        atepc_dataset_ids = gr.Radio(
-                            choices=[dataset.name for dataset in ATEPC.ATEPCDatasetList()[:-1]],
+                        acos_dataset_ids = gr.Radio(
+                            choices=[
+                                dataset.name
+                                for dataset in ABSAInstruction.ACOSDatasetList()
+                            ],
                             value="Laptop14",
                             label="Datasets",
                         )
-                        atepc_inference_button = gr.Button("Let's go!")
+                        acos_inference_button = gr.Button("Let's go!")
 
-                        atepc_output_text = gr.TextArea(label="Example:")
-                        atepc_output_df = gr.DataFrame(label="Prediction Results:")
+                        acos_output_text = gr.TextArea(label="Example:")
+                        acos_output_pred_df = gr.DataFrame(label="Predicted Triplets:")
 
-                        atepc_inference_button.click(
-                            fn=perform_atepc_inference,
-                            inputs=[atepc_input_sentence, atepc_dataset_ids],
-                            outputs=[atepc_output_df, atepc_output_text],
+                        acos_inference_button.click(
+                            fn=inference,
+                            inputs=[acos_input_sentence, acos_dataset_ids, gr.Text("ACOS")],
+                            outputs=[acos_output_pred_df, acos_output_text],
                         )
+        with gr.Row():
+            if triplet_extractor:
+                with gr.Column():
+                    gr.Markdown(
+                        "# <p align='center'>Aspect Sentiment Triplet Extraction !</p>"
+                    )
 
-    gr.Markdown(
-        """### GitHub Repo: [PyABSA V2](https://github.com/yangheng95/PyABSA)
-        ### Author: [Heng Yang](https://github.com/yangheng95) (杨恒)
-        [![Downloads](https://pepy.tech/badge/pyabsa)](https://pepy.tech/project/pyabsa) 
-        [![Downloads](https://pepy.tech/badge/pyabsa/month)](https://pepy.tech/project/pyabsa)
-        """
-    )
-
-demo.launch()
+                    with gr.Row():
+                        with gr.Column():
+                            aste_input_sentence = gr.Textbox(
+                                placeholder="Leave this box blank and choose a dataset will give you a random example...",
+                                label="Example:",
+                            )
+                            gr.Markdown(
+                                "You can find code and dataset at [ASTE examples](https://github.com/yangheng95/PyABSA/tree/v2/examples-v2/aspect_sentiment_triplet_extration)"
+                            )
+                            aste_dataset_ids = gr.Radio(
+                                choices=[
+                                    dataset.name
+                                    for dataset in ASTE.ASTEDatasetList()[:-1]
+                                ],
+                                value="Restaurant14",
+                                label="Datasets",
+                            )
+                            aste_inference_button = gr.Button("Let's go!")
+
+                            aste_output_text = gr.TextArea(label="Example:")
+                            aste_output_pred_df = gr.DataFrame(
+                                label="Predicted Triplets:"
+                            )
+                            aste_output_true_df = gr.DataFrame(
+                                label="Original Triplets:"
+                            )
+
+                            aste_inference_button.click(
+                                fn=inference,
+                                inputs=[aste_input_sentence, aste_dataset_ids, gr.Text("ASTE")],
+                                outputs=[
+                                    aste_output_pred_df,
+                                    aste_output_true_df,
+                                    aste_output_text,
+                                ],
+                            )
+            if aspect_extractor:
+                with gr.Column():
+                    gr.Markdown(
+                        "# <p align='center'>Multilingual Aspect-based Sentiment Analysis !</p>"
+                    )
+                    with gr.Row():
+                        with gr.Column():
+                            atepc_input_sentence = gr.Textbox(
+                                placeholder="Leave this box blank and choose a dataset will give you a random example...",
+                                label="Example:",
+                            )
+                            gr.Markdown(
+                                "You can find the datasets at [github.com/yangheng95/ABSADatasets](https://github.com/yangheng95/ABSADatasets/tree/v1.2/datasets/text_classification)"
+                            )
+                            atepc_dataset_ids = gr.Radio(
+                                choices=[
+                                    dataset.name
+                                    for dataset in ATEPC.ATEPCDatasetList()[:-1]
+                                ],
+                                value="Laptop14",
+                                label="Datasets",
+                            )
+                            atepc_inference_button = gr.Button("Let's go!")
+
+                            atepc_output_text = gr.TextArea(label="Example:")
+                            atepc_output_df = gr.DataFrame(label="Prediction Results:")
+
+                            atepc_inference_button.click(
+                                fn=inference,
+                                inputs=[
+                                    atepc_input_sentence,
+                                    atepc_dataset_ids,
+                                    gr.Text("ATEPC"),
+                                ],
+                                outputs=[atepc_output_df, atepc_output_text],
+                            )
+
+        gr.Markdown(
+            """### GitHub Repo: [PyABSA V2](https://github.com/yangheng95/PyABSA)
+            ### Author: [Heng Yang](https://github.com/yangheng95) (杨恒)
+            [![Downloads](https://pepy.tech/badge/pyabsa)](https://pepy.tech/project/pyabsa) 
+            [![Downloads](https://pepy.tech/badge/pyabsa/month)](https://pepy.tech/project/pyabsa)
+            """
+        )
+
+    demo.launch()