Output scores

app.py

@@ -20,14 +20,17 @@ def predict(image1, image2, text):
      outputs = model(input_ids=encoding_1.input_ids, pixel_values=encoding_1.pixel_values, pixel_values_2=encoding_2.pixel_values)
      
     logits = outputs.logits
-    idx = logits.argmax(-1).item()
-    predicted_answer = model.config.id2label[idx]
+    probs = torch.nn.functional.softmax(logits, dim=1)
+    
+    output = dict()
+    for label, id in model.config.label2id.items():
+        output[label] = probs[:,id].item()
    
-    return predicted_answer
+    return output
    
 images = [gr.inputs.Image(type="pil"), gr.inputs.Image(type="pil")]
 text = gr.inputs.Textbox(lines=2, label="Sentence")
-answer = gr.outputs.Textbox(label="Predicted answer")
+label = gr.outputs.Label(num_top_classes=2)
 
 example_sentence_1 = "The left image contains twice the number of dogs as the right image, and at least two dogs in total are standing."
 example_sentence_2 = "One image shows exactly two brown acorns in back-to-back caps on green foliage."
@@ -39,7 +42,7 @@ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2102.033
 
 interface = gr.Interface(fn=predict, 
                          inputs=images + [text], 
-                         outputs=answer, 
+                         outputs=label, 
                          examples=examples, 
                          title=title,
                          description=description,