Return mtga_card_id in predict output

app.py

@@ -26,14 +26,28 @@ def main():
     def format_for_gallery(predictions):
         # Assuming predictions is like [("Black Lotus", 0.4), ("Mox Ruby", 0.5)]
         return [
-            (card_name_to_image_url(card_name), f"{card_name} ({probability*100:.0f}%)")
-            for card_name, probability in predictions
+            (card_name_to_image_url(card_name), f"{card_name} ({probability*100:.0f}%)", mtga_card_id)
+            for mtga_card_id, card_name, probability in predictions
         ]
 
     def predict_fn(set_code, input):
         state = json.loads(input.replace("'", "\""))
-        return format_for_gallery(predictor.predict(set_code, state))
-    
+        predictions = predictor.predict(set_code, state)
+        
+        # For gallery
+        gallery_output = [
+            (card_name_to_image_url(card_name), f"{card_name} ({probability*100:.0f}%)")
+            for mtga_card_id, card_name, probability in predictions
+        ]
+        
+        # For probabilities list
+        prob_output = [
+            (mtga_card_id, float(probability))
+            for mtga_card_id, _, probability in predictions
+        ]
+        
+        return gallery_output, prob_output
+
     def update_example(set_choice):
         return json.dumps(examples.get(set_choice, {}), indent=2)
     
@@ -89,13 +103,17 @@ def main():
             outputs=[pack_gallery, pool_gallery]
         )
 
-        preidctions_gallery = gr.Gallery(
-            label="Predictions",
-            #height=300,
-            columns=3,
-            object_fit='scale-down',
-            allow_preview=False
-        )
+        with gr.Row():
+            preidctions_gallery = gr.Gallery(
+                label="Predictions",
+                #height=300,
+                columns=3,
+                object_fit='scale-down',
+                allow_preview=False
+            )
+
+            predictions_probs = gr.JSON(label="Card IDs and Probabilities")
+
         # Update text input when dropdown changes
         set_dropdown.change(
             fn=update_example,
@@ -107,7 +125,7 @@ def main():
         gr.Button("Predict").click(
             fn=predict_fn,
             inputs=[set_dropdown, text_input],
-            outputs=[preidctions_gallery]
+            outputs=[preidctions_gallery, predictions_probs]
         )
 
     demo.launch()

src/mtga.py

@@ -49,5 +49,8 @@ def mtga_id_to_card_name(id):
 def card_name_to_image_url(name):
     return cards_by_name[name]['image_uris']['normal']
 
+def card_name_to_mtga_id(name):
+    return cards_by_name[name]['arena_id']
+
 def mtga_id_to_image_url(id):
     return cards_by_mtga_id[id]['image_uris']['normal']

src/predictor.py

@@ -1,6 +1,6 @@
 import pandas as pd
 from fastai.tabular.all import *
-from .mtga import mtga_id_to_card_name
+from .mtga import mtga_id_to_card_name, card_name_to_mtga_id
 
 class MTGPickPredictor:
     def __init__(self):
@@ -19,8 +19,9 @@ class MTGPickPredictor:
         result = ""
         output = []
         for _, (prob, idx) in enumerate(zip(topk_values, topk_indices)):
-            result = result + (f"{self._card_names[set_name][idx]}: {prob*100:.0f}%\n")
-            output.append((self._card_names[set_name][idx], prob))
+            card_name = self._card_names[set_name][idx]
+            result = result + (f"{card_name}: {prob*100:.0f}%\n")
+            output.append((card_name_to_mtga_id(card_name), card_name, prob))
             
         print(result)