Updated app.py to include optional reviewers

app.py

@@ -45,7 +45,10 @@ print(f"Model loaded successfully with {len(id2label)} reviewers")
 
 
 def predict_reviewers(
-    pr_title: str, files_input: str, threshold: float = DEFAULT_THRESHOLD
+    pr_title: str,
+    files_input: str,
+    threshold: float = DEFAULT_THRESHOLD,
+    custom_mapping: str = "",
 ) -> tuple[str, str]:
     """
     Predict reviewers for a PR based on title and modified files.
@@ -54,6 +57,7 @@ def predict_reviewers(
         pr_title: The PR title/description
         files_input: Comma or semicolon separated list of modified files
         threshold: Prediction threshold (0-1)
+        custom_mapping: Optional JSON mapping of label IDs to names
 
     Returns:
         tuple: (formatted_predictions, all_scores_json)
@@ -78,6 +82,18 @@ def predict_reviewers(
     if threshold < 0 or threshold > 1:
         return "⚠️ Threshold must be between 0 and 1", ""
 
+    # Parse custom mapping if provided
+    label_mapping = id2label  # Default to model's labels
+    if custom_mapping and custom_mapping.strip():
+        try:
+            parsed_mapping = json.loads(custom_mapping)
+            # Convert string keys to integers
+            label_mapping = {int(k): v for k, v in parsed_mapping.items()}
+        except json.JSONDecodeError:
+            return "⚠️ Invalid JSON format for custom mapping", ""
+        except (ValueError, TypeError):
+            return "⚠️ Custom mapping must have numeric keys", ""
+
     # Format input for the model
     files_text = f"files: {', '.join(files_list)}"
 
@@ -102,7 +118,7 @@ def predict_reviewers(
     all_scores = {}
 
     for idx, prob in enumerate(probabilities):
-        reviewer_name = id2label[idx]
+        reviewer_name = label_mapping.get(idx, f"label_{idx}")
         all_scores[reviewer_name] = float(prob)
 
         if prob > threshold:
@@ -153,11 +169,17 @@ def predict_reviewers(
 
 # Example inputs
 examples = [
-    ["Fix authentication bug in user service", "auth.py, user.py, test_auth.py", 0.5],
+    [
+        "Fix authentication bug in user service",
+        "auth.py, user.py, test_auth.py",
+        0.5,
+        "",
+    ],
     [
         "Add new payment gateway integration",
         "gateway.py; payment_routes.py; config.py",
         0.5,
+        "",
     ],
 ]
 
@@ -197,18 +219,41 @@ with gr.Blocks(title="PR Reviewer Assignment", theme=gr.themes.Soft()) as demo:
                 info="Only show predictions above this confidence score",
             )
 
+            with gr.Accordion("Custom Label Mapping (Optional)", open=False):
+                gr.Markdown(
+                    """
+                    If your deployed model has generic labels (e.g., `label_0`, `label_1`),
+                    you can paste your own ID to name mapping here in JSON format.
+                    
+                    **Example format:**
+                    ```json
+                    {
+                      "0": "John Doe",
+                      "1": "Jane Smith",
+                      "2": "Bob Johnson"
+                    }
+                    ```
+                    """
+                )
+                custom_mapping_input = gr.Code(
+                    label="Custom ID to Label Mapping (JSON)",
+                    language="json",
+                    lines=10,
+                    value="",
+                )
+
             predict_btn = gr.Button("Predict Reviewers", variant="primary", size="lg")
 
         with gr.Column(scale=3):
             prediction_output = gr.Markdown(label="Predictions")
 
-    with gr.Accordion("Detailed JSON Output", open=False):
+    with gr.Accordion("📋 Detailed JSON Output", open=False):
         json_output = gr.JSON(label="Full Prediction Details")
 
     # Connect the button
     predict_btn.click(
         fn=predict_reviewers,
-        inputs=[pr_title_input, files_input, threshold_input],
+        inputs=[pr_title_input, files_input, threshold_input, custom_mapping_input],
         outputs=[prediction_output, json_output],
     )
 
@@ -216,7 +261,7 @@ with gr.Blocks(title="PR Reviewer Assignment", theme=gr.themes.Soft()) as demo:
     gr.Markdown("### Example Inputs")
     gr.Examples(
         examples=examples,
-        inputs=[pr_title_input, files_input, threshold_input],
+        inputs=[pr_title_input, files_input, threshold_input, custom_mapping_input],
         outputs=[prediction_output, json_output],
         fn=predict_reviewers,
         cache_examples=False,