aa

app.py

@@ -1,24 +1,75 @@
-import streamlit as st
-
 from japanese.embedding import encode_sentences, get_cadidate_embeddings
 from japanese.tokenizer import extract_keyphrase_candidates
 from japanese.ranker import DirectedCentralityRnak
 
+import torch
 from transformers import AutoTokenizer
-from transformers import AutoModel
-
+from transformers import AutoModel, AutoModelForMaskedLM
 
-if __name__ == '__main__':
+def extract_keyphrase(text):
     # load model
     model = AutoModel.from_pretrained('cl-tohoku/bert-base-japanese')
     tokenizer = AutoTokenizer.from_pretrained('cl-tohoku/bert-base-japanese')
 
-    text = st.text_input("origin", "紀元前509年、第7代の王タルクィニウス・スペルブスを追放し共和制を敷いたローマだが、問題は山積していた。まず、王に代わった執政官（コンスル）が元老院の意向で決められるようになったこと、またその被選挙権が40歳以上に限定されていたことから、若い市民を中心としてタルクィニウスを王位に復する王政復古の企みが起こった。これは失敗して、初代執政官ルキウス・ユニウス・ブルトゥスは、彼自身の息子ティトゥスを含む陰謀への参加者を処刑した。ラテン同盟諸都市やエトルリア諸都市との同盟は、これらの都市とローマ王との同盟という形であったため、王の追放で当然に同盟は解消され、対立関係となった。")
     tokens, keyphrases = extract_keyphrase_candidates(text, tokenizer)
 
     document_embs = encode_sentences([tokens], tokenizer, model)
     document_feats = get_cadidate_embeddings([keyphrases], document_embs, [tokens])
     ranker = DirectedCentralityRnak(document_feats, beta=0.1, lambda1=1, lambda2=0.9, alpha=1.2, processors=8)
-    phrases = ranker.extract_summary()
-    st.write(phrases)
+    return ranker.extract_summary()[0]
+
+
+def preparation(tokenized_text, mask):
+    # [CLS],[SEP]の挿入
+    tokenized_text.insert(0, '[CLS]')  # 単語リストの先頭に[CLS]を付ける
+    tokenized_text.append('[SEP]')  # 単語リストの最後に[SEP]を付ける
+
+    maru = []
+    for i, word in enumerate(tokenized_text):
+        if word == '。' and i != len(tokenized_text) - 2:  # 「。」の位置検出
+            maru.append(i)
+
+    for i, loc in enumerate(maru):
+        tokenized_text.insert(loc + 1 + i, '[SEP]')  # 単語リストの「。」の次に[SEP]を挿入する
+
+    # 「□」を[MASK]に置き換え　
+    mask_index = []
+    for index, word in enumerate(tokenized_text):
+        if word == mask:  # 「□」の位置検出
+            tokenized_text[index] = '[MASK]'
+            mask_index.append(index)
+
+    return tokenized_text, mask_index
+
+
+def mask_prediction(text, mask_word):
+    model = AutoModelForMaskedLM.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
+    tokenizer = AutoTokenizer.from_pretrained('cl-tohoku/bert-base-japanese')
+
+    tokens, _ = extract_keyphrase_candidates(text, tokenizer)
+    tokenized_text = tokenizer.tokenize(text)
+    tokenized_text, mask_index = preparation(tokenized_text, mask_word)  # [CLS],[SEP],[MASK]の追加
+    tokens = tokenizer.convert_tokens_to_ids(tokenized_text)  # IDリストに変換
+    tokens_tensor = torch.tensor([tokens])  # IDテンソルに変換
+
+    model.eval()
+    with torch.no_grad():
+        outputs = model(tokens_tensor)
+        predictions = outputs[0]
+        for i in range(len(mask_index)):
+            _, predicted_indexes = torch.topk(predictions[0, mask_index[i]], k=5)
+            predicted_tokens = tokenizer.convert_ids_to_tokens(predicted_indexes.tolist())
+            return predicted_tokens
+
+
+if __name__ == '__main__':
+    text = st.text_input("origin", "ギリシア人ポリュビオスは，著書『歴史』の中で，ローマ共和政の国制（政治体制）を優れたものと評価している。彼によれば，その国制には，コンスルという王制的要素，元老院という共和制的要素，民衆という民主制的要素が存在しており，これら三者が互いに協調や牽制をしあって均衡しているというのである。ローマ人はこの政治体制を誇りとしており，それは，彼らが自らの国家を指して呼んだ「ローマの元老院と民衆」という名称からも読み取ることができる。")
+    phrases = extract_keyphrase(text)
+    for phrase in phrases:
+        for word in phrase.split("_"):
+            distracters = mask_prediction(text, word)
+            if distracters is None:
+                continue
+            for distracter in distracters:
+                st.write(text.replace(word, distracter))