Spaces:
Runtime error
Runtime error
| import torch | |
| from transformers import XLMRobertaTokenizer | |
| from model.xlmr_for_token_classification import XLMRForTokenClassification | |
| from utils.data_utils import NerProcessor | |
| import gradio as gr | |
| BACKGROUND_COLOR_MAP = { | |
| 'B-PER': '#FFDDDD', | |
| 'I-PER': '#FFAAAA', | |
| 'B-ORG': '#DDDDFF', | |
| 'I-ORG': '#AAAAFF', | |
| 'B-LOC': '#DDFFDD', | |
| 'I-LOC': '#AAFFAA', | |
| 'B-MISC': '#DDDDFF', | |
| 'I-MISC': '#AAAAFF', | |
| 'O': 'white' | |
| } | |
| 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 | |
| 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() | |
| 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' | |
| for token, label in zip(tokens, labels): | |
| if token in ["<s>", "</s>"]: | |
| continue | |
| if token.startswith("▁"): | |
| if current_word: | |
| result.append((current_word, current_label)) | |
| current_word = token[1:] | |
| current_label = label | |
| else: | |
| current_word += token | |
| # If the label is not 'O', update the current label to the new label | |
| if label != 'O': | |
| current_label = label | |
| if current_word: | |
| result.append((current_word, current_label)) | |
| return result | |
| def format_predictions(predictions): | |
| formatted_text = "<div style='font-size: 22px;'>" # Increase the font size for the entire text | |
| for word, label in predictions: | |
| if label == 'O': | |
| formatted_text += f'<span>{word}</span> ' | |
| else: | |
| color = BACKGROUND_COLOR_MAP.get(label, 'white') | |
| label_style = f'background-color: {color}; font-weight: bold;' | |
| word_style = f'background-color: {color}; font-size: 22px; padding: 2px;' # Increase font size for each word | |
| formatted_text += f'<span style="{word_style}">{word} <span style="{label_style}">[{label}]</span></span> ' | |
| formatted_text += "</div>" | |
| return formatted_text.strip() | |
| # 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) | |
| def ner(text): | |
| predictions = predict_ner(model, tokenizer, text, label_map, device) | |
| return format_predictions(predictions) | |
| # Define Gradio interface | |
| interface = gr.Interface( | |
| fn=ner, | |
| inputs=gr.Textbox(lines=2, placeholder="Enter text here..."), | |
| outputs=gr.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." | |
| ) | |
| if __name__ == "__main__": | |
| interface.launch() |