Update app.py

app.py

@@ -4,35 +4,12 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM
 import torch
 import gradio as gr
 import re
-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>",
-        "<upper_body_clothing>", "<hair_style>", "<hair_length>", "<headwear>",
-        "<texture>", "<pattern>", "<grid>", "<zone>", "<offset>",
-        "<object_left>", "<object_right>", "<relation>", "<intent>", "<style>",
-        "<fabric>", "<jewelry>"
-    )
-
-config = SymbolicConfig()
-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]"
-
 def mask_and_predict(text: str, selected_roles: list[str]):
+    MASK_TOKEN = tokenizer.mask_token or "[MASK]"
     results = []
     masked_text = text
     token_ids = tokenizer.encode(text, return_tensors="pt").cuda()
@@ -64,16 +41,32 @@ def mask_and_predict(text: str, selected_roles: list[str]):
     accuracy = sum(1 for r in results if r["Match"] == "✅") / max(len(results), 1)
     return results, f"Accuracy: {accuracy:.1%}"
 
-def build_interface():
-    role_checkboxes = [gr.Checkbox(label=role, value=False) for role in config.symbolic_roles]
+symbolic_roles = [
+    "<subject>", "<subject1>", "<subject2>", "<pose>", "<emotion>",
+    "<surface>", "<lighting>", "<material>", "<accessory>", "<footwear>",
+    "<upper_body_clothing>", "<hair_style>", "<hair_length>", "<headwear>",
+    "<texture>", "<pattern>", "<grid>", "<zone>", "<offset>",
+    "<object_left>", "<object_right>", "<relation>", "<intent>", "<style>",
+    "<fabric>", "<jewelry>"
+]
+
+REPO_ID = "AbstractPhil/bert-beatrix-2048"
+REVISION = "main"
+tokenizer = AutoTokenizer.from_pretrained(REPO_ID, revision=REVISION)
+model = AutoModelForMaskedLM.from_pretrained(
+    REPO_ID,
+    revision=REVISION,
+    trust_remote_code=True
+).eval().cuda()
 
+def build_interface():
     with gr.Blocks() as demo:
         gr.Markdown("## 🔎 Symbolic BERT Inference Test")
         with gr.Row():
             with gr.Column():
                 input_text = gr.Textbox(label="Symbolic Input Caption", lines=3)
                 selected_roles = gr.CheckboxGroup(
-                    choices=config.symbolic_roles,
+                    choices=symbolic_roles,
                     label="Mask these symbolic roles"
                 )
                 run_btn = gr.Button("Run Mask Inference")
@@ -85,7 +78,6 @@ def build_interface():
 
     return demo
 
-
 if __name__ == "__main__":
     demo = build_interface()
     demo.launch()