Update to 710M Character Model, add RTD

app.py

@@ -1,8 +1,9 @@
-from transformers import pipeline, AutoTokenizer
+from transformers import pipeline, AutoTokenizer, ElectraForPreTraining
 import pandas as pd
 import numpy as np
 import torch
 import streamlit as st
+from annotated_text import annotated_text
 
 USE_GPU = True
 
@@ -11,64 +12,127 @@ if USE_GPU and torch.cuda.is_available():
 else: 
     device = torch.device('cpu')
 
-MODEL_NAME_CHINESE = "IDEA-CCNL/Erlangshen-DeBERTa-v2-186M-Chinese-SentencePiece"
-#MODEL_NAME_CHINESE = "IDEA-CCNL/Erlangshen-DeBERTa-v2-97M-CWS-Chinese"
+MODEL_NAME_CHINESE = "IDEA-CCNL/Erlangshen-DeBERTa-v2-710M-Chinese"
+RTD_MODEL_NAME_CHINESE = "hfl/chinese-electra-180g-large-discriminator"
 
-
-WORD_PROBABILITY_THRESHOLD = 0.02
-TOP_K_WORDS = 200
-
-CHINESE_WORDLIST = ['一定','一样','不得了','主观','从此','便于','俗话','倒霉','候选','充沛','分别','反倒','只好','同情','吹捧','咳嗽','围绕','如意','实行','将近','就职','应该','归还','当面','忘记','急忙','恢复','悲哀','感冒','成长','截至','打架','把握','报告','抱怨','担保','拒绝','拜访','拥护','拳头','拼搏','损坏','接待','握手','揭发','攀登','显示','普遍','未免','欣赏','正式','比如','流浪','涂抹','深刻','演绎','留念','瞻仰','确保','稍微','立刻','精心','结算','罕见','访问','请示','责怪','起初','转达','辅导','过瘾','运动','连忙','适合','遭受','重叠','镇静']
+WORD_PROBABILITY_THRESHOLD = 0.05
+TOP_K_WORDS = 10
 
 @st.cache_resource
 def get_model_chinese():
     return pipeline("fill-mask", MODEL_NAME_CHINESE, device = device)
 
 @st.cache_resource
-def get_allowed_tokens():
-    df = pd.read_csv('allowed_token_ids.csv')
-    return set(list(df['token']))
+def get_rtd_tokenizer_chinese():
+    return AutoTokenizer.from_pretrained(RTD_MODEL_NAME_CHINESE)
+
+@st.cache_resource
+def get_rtd_model_chinese():
+    return ElectraForPreTraining.from_pretrained(RTD_MODEL_NAME_CHINESE)
+
+@st.cache_resource
+def get_wordlist_chinese():
+    df = pd.read_csv('wordlist_chinese_v2.csv')
+    wordlist = df[df.assess == True]
+    return wordlist['Chinese'].tolist()
+
+@st.cache_resource
+def get_allowed_words():
+    df = pd.read_csv('allowed_words.csv')
+    return set(list(df['word']))
 
 def assess_chinese(word, sentence):
     print("Assessing Chinese")
-    allowed_token_ids = get_allowed_tokens()
-    
+    number_of_chars = len(word)
+    assert number_of_chars == 2
+
+    allowed_words = get_allowed_words()    
     if sentence.lower().find(word.lower()) == -1:
         print('Sentence does not contain the word!')
         return
 
-    text = sentence.replace(word.lower(), "<mask>")
+    text = sentence.replace(word.lower(), "[MASK]"*number_of_chars)
 
-    top_k_prediction = mask_filler_chinese(text, top_k=TOP_K_WORDS)
-    target_word_prediction = mask_filler_chinese(text, targets = word)
+    top_k_prediction = []
+    candidates = mask_filler_chinese(text, top_k=TOP_K_WORDS)[0]
+    for candidate in candidates: 
+        temp_text = text.replace("[MASK]", candidate['token_str'], 1)    
+        second_predictions = mask_filler_chinese(temp_text, top_k=5)
+        for prediction in second_predictions:
+            prediction['token_str'] = candidate['token_str'] + prediction['token_str']
+            prediction['score'] = candidate['score'] * prediction['score']
+            
+        top_k_prediction.extend(second_predictions)
+    top_k_prediction = sorted(top_k_prediction, key = lambda x: x['score'], reverse = True)[:(TOP_K_WORDS*5)]
 
     norm_factor = 0
     for output in top_k_prediction: 
-        if output['token'] not in allowed_token_ids:
+        if output['token_str'] not in allowed_words:
             norm_factor += output['score']
 
     top_k_prediction_new = []
     for output in top_k_prediction:
-        if output['token'] in allowed_token_ids:        
+        if output['token_str'] in allowed_words:        
             output['score'] = output['score']/(1-min(0.5,norm_factor))
             top_k_prediction_new.append(output)
-
-    target_word_prediction[0]['score'] = target_word_prediction[0]['score'] / (1-min(0.5,norm_factor))
-    score = target_word_prediction[0]['score']
+    print (f"NORM_FACTOR: {norm_factor}")
+    
+    # Get target word prediction
+    temp_text = text
+    output1 = mask_filler_chinese(text, targets=word[0])[0][0]
+    temp_text = text.replace("[MASK]", word[0], 1)
+    output2 = mask_filler_chinese(temp_text, targets = word[1])[0]
+    output2['token_str'] = output1['token_str'] + output2['token_str']
+    output2['score'] = output1['score'] * output2['score']
+    target_word_prediction = output2
+
+    target_word_prediction['score'] = target_word_prediction['score'] / (1-min(0.5,norm_factor))
+    score = target_word_prediction['score']
 
     # append the original word if its not found in the results
     top_k_prediction_filtered = [output for output in top_k_prediction_new if \
                                  output['token_str'] == word]
     if len(top_k_prediction_filtered) == 0:
-        top_k_prediction_new.extend(target_word_prediction)
+        top_k_prediction_new.extend([target_word_prediction])
 
     return top_k_prediction_new, score
 
 def assess_sentence(word, sentence):
     return assess_chinese(word, sentence)
     
+def get_annotated_sentence(sentence, errors):
+    if len(errors) == 0:
+        return sentence
+
+    output = ["Input sentence: "]
+
+    wrong_char_indices = [e[0].item() for e in errors]
+    curr_ind = 0
+    for i in range(len(wrong_char_indices)):
+        output.append(sentence[curr_ind:wrong_char_indices[i]])
+        output.append((sentence[wrong_char_indices[i]], "", "#F8C8DC"))
+#        output.append((sentence[wrong_char_indices[i]], " ", "#ff4b4b"))
+        curr_ind = wrong_char_indices[i] + 1
+    output.append(sentence[curr_ind:])
+    print(output)
+
+    return output
+
+def get_word_errors(word, sentence):
+    tokens = rtd_tokenizer_chinese(sentence, return_tensors = 'pt', return_offsets_mapping = True)
+    scores = rtd_model_chinese(**rtd_tokenizer_chinese(sentence, return_tensors = 'pt'))[0][0]
+
+    errors = []
+    for i in range(len(scores)):
+        if scores[i] > 0:
+            errors.append(tokens['offset_mapping'][0][i])
+    
+    print(errors)
+    return errors
+
+
 def get_chinese_word():
-    possible_words = CHINESE_WORDLIST
+    possible_words = get_wordlist_chinese()
     word = np.random.choice(possible_words)
     return word
 
@@ -77,6 +141,8 @@ def get_word():
 
 mask_filler_chinese = get_model_chinese()
 #wordlist_chinese = get_wordlist_chinese()
+rtd_tokenizer_chinese = get_rtd_tokenizer_chinese()
+rtd_model_chinese = get_rtd_model_chinese()
 
 def highlight_given_word(row):
     color = '#ACE5EE' if row.Words == target_word else 'white'
@@ -101,14 +167,17 @@ def get_top_5_results(top_k_prediction):
     return top_5_df
 
 #### Streamlit Page
-st.title("造句 Auto-marking Demo")
+st.title("造句 Self-marking Demo")
 
 if 'target_word' not in st.session_state:
     st.session_state['target_word'] = get_word()
 target_word = st.session_state['target_word']
+target_word_ind = get_wordlist_chinese().index(target_word)
+
+#st.write("Target word: ", target_word)
+target_word = st.selectbox("Choose a word:", get_wordlist_chinese(), index = target_word_ind)
 
-st.write("Target word: ", target_word)
-if st.button("Get new word"):
+if st.button("Get random word"):
     st.session_state['target_word'] = get_word()
     st.experimental_rerun()
 
@@ -122,16 +191,23 @@ if st.button("Grade"):
     with open('./result01.json', 'w') as outfile:
         outfile.write(str(top_k_prediction))
 
-    st.write(f"Probability: {score:.2%}")
-    st.write(f"Target probability: {WORD_PROBABILITY_THRESHOLD:.2%}")
+    errors = get_word_errors(target_word, sentence)
+    annotated_sentence = get_annotated_sentence(sentence, errors)
+
+    annotated_text(annotated_sentence)
+
+    st.write(f"Probability score: {score:.1%}. (Target: {WORD_PROBABILITY_THRESHOLD:.1%})")
+ #   st.write(f"Target probability: {WORD_PROBABILITY_THRESHOLD:.1%}")
     predictions_df = get_top_5_results(top_k_prediction)
     df_style = predictions_df.style.apply(highlight_given_word, axis=1)
 
     if (score >= WORD_PROBABILITY_THRESHOLD):
 #        st.balloons()
-        st.success("Yay good job! 🕺 Practice again with other words", icon="✅")
-        st.table(df_style)
+        if (len(errors) == 0):
+            st.success("Yay good job! 🕺 Practice again with other words", icon="✅")
+        else: 
+            st.warning("Potential word errors detected. Try again?")
     else:
-        st.warning("Hmmm.. maybe try again?")
-        st.table(df_style)
+        st.warning("Probability score too low. Maybe try again?")
+    st.table(df_style)