Views
No views yet
1import os
2import json
3import torch
4import torch.nn.functional as F
5from vllm import LLM
6from vllm.config import PoolerConfig
7from huggingface_hub import hf_hub_download
8
9model_path = 'CjangCjengh/WN-zh-vi-sim-v0.3-GPTQ-Int4'
10zh_text_path = 'zh.txt'
11vi_text_path = 'vi.txt'
12output_path = 'output.json'
13save_interval = 50
14device = 'cuda'
15cpu_offload_gb = 0
16
17lm_head_filename = 'yes_no_lm_head.pt'
18lm_head_path = hf_hub_download(repo_id=model_path, filename=lm_head_filename, local_dir='.')
19
20zh_idx = 0
21vi_idx = 0
22max_extra_lines = 5
23
24align_list = []
25if os.path.exists(output_path):
26 with open(output_path,'r',encoding='utf-8') as f:
27 align_list = json.load(f)
28 zh_idx = sum([i['zh'].count('\n')+1 for i in align_list if i['zh']])
29 vi_idx = sum([i['vi'].count('\n')+1 for i in align_list if i['vi']])
30
31lm_head = torch.load(lm_head_path)
32lm_head.to(device)
33
34llm = LLM(model=model_path, cpu_offload_gb=cpu_offload_gb, enforce_eager=True, task='embed', override_pooler_config=PoolerConfig(pooling_type='ALL'))
35
36zh_lines = open(zh_text_path,'r',encoding='utf-8').readlines()
37vi_lines = open(vi_text_path,'r',encoding='utf-8').readlines()
38zh_lines = [l.strip() for l in zh_lines]
39vi_lines = [l.strip() for l in vi_lines]
40
41
42def get_sim_score(src_text, tgt_text):
43 text = f'<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n下面的中文段落和越南语段落的内容是否完全对应，不存在缺漏？（回答Yes或No）\n\n中文：\n{src_text}\n\n越南语：{tgt_text}<|im_end|>\n<|im_start|>assistant\n'
44 outputs = llm.encode(text)
45 with torch.inference_mode():
46 hidden_states = outputs[0].outputs.data[-1].to(dtype=lm_head.dtype).to(device)
47 logits = torch.matmul(lm_head, hidden_states)
48 result = F.softmax(logits, dim=0).tolist()
49 return result[0]
50
51def generate_pairs():
52 visited = set()
53 size = 1
54 while True:
55 for x in range(size + 1):
56 for y in range(size + 1):
57 if (x, y) not in visited:
58 visited.add((x, y))
59 yield (x+1, y+1)
60 size += 1
61
62while zh_idx < len(zh_lines) and vi_idx < len(vi_lines):
63 for zh_i, vi_i in generate_pairs():
64 zh_text = ''.join(zh_lines[zh_idx:zh_idx+zh_i])
65 vi_text = ' '.join(vi_lines[vi_idx:vi_idx+vi_i])
66 score = get_sim_score(zh_text, vi_text)
67 if score > 0.5:
68 break
69 if zh_i == 1 and vi_i == 1:
70 continue
71 score = get_sim_score(zh_lines[zh_idx+zh_i-1], vi_lines[vi_idx+vi_i-1])
72 if score > 0.5:
73 zh_i -= 1
74 vi_i -= 1
75 break
76 if zh_i > max_extra_lines or vi_i > max_extra_lines:
77 end_flag = False
78 for zh_i, vi_i in generate_pairs():
79 if zh_i == 1 and vi_i == 1:
80 continue
81 for zh_i_offset, vi_i_offset in generate_pairs():
82 zh_start = zh_idx+zh_i-1
83 vi_start = vi_idx+vi_i-1
84 if zh_i+zh_i_offset > max_extra_lines and vi_i+vi_i_offset > max_extra_lines:
85 break
86 if zh_i+zh_i_offset > max_extra_lines or vi_i+vi_i_offset > max_extra_lines:
87 continue
88 zh_text = ''.join(zh_lines[zh_start:zh_start+zh_i_offset])
89 vi_text = ' '.join(vi_lines[vi_start:vi_start+vi_i_offset])
90 score = get_sim_score(zh_text, vi_text)
91 if score > 0.5:
92 zh_i -= 1
93 vi_i -= 1
94 end_flag = True
95 break
96 if end_flag:
97 break
98 if zh_i > max_extra_lines or vi_i > max_extra_lines:
99 with open(output_path,'w',encoding='utf-8') as f:
100 json.dump(align_list, f, ensure_ascii=False, indent=0)
101 raise Exception(f'Error! zh line No.{zh_idx+1} vi line No.{vi_idx+1}')
102
103 zh_text = '\n'.join(zh_lines[zh_idx:zh_idx+zh_i])
104 vi_text = '\n'.join(vi_lines[vi_idx:vi_idx+vi_i])
105 align_list.append({'zh':zh_text, 'vi':vi_text})
106
107 new_align = [list(range(zh_idx, zh_idx+zh_i)), list(range(vi_idx, vi_idx+vi_i))]
108 print(new_align)
109
110 if len(align_list) % save_interval == 0:
111 with open(output_path,'w',encoding='utf-8') as f:
112 json.dump(align_list, f, ensure_ascii=False, indent=0)
113
114 zh_idx += zh_i
115 vi_idx += vi_i
116
117with open(output_path,'w',encoding='utf-8') as f:
118 json.dump(align_list, f, ensure_ascii=False, indent=0)