kumarsushant36/multiLingual_Toxic_Text_Classification
import torch from torch.utils.data import DataLoader, TensorDataset from transformers import MBart50TokenizerFast, MBartForSequenceClassification import warnings warnings.filterwarnings('ignore')
device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu')
#more language can be added here getlangcode = { "hindi":"hiIN", "arabic":"arAR", "french":"frXX", "italian":"itIT", "korean":"koKR", "chinese":"zhCN", "english":"en_XX" }
modelname = "kumarsushant36/multiLingualToxicTextClassification" BartTokenizer = MBart50TokenizerFast.frompretrained(modelname) BartModel = MBartForSequenceClassification.frompretrained( modelname).to(device)
def predictuserinput(inputtext, lang, model=BartModel, tokenizer=BartTokenizer, device=device): userinput = [input_text]
tokenizer.srclang = getlangcode[lang] userencodings = tokenizer( userinput, truncation=True, padding=True, returntensors="pt")
userdataset = TensorDataset( userencodings['inputids'], userencodings['attention_mask'])
userloader = DataLoader(userdataset, batch_size=1, shuffle=False)
model.eval() with torch.nograd(): for batch in userloader: inputids, attentionmask = [t.to(device) for t in batch] outputs = model(inputids, attentionmask=attention_mask) logits = outputs.logits predictions = torch.sigmoid(logits)
predictedlabels = (predictions.cpu().numpy() > 0.5).astype(int) labelslist = ['toxic', 'severetoxic', 'obscene', 'threat', 'insult', 'identityhate'] result = dict(zip(labelslist, predictedlabels[0])) return result
text = "You are very bad" srclang = "english" result = predictuserinput(inputtext=text, lang=src_lang) print(result)
text = "너 정말 나쁘다" srclang = "korean" result = predictuserinput(inputtext=text, lang=src_lang) print(result)
output:
{'toxic': 1, 'severetoxic': 0, 'obscene': 0, 'threat': 0, 'insult': 0, 'identityhate': 0}
{'toxic': 1, 'severetoxic': 0, 'obscene': 0, 'threat': 0, 'insult': 0, 'identityhate': 0}
