Update app.py

app.py

@@ -1,6 +1,6 @@
 # app.py
 
-from transformers import AutoTokenizer, AutoModelForMaskedLM
+from bert_handler import create_handler_from_checkpoint
 import torch
 import gradio as gr
 import re
@@ -50,15 +50,10 @@ symbolic_roles = [
     "<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,
-    ignore_mismatched_sizes=True
-).eval().cuda()
+# Load from official hosted checkpoint
+checkpoint_path = "./bert-beatrix-2048"
+handler, model, tokenizer = create_handler_from_checkpoint(checkpoint_path)
+model = model.eval().cuda()
 
 def build_interface():
     with gr.Blocks() as demo: