Update app.py

app.py

@@ -17,24 +17,20 @@ BACKGROUND_COLOR_MAP = {
 }
 
 def load_model(pretrained_path, model_path, device):
-    # Load the processor and label map
     processor = NerProcessor()
     label_list = processor.get_labels()
     label_map = {i: label for i, label in enumerate(label_list, 1)}
 
-    # Load the tokenizer
     tokenizer = XLMRobertaTokenizer.from_pretrained(pretrained_path)
 
-    # Initialize the model
     model = XLMRForTokenClassification(
         pretrained_path=pretrained_path,
         n_labels=len(label_list) + 1,
-        hidden_size=768,  # or 1024 depending on your model size
+        hidden_size=768,
         dropout_p=0.1,
         device=device
     )
 
-    # Load the model state
     model.load_state_dict(torch.load(model_path, map_location=device))
     model.to(device)
     model.eval()
@@ -42,23 +38,18 @@ def load_model(pretrained_path, model_path, device):
     return model, tokenizer, label_map
 
 def predict_ner(model, tokenizer, text, label_map, device):
-    # Tokenize input text
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
     input_ids = inputs["input_ids"].to(device)
     attention_mask = inputs["attention_mask"].to(device)
 
-    # Predict NER labels
     with torch.no_grad():
         logits = model(inputs_ids=input_ids, labels=None, labels_mask=None, valid_mask=None)
 
-    # Convert logits to label IDs
     label_ids = torch.argmax(logits, dim=2).detach().cpu().numpy()
 
-    # Convert label IDs to label names
     tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
     labels = [label_map[label_id] for label_id in label_ids[0]]
 
-    # Filter out unwanted tokens and remove prefix "▁"
     result = []
     current_word = ''
     current_label = 'O'
@@ -77,22 +68,21 @@ def predict_ner(model, tokenizer, text, label_map, device):
     if current_word:
         result.append((current_word, current_label))
 
-    # Remove punctuation from results
-    result = [(word, label) for word, label in result if word and word.isalnum()]
+    result = [(word, label) for word, label in result if word]
 
     return result
 
 def format_output(text, predictions):
     formatted_text = ""
     for word, label in predictions:
-        if label != 'O':
-            color = BACKGROUND_COLOR_MAP.get(label, 'white')
-            formatted_text += f"<span style='background-color:{color}; border:1px solid {color}; padding:2px;'>{word} <b>{label}</b></span> "
-        else:
+        if label == 'O':
             formatted_text += f"{word} "
-    return formatted_text.strip()
+        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} ({label})</span> "
+    
+    return formatted_text
 
-# Load the model and tokenizer
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model, tokenizer, label_map = load_model("./pretrained_models/xlmr.base", "./model.pt", device)
 
@@ -100,7 +90,6 @@ def ner(text):
     predictions = predict_ner(model, tokenizer, text, label_map, device)
     return format_output(text, predictions)
 
-# Define Gradio interface
 interface = gr.Interface(
     fn=ner,
     inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),