Refactor model and data downloading in app.py to include Hugging Face token support

app.py

@@ -102,18 +102,16 @@ def load_model_and_tokenizer():
 
             status_text.text("Loading fine-tuned model from Hugging Face...")
             progress_bar.progress(50)
-            # Try to download and load fine-tuned model from Hugging Face
             try:
-                # Download the checkpoint file from Hugging Face
                 from huggingface_hub import hf_hub_download
-                
+                hf_token = os.environ.get("HF_TOKEN")
                 status_text.text("⬇️ Downloading model from saketh11/ColiFormer...")
                 model_path = hf_hub_download(
                     repo_id="saketh11/ColiFormer",
                     filename="balanced_alm_finetune.ckpt",
-                    cache_dir="./hf_cache"
+                    cache_dir="./hf_cache",
+                    token=hf_token
                 )
-                
                 status_text.text("🔄 Loading downloaded model...")
                 st.session_state.model = load_model(
                     model_path=model_path,
@@ -127,10 +125,7 @@ def load_model_and_tokenizer():
                 status_text.text("Loading base model as fallback...")
                 st.session_state.model = BigBirdForMaskedLM.from_pretrained("adibvafa/CodonTransformer")
                 if isinstance(st.session_state.model, torch.nn.Module):
-                    if isinstance(st.session_state.model, torch.nn.Module):
-                        st.session_state.model.to(st.session_state.device)
-                    else:
-                        st.warning("Fallback model loaded is not a PyTorch module. Cannot move to device.")
+                    st.session_state.model = st.session_state.model.to(st.session_state.device)
                 else:
                     st.warning("Fallback model loaded is not a PyTorch module. Cannot move to device.")
                 st.session_state.model_type = "base"
@@ -145,17 +140,18 @@ def load_model_and_tokenizer():
 def download_reference_data():
     """Download and cache reference data from Hugging Face"""
     try:
-        # Download the processed genes file from Hugging Face
+        from huggingface_hub import hf_hub_download
+        hf_token = os.environ.get("HF_TOKEN")
         file_path = hf_hub_download(
             repo_id="saketh11/ColiFormer-Data",
             filename="ecoli_processed_genes.csv",
-            repo_type="dataset"
+            repo_type="dataset",
+            token=hf_token
         )
         df = pd.read_csv(file_path)
         return df['dna_sequence'].tolist()
     except Exception as e:
         st.warning(f"Could not download reference data from Hugging Face: {e}")
-        # Fallback to minimal sequences
         return [
             "ATGGCGAAAGCGCTGTATCGCGAAAGCGCTGTATCGCGAAAGCGCTGTATCGC",
             "ATGAAATTTATTTATTATTATAAATTTATTTATTATTATAAATTTATTTAT",
@@ -166,11 +162,13 @@ def download_reference_data():
 def download_tai_weights():
     """Download and cache tAI weights from Hugging Face"""
     try:
-        # Download the tAI weights file from Hugging Face
+        from huggingface_hub import hf_hub_download
+        hf_token = os.environ.get("HF_TOKEN")
         file_path = hf_hub_download(
             repo_id="saketh11/ColiFormer-Data",
             filename="organism_tai_weights.json",
-            repo_type="dataset"
+            repo_type="dataset",
+            token=hf_token
         )
         with open(file_path, 'r') as f:
             all_weights = json.load(f)