Upload app.py

app.py

@@ -9,7 +9,6 @@ import pandas as pd
 import torch
 
 
-
 id2label = {0: "NEGATIVE", 1: "POSITIVE"}
 label2id = {"NEGATIVE": 0, "POSITIVE": 1}
 
@@ -78,7 +77,10 @@ kor_model = AutoModelForSequenceClassification.from_pretrained(
 
 def builder(lang, text):
     percent_kor, percent_eng = 0, 0
+    text_list = text.split(' ')
+
 
+    # [ output_1 ]
     if lang == 'Any':
         pred = LANGUAGE.predict_lang(text)
         if '__label__en' in pred[0]:
@@ -101,6 +103,7 @@ def builder(lang, text):
         if percent_kor==0: percent_kor=1
         
 
+    # [ output_2 ]
     inputs = tokenized_data(tokenizer, text)
     model.eval()
     with torch.no_grad():
@@ -111,25 +114,45 @@ def builder(lang, text):
     output = m(logits)
     # print(logits, output)
 
-    prediction = torch.argmax(logits, axis=1)
+
+    # [ output_3 ]
+    output_analysis = []
+    for word in text_list:
+        tokenized_word = tokenized_data(tokenizer, word)
+        with torch.no_grad():
+            logit = model(input_ids=tokenized_word['input_ids'], 
+                attention_mask=tokenized_word['attention_mask']).logits
+        word_output = m(logit)
+        if word_output[0][1] > 0.95:
+            output_analysis.append( (word, '+') )
+        elif word_output[0][1] < 0.05:
+            output_analysis.append( (word, '-') )
+        else:
+            output_analysis.append( (word, None) )
     
-    return [ {'Kor': percent_kor, 'Eng': percent_eng, 'Other': 1-(percent_kor+percent_eng)}, {id2label[1]: output[0][1].item(), id2label[0]: output[0][0].item()} ]
+    return [ {'Kor': percent_kor, 'Eng': percent_eng, 'Other': 1-(percent_kor+percent_eng)}, 
+            {id2label[1]: output[0][1].item(), id2label[0]: output[0][0].item()}, 
+            output_analysis ]
+            
+    # prediction = torch.argmax(logits, axis=1)
     return id2label[prediction.item()]
 
 
-
-demo = gr.Interface(builder, inputs=[gr.inputs.Dropdown(['Any', 'Eng', 'Kor']), "text"], 
-                            outputs=[ gr.Label(num_top_classes=3, label='Lang'), gr.Label(num_top_classes=2, label='Result') ], 
-                            # outputs='label',
-                            title=title, 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)
 
 
+
+demo = gr.Interface(builder, inputs=[gr.inputs.Dropdown(['Any', 'Eng', 'Kor']), "text"], 
+                    outputs=[ gr.Label(num_top_classes=3, label='Lang'), 
+                            gr.Label(num_top_classes=2, label='Result'),
+                            gr.HighlightedText(label="Analysis", combine_adjacent=False).style(color_map={"+": "red", "-": "green"}) ],
+                    # outputs='label',
+                    title=title, description=description, examples=examples)
+
+
 if __name__ == "__main__":
     # print(examples)
     demo.launch()