Update app.py

app.py

@@ -1,9 +1,10 @@
 import torch
 from transformers import XLMRobertaTokenizer
 from model.xlmr_for_token_classification import XLMRForTokenClassification
-from utils.data_utils import NerProcessor
 import gradio as gr
+import os
 
+# Background color map for named entities
 BACKGROUND_COLOR_MAP = {
     'B-PER': '#FFDDDD',
     'I-PER': '#FFAAAA',
@@ -17,28 +18,24 @@ 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
+    
+    # Load the model
     model = XLMRForTokenClassification(
         pretrained_path=pretrained_path,
-        n_labels=len(label_list) + 1,
-        hidden_size=768,  # or 1024 depending on your model size
+        n_labels=6,  # Adjust number of labels based on your use case
+        hidden_size=768,  # Adjust based on your model size
         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()
-
+    
+    # Define label map
+    label_map = {0: 'O', 1: 'B-PER', 2: 'I-PER', 3: 'B-ORG', 4: 'I-ORG', 5: 'B-LOC', 6: 'I-LOC'}
+    
     return model, tokenizer, label_map
 
 def predict_ner(model, tokenizer, text, label_map, device):
@@ -49,54 +46,42 @@ def predict_ner(model, tokenizer, text, label_map, device):
 
     # Predict NER labels
     with torch.no_grad():
-        logits = model(input_ids=input_ids, labels=None, labels_mask=None, valid_mask=None)
-
+        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]]
-
-    # Build the result with formatting
-    formatted_result = ""
-    current_word = ''
-    current_label = 'O'
+    
+    # Generate highlighted text
+    highlighted_text = ""
     for token, label in zip(tokens, labels):
         if token in ["<s>", "</s>"]:
             continue
-        if token.startswith("▁"):
-            if current_word:
-                background_color = BACKGROUND_COLOR_MAP.get(current_label, 'white')
-                formatted_result += f'<span style="background-color: {background_color};">{current_word}</span> '
-            current_word = token[1:]
-            current_label = label
-        else:
-            current_word += token
-        if label != 'O':
-            current_label = label
-    if current_word:
-        background_color = BACKGROUND_COLOR_MAP.get(current_label, 'white')
-        formatted_result += f'<span style="background-color: {background_color};">{current_word}</span>'
-
-    return formatted_result
-
-# Load the model and tokenizer
+        token = token.replace("▁", " ")
+        color = BACKGROUND_COLOR_MAP.get(label, "white")
+        highlighted_text += f'<span style="background-color: {color}">{token}</span> '
+    
+    return highlighted_text
+
+# Load model and tokenizer
+pretrained_path = "./pretrained_models/xlmr.base"
+model_path = "./model.pt"
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model, tokenizer, label_map = load_model("./pretrained_models/xlmr.base", "./model.pt", device)
+model, tokenizer, label_map = load_model(pretrained_path, model_path, device)
 
+# Define Gradio interface
 def ner(text):
-    formatted_text = predict_ner(model, tokenizer, text, label_map, device)
-    return formatted_text
+    return predict_ner(model, tokenizer, text, label_map, device)
 
-# Define Gradio interface
-interface = gr.Interface(
+iface = gr.Interface(
     fn=ner,
-    inputs=gr.Textbox(lines=2, placeholder="Enter text here..."),
+    inputs="text",
     outputs="html",
-    title="Named Entity Recognition with XLM-RoBERTa",
-    description="This app uses a fine-tuned XLM-RoBERTa model for Named Entity Recognition (NER). Enter a text to see the NER predictions with background color highlights for named entities."
+    title="XLM-RoBERTa Named Entity Recognition",
+    description="Enter text to get Named Entity Recognition results with highlighted entities.",
 )
 
-if __name__ == "__main__":
-    interface.launch()
+iface.launch()