testing oauth

app.py

@@ -1,5 +1,6 @@
 import itertools
 import json
+import os
 import tempfile
 
 import biotite.structure as bs
@@ -15,8 +16,8 @@ from biotite.database import rcsb
 from biotite.sequence import io as seqio
 from biotite.structure import filter_amino_acids, io, spread_residue_wise, to_sequence
 from gradio_molecule3d import Molecule3D
-from huggingface_hub import get_hf_file_metadata
-from huggingface_hub.utils import GatedRepoError
+from huggingface_hub import HfApi
+from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError
 from matplotlib.cm import ScalarMappable
 from matplotlib.colors import Normalize
 
@@ -117,7 +118,6 @@ def predict_rocketshp(
     structure_code: str | None,
     structure_file: str | None,
     chain_id: str | None,
-    oauth_token: gr.OAuthToken | None = None,
 ):
     print(f"sequence text: {sequence}")
     print(f"sequence file: {sequence_file}")
@@ -125,12 +125,13 @@ def predict_rocketshp(
     print(f"structure file: {structure_file}")
     print(f"model variant: {model_variant}")
 
-    if oauth_token is None:
-        raise gr.Error("Please log in to use this Space")
-    token_value = oauth_token.token
-    check_permissions(token_value)
+    is_authorized, token = check_user_access()
+    if not is_authorized:
+        raise gr.Error("Failed to authorize repository access.")
 
+    # Load the model
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    model = RocketSHP.load_from_checkpoint(model_variant).to(device)
     is_sequence_model = "seq" in model_variant or "mini" in model_variant
 
     if is_sequence_model:
@@ -145,7 +146,7 @@ def predict_rocketshp(
             raise gr.Error("Sequence input is required for the selected model.")
 
         struct_features = None
-        seq_features = load_sequence(sequence, device=device)
+        seq_features = load_sequence(sequence, device=device, HF_TOKEN=token)
 
     else:
         if structure_file is None:
@@ -206,9 +207,6 @@ def predict_rocketshp(
         sequence = str(to_sequence(structure)[0][0])
         seq_features = load_sequence(sequence, device=device)
 
-    # Load the model
-    model = RocketSHP.load_from_checkpoint(model_variant).to(device)
-
     # Make predictions
     with torch.no_grad():
         try:
@@ -410,22 +408,36 @@ def visualize_network(
     return fig, bc_highlight, comm_highlight, out_cluster_file_name
 
 
-def check_permissions(token: str | None = None) -> None:
-    if token is None:
-        raise gr.Error("Please log in to use this Space")
+def check_user_access(profile: gr.OAuthProfile | None):
+    """Check if user is logged in and has access to private repo"""
+    if profile is None:
+        return False, "Please log in to use this Space"
+
     try:
-        url = huggingface_hub.hf_hub_url(
+        # Try to access a file from your private repo
+        api = HfApi(token=profile.oauth_info.access_token)
+
+        # Test access by trying to get repo info
+        api.repo_info(
             repo_id="EvolutionaryScale/esm3-sm-open-v1",
-            repo_type="model",
-            filename="config.json",
+            repo_type="model",  # or "dataset" or "space"
+            token=profile.oauth_info.access_token,
         )
-        get_hf_file_metadata(url=url)
-        return
+
+        return True, profile.oauth_info.access_token
+
     except GatedRepoError:
-        raise gr.Error(
-            "You must have access to ... to run this Space. Please go through the gating process and come back."
+        return (
+            False,
+            "You need to request access to the private repository at https://huggingface.co/username/private-repo-name",
         )
 
+    except RepositoryNotFoundError:
+        return False, "You don't have access to the required repository"
+
+    except Exception as e:
+        return False, f"Error checking access: {str(e)}"
+
 
 reps = [
     {