import os import sys import gradio as gr import html from tqdm import tqdm import torch from transformers import MBartForConditionalGeneration, AutoTokenizer, AutoModel, AutoModelForQuestionAnswering, AutoModelForTokenClassification, pipeline from torch import nn import torch.nn.functional as F from underthesea import word_tokenize device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") # Load multi task model bartpho_mt_base = MBartForConditionalGeneration.from_pretrained("mc0c0z/BARTPho-multi-task") bartpho_mt_base_tokenizer = AutoTokenizer.from_pretrained("mc0c0z/BARTPho-multi-task") bartpho_mt_base.to(device) bartpho_mt = MBartForConditionalGeneration.from_pretrained("mc0c0z/BARTPho-Large-multi-task") bartpho_mt_tokenizer = AutoTokenizer.from_pretrained("mc0c0z/BARTPho-Large-multi-task") bartpho_mt.to(device) def segmenter(text): text = html.unescape(text) tokens = word_tokenize(text) result = [] for token in tokens: if ' ' in token: result.append(token.replace(' ', '_')) else: result.append(token) return result class MultiTaskModel: def __init__(self, model, tokenizer, device): self.model = model self.tokenizer = tokenizer self.device = device def get_prompt(self, task): if task == 'sa': return "Classify the sentiment: " elif task == 'mt-en-vi': return "Translate English to Vietnamese: " elif task == 'mt-vi-en': return "Translate Vietnamese to English: " else: return "" def inference(self, task, sentence, device): # Tiền xử lý câu đầu vào tương tự như trong CustomDataset tokenized_text = segmenter(sentence) source = self.get_prompt(task) + " ".join(tokenized_text) # Tokenize input inputs = self.tokenizer(source, padding='max_length', truncation=True, max_length=128, return_tensors='pt') # Di chuyển input sang device input_ids = inputs["input_ids"].to(device) attention_mask = inputs["attention_mask"].to(device) # Sinh dự đoán self.model.eval() with torch.no_grad(): generated_output = self.model.generate(input_ids, attention_mask=attention_mask, max_length=128) # Giải mã dự đoán prediction = self.tokenizer.decode(generated_output[0], skip_special_tokens=True) if task == 'sa': class_names = ["Negative", "Positive"] return class_names[int(prediction[0])] return html.unescape(prediction) #Load SA model class CustomModel(nn.Module): def __init__(self, bert_model): super(CustomModel, self).__init__() self.bert = bert_model self.mlp = nn.Sequential( nn.Linear(768 * 5, 512), # 768*5 cho BERT nn.ReLU(), nn.Linear(512, 256), nn.ReLU(), nn.Linear(256, 3) # num_classes là số lượng lớp trong bài toán ) def forward(self, input_ids, attention_mask): outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) # Lấy 5 lớp ẩn cuối cùng của token [CLS] last_hidden_states = outputs.hidden_states[-5:] cls_embeddings = torch.cat([state[:, 0, :] for state in last_hidden_states], dim=1) # Đưa qua MLP logits = self.mlp(cls_embeddings) return logits ## PhoBERT phobert_sa = AutoModel.from_pretrained("vinai/phobert-base", output_hidden_states=True) phobert_sa_tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base") phobert_sa = CustomModel(phobert_sa) phobert_sa.load_state_dict(torch.load('sa_model\phobert_sentiment_analysis.pth', map_location=device)) phobert_sa.to(device) ## PhoBERTv2 phobertv2_sa = AutoModel.from_pretrained("vinai/phobert-base-v2", output_hidden_states=True) phobertv2_sa_tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base-v2") phobertv2_sa = CustomModel(phobertv2_sa) phobertv2_sa.load_state_dict(torch.load('sa_model\phobertv2_sentiment_analysis.pth', map_location=device)) phobertv2_sa.to(device) ## Multilingual BERT m_bert_sa = AutoModel.from_pretrained("google-bert/bert-base-multilingual-cased", output_hidden_states=True) m_bert_sa_tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-multilingual-cased") m_bert_sa = CustomModel(m_bert_sa) m_bert_sa.load_state_dict(torch.load('sa_model\\bert_model_sentiment_analysis.pth', map_location=device)) m_bert_sa.to(device) # Load Q&A model roberta_qa = AutoModelForQuestionAnswering.from_pretrained("HungLV2512/Vietnamese-QA-fine-tuned") roberta_qa_tokenizer = AutoTokenizer.from_pretrained("HungLV2512/Vietnamese-QA-fine-tuned") roberta_qa.to(device) # Load NER model label_map = { 'B-LOC': 0, 'B-MISC': 1, 'B-ORG': 2, 'B-PER': 3, 'I-LOC': 4, 'I-MISC': 5, 'I-ORG': 6, 'I-PER': 7, 'O': 8 } ## PhoBERT phobert_ner = AutoModelForTokenClassification.from_pretrained("DrRinS/NER-PhoBERT", num_labels=len(label_map)) phobert_ner_tokenizer = AutoTokenizer.from_pretrained("DrRinS/NER-PhoBERT") phobert_ner.to(device) ## PhoBERTv2 phobertv2_ner = AutoModelForTokenClassification.from_pretrained("DrRinS/NER-PhoBERTv2", num_labels=len(label_map)) phobertv2_ner_tokenizer = AutoTokenizer.from_pretrained("DrRinS/NER-PhoBERTv2") phobertv2_ner.to(device) ## Multilingual BERT m_bert_ner = AutoModelForTokenClassification.from_pretrained("DrRinS/NER_MultilingualBERT", num_labels=len(label_map)) m_bert_ner_tokenizer = AutoTokenizer.from_pretrained("DrRinS/NER_MultilingualBERT") m_bert_ner.to(device) # Inference function def sentiment_inference(model, tokenizer, text, device): # Segment the input text text = " ".join(segmenter(text)) # Tokenize the segmented text inputs = tokenizer( text, padding='max_length', truncation=True, max_length=128, return_tensors='pt' ) # Move inputs to the correct device input_ids = inputs['input_ids'].to(device) attention_mask = inputs['attention_mask'].to(device) # Ensure inputs have the correct shape input_ids = input_ids.unsqueeze(0) if input_ids.dim() == 1 else input_ids attention_mask = attention_mask.unsqueeze(0) if attention_mask.dim() == 1 else attention_mask # Perform inference model.eval() with torch.no_grad(): outputs = model(input_ids, attention_mask) _, preds = torch.max(outputs, dim=1) # Map predictions to class names class_names = ["Negative", "Positive", "Neutral"] return class_names[preds.cpu().item()] def multitask_inference(model, tokenizer, text, task, device): multitask_model = MultiTaskModel(model, tokenizer, device) return multitask_model.inference(task, text, device) def qa_inference(model, tokenizer, question, context, device): qa_pipeline = pipeline('question-answering', model=model, tokenizer=tokenizer) res = qa_pipeline(question=question, context=context) return res['answer'] def ner_inference(model, tokenizer, text, device): predictions = [] # Tokenize the segmented text inputs = tokenizer( text, padding='max_length', truncation=True, max_length=128, return_tensors='pt' ) # Move inputs to the correct device input_ids = inputs['input_ids'].to(device) attention_mask = inputs['attention_mask'].to(device) # Perform inference model.eval() with torch.no_grad(): outputs = model(input_ids, attention_mask) _, preds = torch.max(outputs.logits, dim=2) # Convert predictions to labels id_to_label = {v: k for k, v in label_map.items()} predictions = preds[attention_mask.bool()].cpu().numpy().flatten() labels = [id_to_label[p] for p in predictions] # Decode the input ids to tokens tokens = tokenizer.convert_ids_to_tokens(input_ids[0], skip_special_tokens=True) # Combine tokens with their NER labels ner_tags = list(zip(tokens, labels)) return ner_tags def process_input(input_text, context, task): results = {} if task == "Sentiment Analysis": results["PhoBERT"] = sentiment_inference(phobert_sa, phobert_sa_tokenizer, input_text, device) results["PhoBERTv2"] = sentiment_inference(phobertv2_sa, phobertv2_sa_tokenizer, input_text, device) results["Multilingual BERT"] = sentiment_inference(m_bert_sa, m_bert_sa_tokenizer, input_text, device) results["BARTPho Base"] = multitask_inference(bartpho_mt_base, bartpho_mt_base_tokenizer, input_text, "sa", device) results["BARTPho Large"] = multitask_inference(bartpho_mt, bartpho_mt_tokenizer, input_text, "sa", device) elif task == "English to Vietnamese": results["BARTPho Base"] = multitask_inference(bartpho_mt_base, bartpho_mt_base_tokenizer, input_text, "mt-en-vi", device) results["BARTPho Large"] = multitask_inference(bartpho_mt, bartpho_mt_tokenizer, input_text, "mt-en-vi", device) elif task == "Vietnamese to English": results["BARTPho Base"] = multitask_inference(bartpho_mt_base, bartpho_mt_base_tokenizer, input_text, "mt-vi-en", device) results["BARTPho Large"] = multitask_inference(bartpho_mt, bartpho_mt_tokenizer, input_text, "mt-vi-en", device) elif task == "Question Answering": results["RoBERTa"] = qa_inference(roberta_qa, roberta_qa_tokenizer, input_text, context, device) elif task == "Named Entity Recognition": results["PhoBERT"] = ner_inference(phobert_ner, phobert_ner_tokenizer, input_text, device) results["PhoBERTv2"] = ner_inference(phobertv2_ner, phobertv2_ner_tokenizer, input_text, device) results["Multilingual BERT"] = ner_inference(m_bert_ner, m_bert_ner_tokenizer, input_text, device) return results with gr.Blocks() as iface: gr.Markdown("# Multi-task NLP Demo") gr.Markdown("Perform sentiment analysis, machine translation, question answering, or named entity recognition using various models.") with gr.Row(): task = gr.Radio(["Sentiment Analysis", "Question Answering", "Named Entity Recognition", "English to Vietnamese", "Vietnamese to English"], label="Task") with gr.Row(): input_text = gr.Textbox(label="Input Text") context = gr.Textbox(label="Context", visible=False) output = gr.JSON(label="Results") submit = gr.Button("Submit") def on_task_change(task): if task == "Question Answering": return { input_text: gr.update(label="Question", visible=True), context: gr.update(visible=True) } else: return { input_text: gr.update(label="Input Text", visible=True), context: gr.update(visible=False) } task.change(on_task_change, task, [input_text, context]) submit.click( process_input, inputs=[input_text, context, task], outputs=output ) if __name__ == "__main__": iface.launch(share=True)