Views
No views yet
1from transformers import AutoTokenizer, AutoModelForMaskedLM
2from huggingface_hub import HfApi, login
3with open('./api_key/HGF_TOKEN.txt', 'r') as hgf:
4 login(token=hgf.read())
5api = HfApi()
6
7model_id = "x2bee/KoModernBERT-base-mlm-v01"
8tokenizer = AutoTokenizer.from_pretrained(model_id)
9model = AutoModelForMaskedLM.from_pretrained(model_id).to("cuda")
10
11def modern_bert_convert_with_multiple_masks(text: str, top_k: int = 1, select_method:str = "Logit") -> str:
12 if "[MASK]" not in text:
13 raise ValueError("MLM Model should include '[MASK]' in the sentence")
14
15 while "[MASK]" in text:
16 inputs = tokenizer(text, return_tensors="pt").to("cuda")
17 outputs = model(**inputs)
18
19 input_ids = inputs["input_ids"][0].tolist()
20 mask_indices = [i for i, token_id in enumerate(input_ids) if token_id == tokenizer.mask_token_id]
21
22 current_mask_index = mask_indices[0]
23
24 logits = outputs.logits[0, current_mask_index]
25
26 top_k_tokens = logits.topk(top_k).indices.tolist()
27 top_k_logits, top_k_indices = logits.topk(top_k)
28
29 if select_method == "Logit":
30 probabilities = torch.softmax(top_k_logits, dim=0).tolist()
31 predicted_token_id = random.choices(top_k_indices.tolist(), weights=probabilities, k=1)[0]
32 predicted_token = tokenizer.decode([predicted_token_id]).strip()
33
34 elif select_method == "Random":
35 predicted_token_id = random.choice(top_k_tokens)
36 predicted_token = tokenizer.decode([predicted_token_id]).strip()
37
38 elif select_method == "Best":
39 predicted_token_id = top_k_tokens[0]
40 predicted_token = tokenizer.decode([predicted_token_id]).strip()
41
42 else:
43 raise ValueError("select_method should be one of ['Logit', 'Random', 'Best']")
44
45 text = text.replace("[MASK]", predicted_token, 1)
46
47 print(f"Predicted: {predicted_token} | Current text: {text}")
48
49 return texttext = "30일 전남 무안국제[MASK] 활주로에 전날 발생한 제주항공 [MASK] 당시 기체가 [MASK]착륙하면서 강한 마찰로 생긴 흔적이 남아 있다. 이 참사로 [MASK]과 승무원 181명 중 179명이 숨지고 [MASK]는 형체를 알아볼 수 없이 [MASK]됐다. [MASK] 규모와 [MASK] 원인 등에 대해 다양한 [MASK]이 제기되고 있는 가운데 [MASK]에 설치된 [MASK](착륙 유도 안전시설)가 [MASK]를 키웠다는 [MASK]이 나오고 있다."
result = mbm.modern_bert_convert_with_multiple_masks(text, top_k=1)
'30일 전남 무안국제터미널 활주로에 전날 발생한 제주항공 사고 당시 기체가 무단착륙하면서 강한 마찰로 생긴 흔적이 남아 있다. 이 참사로 승객과 승무원 181명 중 179명이 숨지고 일부는 형체를 알아볼 수 없이 실종됐다. 사고 규모와 사고 원인 등에 대해 다양한 의혹이 제기되고 있는 가운데 기내에 설치된 ESC(착륙 유도 안전시설)가 사고를 키웠다는 주장이 나오고 있다.'text = "중국의 수도는 [MASK]이다"
result = mbm.modern_bert_convert_with_multiple_masks(text, top_k=1)
'중국의 수도는 베이징이다'
text = "일본의 수도는 [MASK]이다"
result = mbm.modern_bert_convert_with_multiple_masks(text, top_k=1)
'일본의 수도는 도쿄이다'
text = "대한민국의 가장 큰 도시는 [MASK]이다"
result = mbm.modern_bert_convert_with_multiple_masks(text, top_k=1)
'대한민국의 가장 큰 도시는 서울이다'