Upload 2 files

app.py

@@ -1,21 +1,11 @@
 import gradio as gr
 from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
-import pandas as pd
 import random
+import numpy as np
 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}
@@ -26,7 +16,6 @@ description = "It is a program that classifies whether it is positive or negativ
 
 
 
-
 def tokenized_data(tokenizer, inputs):
     return tokenizer.batch_encode_plus(
         [inputs],
@@ -37,12 +26,9 @@ def tokenized_data(tokenizer, inputs):
 
 
 
-
-examples_eng = ["the greatest musicians ", "cold movie "]
-examples_kor = ["긍정", "부정"]
-
 examples = []
 df = pd.read_csv('examples.csv', sep='\t', index_col='Unnamed: 0')
+random.seed(100)
 for i in range(2):
     idx = random.randint(0, 50)
     examples.append(['Eng', df.iloc[idx, 0]])
@@ -60,13 +46,13 @@ eng_model = AutoModelForSequenceClassification.from_pretrained(
 )
 
 
-kor_model_name = "klue_roberta-small"
+kor_model_name = "klue/roberta-small"
 kor_step = 2400
-kor_tokenizer = AutoTokenizer.from_pretrained(kor_model_name.replace('_', '/'))
-kor_file_name = "{}-{}.pt".format(kor_model_name, kor_step)
+kor_tokenizer = AutoTokenizer.from_pretrained(kor_model_name)
+kor_file_name = "{}-{}.pt".format(kor_model_name.replace('/', '_'), kor_step)
 kor_state_dict = torch.load(kor_file_name)
 kor_model = AutoModelForSequenceClassification.from_pretrained(
-    kor_model_name.replace('_', '/'), num_labels=2, id2label=id2label, label2id=label2id,
+    kor_model_name, num_labels=2, id2label=id2label, label2id=label2id,
     state_dict=kor_state_dict
 )
 
@@ -86,28 +72,31 @@ def builder(lang, text):
         logits = model(input_ids=inputs['input_ids'], 
             attention_mask=inputs['attention_mask']).logits
 
+
+    
+    m = torch.nn.Softmax(dim=1)
+    output = m(logits)
+    # print(logits, output)
+
     prediction = torch.argmax(logits, axis=1)
     
+    return {id2label[1]: output[0][1].item(), id2label[0]: output[0][0].item()}
     return id2label[prediction.item()]
 
 
-def builder2(inputs):
-    return eng_model(inputs)
 
-
-demo = gr.Interface(builder, inputs=[gr.inputs.Dropdown(['Eng', 'Kor']), "text"], outputs="text", 
+demo = gr.Interface(builder, inputs=[gr.inputs.Dropdown(['Eng', 'Kor']), "text"], 
+                            # outputs=gr.Label(num_top_classes=2), 
+                            outputs='label',
                             title=title, description=description, examples=examples)
 
-# demo2 = gr.Interface(builder2, inputs="text", outputs="text", 
-#                          title=title, theme="peach",
-#                          allow_flagging="auto",
-#                          description=description, examples=examples)
 
 # demo3 = gr.Interface.load("models/mdj1412/movie_review_score_discriminator_eng", inputs="text", outputs="text", 
 #                          title=title, theme="peach",
 #                          allow_flagging="auto",
 #                          description=description, examples=examples)
-    
+output = []
+
 if __name__ == "__main__":
     # print(examples)
     demo.launch()

requirements.txt

@@ -2,4 +2,5 @@ gradio
 datasets
 transformers
 torch
-pandas
+pandas
+numpy