Update app.py

app.py

@@ -1,6 +1,5 @@
 # app.py
 
-from huggingface_hub import snapshot_download
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 import torch
 import gradio as gr
@@ -9,9 +8,11 @@ from dataclasses import dataclass
 from pathlib import Path
 import spaces
 
+@spaces.GPU
 @dataclass
 class SymbolicConfig:
     repo_id: str = "AbstractPhil/bert-beatrix-2048"
+    revision: str = "main"
     symbolic_roles: list = (
         "<subject>", "<subject1>", "<subject2>", "<pose>", "<emotion>",
         "<surface>", "<lighting>", "<material>", "<accessory>", "<footwear>",
@@ -22,14 +23,15 @@ class SymbolicConfig:
     )
 
 config = SymbolicConfig()
-model_dir = snapshot_download(repo_id=config.repo_id)
-tokenizer = AutoTokenizer.from_pretrained(model_dir)
-model = AutoModelForMaskedLM.from_pretrained(model_dir, trust_remote_code=True).eval().cuda()
-
+tokenizer = AutoTokenizer.from_pretrained(config.repo_id, revision=config.revision)
+model = AutoModelForMaskedLM.from_pretrained(
+    config.repo_id,
+    revision=config.revision,
+    trust_remote_code=True
+).eval().cuda()
 
 MASK_TOKEN = tokenizer.mask_token or "[MASK]"
 
-@spaces.GPU
 def mask_and_predict(text: str, selected_roles: list[str]):
     results = []
     masked_text = text