Update app.py

app.py

@@ -3,12 +3,16 @@ Gradio app for PR Reviewer Assignment Model.
 
 This application provides an interactive interface for predicting PR reviewers
 based on PR title and modified files using a fine-tuned DeBERTa model.
+
+For private models, set the HF_TOKEN environment variable:
+    export HF_TOKEN=your_huggingface_token
 """
 
 import gradio as gr
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 import json
+import os
 
 
 # Model configuration
@@ -18,16 +22,31 @@ MODEL_NAME = (
 MAX_LENGTH = 8192
 DEFAULT_THRESHOLD = 0.5
 
+# Authentication token for private models
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
 
 def load_model():
     """
     Load the model and tokenizer.
 
+    For private models, requires HF_TOKEN environment variable to be set.
+
     Returns:
         tuple: (model, tokenizer, id2label)
     """
-    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    if HF_TOKEN:
+        print(f"Using authentication token for private model: {MODEL_NAME}")
+        model = AutoModelForSequenceClassification.from_pretrained(
+            MODEL_NAME, token=HF_TOKEN
+        )
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
+    else:
+        print(f"No token found, attempting to load public model: {MODEL_NAME}")
+        print("If this is a private model, set HF_TOKEN environment variable")
+        model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
     model.eval()
 
     # Get label mappings from model config