Update app.py

app.py

@@ -43,7 +43,7 @@ def predict_ner(model, tokenizer, text, label_map, device):
     attention_mask = inputs["attention_mask"].to(device)
 
     with torch.no_grad():
-        logits = model(input_ids=input_ids, attention_mask=attention_mask)[0]
+        logits = model(inputs_ids=input_ids, labels=None, labels_mask=None, valid_mask=None)
 
     label_ids = torch.argmax(logits, dim=2).detach().cpu().numpy()
 
@@ -73,7 +73,7 @@ def predict_ner(model, tokenizer, text, label_map, device):
     return result
 
 def format_output(text, predictions):
-    formatted_text = "<div style='border: 1px solid #ccc; padding: 10px; border-radius: 5px;'>"
+    formatted_text = ""
     i = 0
     for word, label in predictions:
         while i < len(text) and text[i].isspace():
@@ -86,11 +86,14 @@ def format_output(text, predictions):
             formatted_text += f"{word} "
         else:
             color = BACKGROUND_COLOR_MAP[label]
-            formatted_text += f"<span style='background-color:{color}; padding: 2px 5px; border: 1px solid {color}; font-weight: bold;'>{word} <span style='background-color:{color}; border: 1px solid {color};'>{label}</span></span> "
+            formatted_text += (
+                f"<span style='background-color:{color}; padding: 2px 5px; border: 1px solid {color}; font-weight: bold;'>"
+                f"{word} <span style='background-color:{color}; border: 1px solid {color}; padding: 2px; border-radius: 3px; font-weight: bold;'>"
+                f"{label}</span></span> "
+            )
         i += len(word)
     
-    formatted_text += "</div>"
-    return formatted_text
+    return f"<div style='border: 2px solid #cccccc; padding: 10px; border-radius: 5px;'>{formatted_text}</div>"
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model, tokenizer, label_map = load_model("./pretrained_models/xlmr.base", "./model.pt", device)