work on test

app.py

@@ -102,9 +102,10 @@ def test():
             correct += pred.eq(target.data.view_as(pred)).sum()
             test_loss /= len(test_loader.dataset)
         test_losses.append(test_loss)
-        print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
+        test_metric = '〽Current test metric - Avg. loss: `{:.4f}`, Accuracy: `{}/{}` (`{:.0f}%`)\n'.format(
         test_loss, correct, len(test_loader.dataset),
-        100. * correct / len(test_loader.dataset)))
+        100. * correct / len(test_loader.dataset))
+        return test_metric
 
 
 
@@ -153,8 +154,12 @@ def image_classifier(inp):
             confidences.update({s:v})
         return confidences
 
+def train_and_test():
+    # Train for one epoch and test
+    train(1,network,optimizer)
+    test_metric = test()
 
-def flag(input_image,correct_result):
+def flag(input_image,correct_result,train):
     # take an image, the wrong result, the correct result.
     # push to dataset. 
     # get size of current dataset 
@@ -197,8 +202,12 @@ def flag(input_image,correct_result):
                 token=HF_TOKEN
             )        
     
-    output = f'<div> Successfully saved to flagged dataset. </div>'
-    return output
+    output = f'<div> ✔ Successfully saved to flagged dataset. </div>'
+    train=True
+    if train:
+        output = f'<div> ✔ Successfully saved to flagged dataset. Training the model on adversarial data! </div>'
+
+    return output,train
 
 
 
@@ -241,8 +250,12 @@ def main():
 
                 flag_btn = gr.Button("Flag")
                 output_result = gr.outputs.HTML()
+                to_train = gr.Variable(value=False)
                 submit.click(image_classifier,inputs = [image_input],outputs=[label_output])
-                flag_btn.click(flag,inputs=[image_input,number_dropdown],outputs=[output_result])
+                flag_btn.click(flag,inputs=[image_input,number_dropdown,to_train],outputs=[output_result,to_train])
+                if to_train.value:
+                    import pdb;pdb.set_trace()
+                    train_and_test()        
                 
             
     block.launch()