1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4model_id = "DAMI-Lab/ARI-8B"
5
6tokenizer = AutoTokenizer.from_pretrained(model_id)
7model = AutoModelForCausalLM.from_pretrained(
8 model_id,
9 torch_dtype="auto",
10 device_map="auto",
11)
12model.eval()
13
14system_prompt = '''### Task
15You are an expert in restoring damaged Hanja characters. Restore each [Dn] with exactly the original Hanja character. Each [Dn] corresponds to exactly one Hanja character.
16
17### Requirements
18Base your restoration on the document’s overall context and meaning rather than treating each damaged token in isolation.
19
20### Input & Output
21The input consists of the document itself, its metadata, and the related documents.
22While the related documents are omitted in the shot for conciseness, they are always present in the actual dataset.
23The output must follow the format: {"[Dn]": "the restored Hanja character for [Dn]"}
24
25### Example Input & Output
26[Example 1]
27**Input**
28The document was written at date: 7, month: 6, year: 1771 - 英祖 era.
29Input Document: "傳于[D1]興宗曰, 當自光明殿出, 承旨·侍衛, 來待于建禮門."
30**Output**
31{"[D1]":"李"}
32
33[Example 2]
34**Input**
35The document was written at date: 26, month: 8, year: 1824 - 純祖 era.
36Input Document: "[D1][D2]口傳政事, 副護軍單李紀淵."
37**Output**
38{"[D1]":"兵","[D2]":"曹"}
39
40[Example 3]
41**Input**
42The document was written at date: 30, month: 6, year: 1686 - 肅宗 era.
43Input Document: "府[D1]啓, 請[D2]禮·壽進·於[D3]·龍洞·明安公主房折受處[D4]査正事. 入啓."
44**Output**
45{"[D1]": "前", "[D2]": "明", "[D3]": "義", "[D4]": "一"}
46
47[Example 4]
48**Input**
49The document was written at date: 14, month: 7, year: 1654 - 孝宗 era.
50Input Document: "備[D1]記, 國[D2]難事, 而謀避[D3]免, 朝有逆黨而營救掩護, 最在人先, 身爲[D4][D5], 所爲如此, 他[D6][D7][D8]? 右議政具仁垕罷職."
51**Output**
52{"[D1]": "忘", "[D2]": "有", "[D3]": "辭", "[D4]": "大", "[D5]": "臣", "[D6]": "何", "[D7]": "足", "[D8]": "觀"}
53
54[Example 5]
55**Input**
56The document was written at date: 26, month: 4, year: 1656 - 孝宗 era.
57Input Document: "又啓曰, 卽者[D1][D2]官, 使差備[D3][D4], 以勅使之意, [D5][D6][D7]臣等, 使之[D8]待於[D9]宴廳, 臣等依其言, 卽[D10][D11]去, 則大通官四人, 一時[D12]來, 傳[D13]禮部咨文[D14]通於臣等曰, 今此咨文, 從速入啓, 明日內回報, 可也[D15][D16], 故咨文送于政院之意, 敢啓. 傳曰, 知道."
58**Output**
59{"[D1]": "大", "[D2]": "通", "[D3]": "譯", "[D4]": "官", "[D5]": "傳", "[D6]": "言", "[D7]": "於", "[D8]": "來", "[D9]": "西", "[D10]": "爲", "[D11]": "進", "[D12]": "出", "[D13]": "給", "[D14]": "一", "[D15]": "云", "[D16]": "云"}'''
60
61related_documents = [
62 "傳曰, 呈告工判·同敦寧許遞, 今日政差出.",
63 "傳曰, 在外敦寧都正·同敦寧, 許遞, 今日政差出.",
64 "傳曰, 呈告禮曹判書·同敦寧·兵曹參判·同成均許遞, 今日政差出.",
65]
66
67related_text = "\n".join(f"- {doc}" for doc in related_documents)
68
69user_prompt = f'''Use the following documents as references to accurately restore the input document.
70
71Related Documents:
72{related_text}
73
74The input document was written at date: 28, month: 7, year: 1887 - 高宗 era.
75
76Input Document: 傳曰, 同敦寧·[D1][D2]都正竝許遞, 今日政差出.
77'''
78
79messages = [
80 {"role": "system", "content": system_prompt},
81 {"role": "user", "content": user_prompt},
82]
83
84try:
85 prompt = tokenizer.apply_chat_template(
86 messages,
87 tokenize=False,
88 add_generation_prompt=True,
89 enable_thinking=False,
90 )
91except TypeError:
92 prompt = tokenizer.apply_chat_template(
93 messages,
94 tokenize=False,
95 add_generation_prompt=True,
96 )
97
98inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
99
100with torch.inference_mode():
101 output_ids = model.generate(
102 **inputs,
103 max_new_tokens=128,
104 do_sample=False,
105 )
106
107generated_ids = output_ids[0, inputs["input_ids"].shape[-1]:]
108response = tokenizer.decode(generated_ids, skip_special_tokens=True)
109print(response)