Upload app.py

app.py

@@ -4,7 +4,6 @@ import fasttext
 from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
 
-import random
 import numpy as np
 import pandas as pd
 import torch
@@ -27,7 +26,7 @@ class LanguageIdentification:
         self.model = fasttext.load_model(pretrained_lang_model)
 
     def predict_lang(self, text):
-        predictions = self.model.predict(text, k=2) # returns top 2 matching languages
+        predictions = self.model.predict(text, k=200) # returns top 200 matching languages
         return predictions
 
 LANGUAGE = LanguageIdentification()
@@ -46,10 +45,13 @@ def tokenized_data(tokenizer, inputs):
 
 examples = []
 df = pd.read_csv('examples.csv', sep='\t', index_col='Unnamed: 0')
-random.seed(100)
-for i in range(15):
-    idx = random.randint(0, 50)
-    examples.extend([ ['Eng', df.iloc[idx, 0]], ['Kor', df.iloc[idx, 1]] ])
+np.random.seed(100)
+
+idx = np.random.choice(50, size=5, replace=False)
+eng_examples = [ ['Eng', df.iloc[i, 0]] for i in idx ]
+kor_examples = [ ['Kor', df.iloc[i, 1]] for i in idx ]
+examples = eng_examples + kor_examples
+
 
 
 eng_model_name = "roberta-base"
@@ -75,23 +77,31 @@ kor_model = AutoModelForSequenceClassification.from_pretrained(
 
 
 def builder(lang, text):
+    percent_kor, percent_eng = 0, 0
+
     if lang == 'Any':
         pred = LANGUAGE.predict_lang(text)
-        if pred[0][0] == '__label__ko':
-            lang = 'Kor'
-        else: # '__label__en'
+        if '__label__en' in pred[0]:
             lang = 'Eng'
-        # else:
-        #     raise NotImplementedError("It's neither Korean nor English.")
+            idx = pred[0].index('__label__en')
+            percent_eng = pred[1][idx]
+        if '__label__ko' in pred[0]:
+            lang = 'Kor'
+            idx = pred[0].index('__label__ko')
+            percent_kor = pred[1][idx]
+
     if lang == 'Eng':
         model = eng_model
         tokenizer = eng_tokenizer
+        if percent_eng==0: percent_eng=1
+
     if lang == 'Kor':
         model = kor_model
         tokenizer = kor_tokenizer
+        if percent_kor==0: percent_kor=1
         
+
     inputs = tokenized_data(tokenizer, text)
-    
     model.eval()
     with torch.no_grad():
         logits = model(input_ids=inputs['input_ids'], 
@@ -103,13 +113,13 @@ def builder(lang, text):
 
     prediction = torch.argmax(logits, axis=1)
     
-    return {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()} ]
     return id2label[prediction.item()]
 
 
 
 demo = gr.Interface(builder, inputs=[gr.inputs.Dropdown(['Any', 'Eng', 'Kor']), "text"], 
-                            outputs=gr.Label(num_top_classes=2, label='Result', color='CadetBlue'), 
+                            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)
 
@@ -119,6 +129,7 @@ demo = gr.Interface(builder, inputs=[gr.inputs.Dropdown(['Any', 'Eng', 'Kor']),
 #                          allow_flagging="auto",
 #                          description=description, examples=examples)
 
+
 if __name__ == "__main__":
     # print(examples)
     demo.launch()