Spaces:
Sleeping
Sleeping
File size: 3,896 Bytes
0243138 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
# NER.py
from seqeval.metrics import classification_report, accuracy_score, f1_score, precision_score, recall_score
import torch
from transformers import AutoTokenizer, AutoModelForTokenClassification
import gradio as gr
import nltk
from nltk.tokenize import word_tokenize
# Download necessary NLTK data
nltk.download('punkt')
nltk.download('punkt_tab')
# Load the two models
model_id = "Swaraj66/Banglabert-finetuned-ner"
model_id2 = "Swaraj66/Finetuned_RemBERT"
banglabert_tokenizer = AutoTokenizer.from_pretrained(model_id)
banglabert_model = AutoModelForTokenClassification.from_pretrained(model_id)
rembert_tokenizer = AutoTokenizer.from_pretrained(model_id2)
rembert_model = AutoModelForTokenClassification.from_pretrained(model_id2)
# Helper functions
def get_word_logits(model, tokenizer, tokens):
encodings = tokenizer(tokens, is_split_into_words=True, return_tensors="pt", padding=True, truncation=True)
word_ids = encodings.word_ids()
with torch.no_grad():
logits = model(**encodings).logits
selected_logits = []
seen = set()
for idx, word_idx in enumerate(word_ids):
if word_idx is None:
continue
if word_idx not in seen:
selected_logits.append(logits[0, idx])
seen.add(word_idx)
return torch.stack(selected_logits)
def ensemble_predict(tokens):
rembert_logits = get_word_logits(rembert_model, rembert_tokenizer, tokens)
banglabert_logits = get_word_logits(banglabert_model, banglabert_tokenizer, tokens)
min_len = min(rembert_logits.shape[0], banglabert_logits.shape[0])
rembert_logits = rembert_logits[:min_len]
banglabert_logits = banglabert_logits[:min_len]
ensemble_logits = rembert_logits + banglabert_logits
preds = torch.argmax(ensemble_logits, dim=-1)
return preds.tolist()
# Label mapping
id2label = {
0: "O", 1: "B-PER", 2: "I-PER", 3: "B-ORG", 4: "I-ORG", 5: "B-LOC", 6: "I-LOC", 7: "B-MISC", 8: "I-MISC",
"0": "O", "1": "B-PER", "2": "I-PER", "3": "B-ORG", "4": "I-ORG", "5": "B-LOC", "6": "I-LOC", "7": "B-MISC", "8": "I-MISC"
}
# Main NER function
def ner_function(user_input):
words = word_tokenize(user_input)
preds = ensemble_predict(words)
pred_labels_list = [id2label[str(label)] for label in preds]
labeled_words = list(zip(words, pred_labels_list))
entities = []
current_entity = ""
current_label = None
for word, label in labeled_words:
if label.startswith("B-"):
if current_entity and current_label:
entities.append((current_entity.strip(), current_label))
current_entity = word
current_label = label[2:]
elif label.startswith("I-") and current_label == label[2:]:
current_entity += " " + word
else:
if current_entity and current_label:
entities.append((current_entity.strip(), current_label))
current_entity = ""
current_label = None
if current_entity and current_label:
entities.append((current_entity.strip(), current_label))
return entities
# Gradio UI
def build_ui():
with gr.Blocks() as demo:
gr.Markdown("# Named Entity Recognition App Using Ensemble Model (RemBERT + BanglaBERT)\nEnter a sentence to detect named entities.")
with gr.Row():
input_text = gr.Textbox(label="Enter a sentence", placeholder="Type your text here...")
with gr.Row():
submit_btn = gr.Button("Analyze Entities")
with gr.Row():
output_json = gr.JSON(label="Named Entities")
submit_btn.click(fn=ner_function, inputs=input_text, outputs=output_json)
return demo
# Launch the app
app = build_ui()
if __name__ == "__main__":
app.launch()
|