Update app.py

app.py

@@ -8,15 +8,11 @@ from transformers import LayoutLMTokenizerFast, LayoutLMForTokenClassification
 label2id = {"O":0, "COMPANY":1, "DATE":2, "TOTAL":3}
 id2label = {v:k for k,v in label2id.items()}
 
-# MODEL
-model = LayoutLMForTokenClassification.from_pretrained(
-    "microsoft/layoutlm-base-uncased",
-    num_labels=4
-)
+# ✅ USE YOUR MODEL
+MODEL_NAME = "YOUR_USERNAME/sroie-layoutlm"
 
-tokenizer = LayoutLMTokenizerFast.from_pretrained(
-    "microsoft/layoutlm-base-uncased"
-)
+model = LayoutLMForTokenClassification.from_pretrained(MODEL_NAME)
+tokenizer = LayoutLMTokenizerFast.from_pretrained(MODEL_NAME)
 
 # NORMALIZE
 def normalize(box, width, height):
@@ -53,7 +49,8 @@ def process(image):
         return_tensors="pt",
         padding="max_length",
         truncation=True,
-        is_split_into_words=True
+        is_split_into_words=True,
+        max_length=128
     )
 
     with torch.no_grad():