Update model selection for HF Spaces to use Large model

- Remove hardcoded fallback from Large to Medium model
- Simplify download_model_from_hf function to support all model sizes
- Add try-catch for model loading with automatic fallback to Medium if needed
- All Tranception models (Small, Medium, Large) are now available on HF Hub

app.py

@@ -42,19 +42,9 @@ def download_model_from_hf(model_name):
     """Download model from Hugging Face Hub if not present locally"""
     model_path = f"./{model_name}"
     if not os.path.exists(model_path):
-        print(f"Downloading {model_name} model...")
-        try:
-            # For Small and Medium models, they are available on HF Hub
-            if model_name in ["Tranception_Small", "Tranception_Medium"]:
-                return f"PascalNotin/{model_name}"
-            else:
-                # For Large model, we need to download from the original source
-                print("Note: Large model needs to be downloaded from the original source.")
-                print("Using Medium model as fallback...")
-                return "PascalNotin/Tranception_Medium"
-        except Exception as e:
-            print(f"Error downloading {model_name}: {e}")
-            return None
+        print(f"Loading {model_name} model from Hugging Face Hub...")
+        # All models are available on HF Hub
+        return f"PascalNotin/{model_name}"
     return model_path
 
 AA_vocab = "ACDEFGHIKLMNPQRSTVWY"
@@ -239,15 +229,21 @@ def score_and_create_matrix_all_singles(sequence,mutation_range_start=None,mutat
   assert mutation_range_end <= len(sequence), f"End position ({mutation_range_end}) exceeds sequence length ({len(sequence)})"
   
   # Load model with HF Space compatibility
-  if model_type=="Small":
-    model_path = download_model_from_hf("Tranception_Small")
-    model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(pretrained_model_name_or_path=model_path)
-  elif model_type=="Medium":
+  try:
+    if model_type=="Small":
+      model_path = download_model_from_hf("Tranception_Small")
+      model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(pretrained_model_name_or_path=model_path)
+    elif model_type=="Medium":
+      model_path = download_model_from_hf("Tranception_Medium")
+      model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(pretrained_model_name_or_path=model_path)
+    elif model_type=="Large":
+      model_path = download_model_from_hf("Tranception_Large")
+      model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(pretrained_model_name_or_path=model_path)
+  except Exception as e:
+    print(f"Error loading {model_type} model: {e}")
+    print("Falling back to Medium model...")
     model_path = download_model_from_hf("Tranception_Medium")
     model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(pretrained_model_name_or_path=model_path)
-  elif model_type=="Large":
-    model_path = download_model_from_hf("Tranception_Large")
-    model = tranception.model_pytorch.TranceptionLMHeadModel.from_pretrained(pretrained_model_name_or_path=model_path)
   
   # Device selection - for HF Spaces, typically CPU
   if torch.cuda.is_available():