Spaces:
Sleeping
Sleeping
| import pip | |
| pip.main(['install', 'torch']) | |
| pip.main(['install', 'transformers']) | |
| import torch | |
| import torch.nn as nn | |
| import gradio as gr | |
| import transformers | |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
| def load_model(model_name): | |
| # model_name = "Unggi/hate_speech_classifier_KcElectra" | |
| # model | |
| model = AutoModelForSequenceClassification.from_pretrained(model_name) | |
| # tokenizer.. | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| return model, tokenizer | |
| def inference(prompt): | |
| model_name = "Unggi/hate_speech_classifier_KcElectra" #"Unggi/hate_speech_bert" | |
| model, tokenizer = load_model( | |
| model_name = model_name | |
| ) | |
| inputs = tokenizer( | |
| prompt, | |
| return_tensors="pt" | |
| ) | |
| with torch.no_grad(): | |
| logits = model(**inputs).logits | |
| # for binary classification | |
| sigmoid = nn.Sigmoid() | |
| bi_prob = sigmoid(logits) | |
| predicted_class_id = bi_prob.argmax().item() | |
| # μμ | |
| model.config.id2label = {0: 'CLEAN', 1: 'UNCLEAN'} | |
| class_id = model.config.id2label[predicted_class_id] | |
| return "class_id: " + str(class_id) + "\n" + "clean_prob: " + str(bi_prob[0][0].item()) + "\n" + "unclean_prob: " + str(bi_prob[0][1].item()) | |
| demo = gr.Interface( | |
| fn=inference, | |
| inputs="text", | |
| outputs="text", #return κ° | |
| ).launch() |