make the app more 'user friendly'

app.py

@@ -10,56 +10,56 @@ if gr.NO_RELOAD:
 
 DEVICE = 'cpu'
 MODELS = [
-    (
-        'bert-model_1950',
-        lambda: BaseTransferLearningModel(
-            'bert-base-uncased',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/bert-model_1950.pt',
-        ),
-    ),
-    (
-        'bert-model_2000',
-        lambda: BaseTransferLearningModel(
-            'bert-base-uncased',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/bert-model_2000.pt',
-        ),
-    ),
-    (
-        'deberta-base-model_1100',
-        lambda: BaseTransferLearningModel(
-            'microsoft/deberta-base',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/deberta-base-model_4400.pt',
-        ),
-    ),
-    (
-        'deberta-base-model_2000',
-        lambda: BaseTransferLearningModel(
-            'microsoft/deberta-base',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/deberta-base-model_8000.pt',
-        ),
-    ),
-    (
-        'deberta-v3-base-model_1700',
-        lambda: BaseTransferLearningModel(
-            'microsoft/deberta-v3-base',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/deberta-v3-base-model_3400.pt',
-        ),
-    ),
+    # (
+    #     'bert-model_1950',
+    #     lambda: BaseTransferLearningModel(
+    #         'bert-base-uncased',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/bert-model_1950.pt',
+    #     ),
+    # ),
+    # (
+    #     'bert-model_2000',
+    #     lambda: BaseTransferLearningModel(
+    #         'bert-base-uncased',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/bert-model_2000.pt',
+    #     ),
+    # ),
+    # (
+    #     'deberta-base-model_1100',
+    #     lambda: BaseTransferLearningModel(
+    #         'microsoft/deberta-base',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/deberta-base-model_4400.pt',
+    #     ),
+    # ),
+    # (
+    #     'deberta-base-model_2000',
+    #     lambda: BaseTransferLearningModel(
+    #         'microsoft/deberta-base',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/deberta-base-model_8000.pt',
+    #     ),
+    # ),
+    # (
+    #     'deberta-v3-base-model_1700',
+    #     lambda: BaseTransferLearningModel(
+    #         'microsoft/deberta-v3-base',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/deberta-v3-base-model_3400.pt',
+    #     ),
+    # ),
     (
         'deberta-v3-base-model_2000',
         lambda: BaseTransferLearningModel(
@@ -70,58 +70,64 @@ MODELS = [
             state_dict='src/ckpt/deberta-v3-base-model_4000.pt',
         ),
     ),
-    (
-        'distilbert-model_1850',
-        lambda: BaseTransferLearningModel(
-            'distilbert-base-uncased',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/distilbert-model_1850.pt',
-        ),
-    ),
-    (
-        'distilbert-model_2000',
-        lambda: BaseTransferLearningModel(
-            'distilbert-base-uncased',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/distilbert-model_2000.pt',
-        ),
-    ),
-    (
-        'roberta-base-model_1250',
-        lambda: BaseTransferLearningModel(
-            'FacebookAI/roberta-base',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/roberta-base-model_1250.pt',
-        ),
-    ),
-    (
-        'roberta-base-model_2000',
-        lambda: BaseTransferLearningModel(
-            'FacebookAI/roberta-base',
-            [('linear', ['in', 'out']), ('softmax')],
-            2,
-            device=DEVICE,
-            state_dict='src/ckpt/roberta-base-model_2000.pt',
-        ),
-    ),
+    # (
+    #     'distilbert-model_1850',
+    #     lambda: BaseTransferLearningModel(
+    #         'distilbert-base-uncased',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/distilbert-model_1850.pt',
+    #     ),
+    # ),
+    # (
+    #     'distilbert-model_2000',
+    #     lambda: BaseTransferLearningModel(
+    #         'distilbert-base-uncased',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/distilbert-model_2000.pt',
+    #     ),
+    # ),
+    # (
+    #     'roberta-base-model_1250',
+    #     lambda: BaseTransferLearningModel(
+    #         'FacebookAI/roberta-base',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/roberta-base-model_1250.pt',
+    #     ),
+    # ),
+    # (
+    #     'roberta-base-model_2000',
+    #     lambda: BaseTransferLearningModel(
+    #         'FacebookAI/roberta-base',
+    #         [('linear', ['in', 'out']), ('softmax')],
+    #         2,
+    #         device=DEVICE,
+    #         state_dict='src/ckpt/roberta-base-model_2000.pt',
+    #     ),
+    # ),
 ]
 
 
 class WebUI:
 
-    def __init__(self, models: list[(str, Callable)] = [], device: str = 'cpu') -> None:
+    def __init__(
+        self,
+        models: list[(str, Callable)] = [],
+        device: str = 'cpu',
+        debug: bool = False,
+    ) -> None:
         self.models = models
         self.device = device
         self.is_ready = False
         self.model = self.models[0][1]()
         self.is_ready = True
         self.scraper = GenericScraper()
+        self.debug = debug
 
     def _change_model(self, idx: int) -> str:
         if gr.NO_RELOAD:
@@ -142,7 +148,9 @@ class WebUI:
         if self.is_ready == False:
             return 'Model is not yet ready!'
         output = self.model.predict(text, self.device).detach().cpu().numpy()[0]
-        return f'Fake: {output[0]:.10f}, Real: {output[1]:.10f}'
+        if self.debug:
+            return f'Fake: {output[0]:.10f}, Real: {output[1]:.10f}'
+        return f'We think that this is a {"fake" if output[0] > output[1] else "real"} news article with {max(output[0], output[1]) * 100:.2f}% certainty.'
 
     def _scrape(self, url: str) -> str:
         try:
@@ -173,16 +181,18 @@ class WebUI:
                         )
                         btn_submit = gr.Button(value='Submit', variant='primary')
                 with gr.Column():
-                    ddl_model = gr.Dropdown(
-                        label='Model',
-                        choices=[model[0] for model in self.models],
-                        value=self._change_model(0),
-                        type='index',
-                        interactive=True,
-                        filterable=True,
-                    )
+                    if self.debug:
+                        ddl_model = gr.Dropdown(
+                            label='Model',
+                            choices=[model[0] for model in self.models],
+                            value=self._change_model(0),
+                            type='index',
+                            interactive=True,
+                            filterable=True,
+                        )
                     t_out = gr.Textbox(label='Output')
-            ddl_model.change(fn=self._change_model, inputs=ddl_model)
+            if self.debug:
+                ddl_model.change(fn=self._change_model, inputs=ddl_model)
             btn_scrape.click(fn=self._scrape, inputs=t_url, outputs=t_inp)
             btn_submit.click(fn=self._predict, inputs=t_inp, outputs=t_out)
         return ui