|
import gradio as gr |
|
from transformers import AutoModelForSequenceClassification |
|
from transformers import AutoTokenizer |
|
import random |
|
import torch |
|
|
|
|
|
README = """ |
|
# Movie Review Score Discriminator |
|
It is a program that classifies whether it is positive or negative by entering movie reviews. |
|
You can choose between the Korean version and the English version. |
|
## Usage |
|
|
|
""" |
|
|
|
|
|
|
|
|
|
id2label = {0: "NEGATIVE", 1: "POSITIVE"} |
|
label2id = {"NEGATIVE": 0, "POSITIVE": 1} |
|
|
|
|
|
title = "Movie Review Score Discriminator" |
|
description = "It is a program that classifies whether it is positive or negative by entering movie reviews. You can choose between the Korean version and the English version." |
|
|
|
|
|
|
|
|
|
def tokenized_data(tokenizer, inputs): |
|
return tokenizer.batch_encode_plus( |
|
inputs, |
|
return_tensors="pt", |
|
padding="max_length", |
|
max_length=64, |
|
truncation=True) |
|
|
|
|
|
|
|
|
|
examples_eng = ["the greatest musicians ", "cold movie "] |
|
examples_kor = ["๊ธ์ ", "๋ถ์ "] |
|
|
|
examples = [] |
|
df = pd.read_csv('examples.csv', sep='\t', index_col='Unnamed: 0') |
|
for i in range(2): |
|
idx = random.randint(0, 50) |
|
examples.append(df.iloc[idx, 0]) |
|
examples.append(df.iloc[idx, 1]) |
|
|
|
|
|
|
|
|
|
model_kor = gr.Interface.load("models/cardiffnlp/twitter-roberta-base-sentiment") |
|
model_eng = gr.Interface.load("models/mdj1412/movie_review_score_discriminator_eng") |
|
|
|
|
|
|
|
def builder(version, inputs): |
|
if version == 'Eng': |
|
model_name = "roberta-base" |
|
step = 1900 |
|
|
|
else: |
|
model_name = "klue/roberta-small" |
|
step = 2400 |
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
inputs = tokenized_data(tokenizer, inputs) |
|
file_name = "{}-{}.pt".format(model_name, step) |
|
state_dict = torch.load(file_name) |
|
model = AutoModelForSequenceClassification.from_pretrained( |
|
model_name, num_labels=2, id2label=id2label, label2id=label2id, |
|
state_dict=state_dict |
|
) |
|
|
|
model.eval() |
|
with torch.no_grad(): |
|
logits = model(input_ids=inputs['input_ids'], |
|
attention_mask=inputs['attention_mask']).logits |
|
|
|
prediction = torch.argmax(logits, axis=1) |
|
|
|
return id2label[prediction.item()] |
|
|
|
|
|
def builder2(inputs): |
|
return model_eng(inputs) |
|
|
|
|
|
demo = gr.Interface(builder, inputs=[gr.inputs.Dropdown(['Eng', 'Kor']), "text"], outputs="text", |
|
title=title, description=description, examples=[examples]) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
demo.launch() |
|
|
