oauth and pdb struct reading bug fixes

app.py

@@ -17,7 +17,8 @@ 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 HfApi
-from huggingface_hub.utils import RepositoryNotFoundError, GatedRepoError
+from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
+from loguru import logger
 from matplotlib.cm import ScalarMappable
 from matplotlib.colors import Normalize
 
@@ -118,15 +119,17 @@ def predict_rocketshp(
     structure_code: str | None,
     structure_file: str | None,
     chain_id: str | None,
-    profile: gr.OAuthProfile | None,
+    token: gr.OAuthToken | None,
 ):
-    print(f"sequence text: {sequence}")
-    print(f"sequence file: {sequence_file}")
-    print(f"structure code: {structure_code}")
-    print(f"structure file: {structure_file}")
-    print(f"model variant: {model_variant}")
-
-    is_authorized, token = check_user_access(profile)
+    logger.info(f"sequence text: {sequence}")
+    logger.info(f"sequence file: {sequence_file}")
+    logger.info(f"structure code: {structure_code}")
+    logger.info(f"structure file: {structure_file}")
+    logger.info(f"model variant: {model_variant}")
+
+    is_authorized, token, message = check_user_access(token)
+    logger.info(is_authorized)
+    logger.info(message)
     if not is_authorized:
         raise gr.Error("Failed to authorize repository access.")
 
@@ -141,7 +144,7 @@ def predict_rocketshp(
                 gr.Warning("Sequence file provided, ignoring text box.")
 
             sequence = str(seqio.load_sequence(sequence_file))
-            print(sequence)
+            logger.info(sequence)
 
         elif sequence == "":
             raise gr.Error("Sequence input is required for the selected model.")
@@ -160,12 +163,15 @@ def predict_rocketshp(
                 "pdb",
                 target_path=structure_tmp_dir.name,
             )
-            print(structure_tmp_dir)
-            print(structure_file)
+            logger.info(structure_tmp_dir)
+            logger.info(structure_file)
         elif structure_code != "":
             gr.Warning(f"PDB file provided, ignoring PDB code {structure_code}.")
 
-        structure = io.load_structure(structure_file)
+        try:
+            structure = io.load_structure(structure_file)
+        except ValueError:
+            raise gr.Error(f"Invalid PDB Code {structure_code}")
 
         if isinstance(structure, bs.AtomArrayStack):
             gr.Info(
@@ -201,8 +207,8 @@ def predict_rocketshp(
                 f"No amino acid residues found in chain {chain_id} of the provided structure."
             )
 
-        print(len(structure))
-        print(structure[:3])
+        logger.info(len(structure))
+        logger.info(structure[:3])
 
         struct_features = load_structure(structure, device=device)
         sequence = str(to_sequence(structure)[0][0])
@@ -210,6 +216,14 @@ def predict_rocketshp(
 
     # Make predictions
     with torch.no_grad():
+        logger.info(f"Sequence length: {len(sequence)}")
+        logger.info(
+            f"Structure features shape: {struct_features.shape if struct_features is not None else 'N/A'}"
+        )
+        logger.info(
+            f"Sequence features shape: {seq_features.shape if seq_features is not None else 'N/A'}"
+        )
+
         try:
             dynamics_pred = model(
                 {
@@ -268,10 +282,10 @@ def cluster_network(G: nx.Graph, k: int = 5):
     Cluster the network using Girvan-Newman algorithm.
     """
 
-    print(f"Nodes: {G.number_of_nodes()}")
-    print(f"Edges: {G.number_of_edges()}")
-    print(f"Number of connected components: {nx.number_connected_components(G)}")
-    print(f"Connected: {nx.is_connected(G)}")
+    logger.info(f"Nodes: {G.number_of_nodes()}")
+    logger.info(f"Edges: {G.number_of_edges()}")
+    logger.info(f"Number of connected components: {nx.number_connected_components(G)}")
+    logger.info(f"Connected: {nx.is_connected(G)}")
 
     comp = nx.community.girvan_newman(G)
     # limited = itertools.takewhile(lambda c: len(c) <= k, comp)
@@ -409,35 +423,40 @@ def visualize_network(
     return fig, bc_highlight, comm_highlight, out_cluster_file_name
 
 
-def check_user_access(profile: gr.OAuthProfile | None):
+def check_user_access(token: gr.OAuthToken | None):
     """Check if user is logged in and has access to private repo"""
-    if profile is None:
+    if token is None:
         return False, "Please log in to use this Space"
 
     try:
         # Try to access a file from your private repo
-        api = HfApi(token=profile.oauth_info.access_token)
+        api = HfApi(token=token)
 
         # Test access by trying to get repo info
         api.repo_info(
             repo_id="EvolutionaryScale/esm3-sm-open-v1",
             repo_type="model",  # or "dataset" or "space"
-            token=profile.oauth_info.access_token,
+            token=token,
         )
 
-        return True, profile.oauth_info.access_token
+        return (
+            True,
+            token,
+            "Access granted to private repository.",
+        )
 
     except GatedRepoError:
         return (
             False,
+            None,
             "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"
+        return False, None, "You don't have access to the required repository"
 
     except Exception as e:
-        return False, f"Error checking access: {str(e)}"
+        return False, None, f"Error checking access: {str(e)}"
 
 
 reps = [