LLMGuard/get_new_edges.py at main · lvXiangwei/LLMGuard · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

import os
import json
import torch
from torch.utils.data import DataLoader
from utils.model import BertClassifier, BertClaInfModel
from transformers import AutoTokenizer, AutoModel

from utils.util import Dataset, get_raw_text_cora, get_response
from tqdm import tqdm


from torch_geometric.datasets import Planetoid
import torch_geometric.transforms as T
import os

def get_ood_id_links(dataname='cora', threshold=3):
    dataset, text = get_raw_text_cora(use_text=True, seed=0)
    ood_data = []
    ood_label = []
    ood_id_positive_edges = []
    ood_id_negative_edges = []
    file_lst = []
    for ood_idx, filename in enumerate(os.listdir("cora/data")):
        # import ipdb; ipdb.set_trace()
        with open(f'cora/data/{filename}', 'r') as f:
            string = f.read()
        parsed_data = json.loads(string)
        data = 'Title: ' + parsed_data['title'] + '\n' + \
                'Abstract: ' + parsed_data['abstract']
        ood_data.append(data)
        ood_label.append(parsed_data['topic'])
        positive_edges = []
        negative_edges = []
        for id_idx, strength in parsed_data['strength']:
            if strength >= threshold:
                positive_edges.append(id_idx)
            else:
                negative_edges.append(id_idx)
        ood_id_positive_edges.append(positive_edges)
        ood_id_negative_edges.append(negative_edges)

    return ood_data, text, ood_id_positive_edges, ood_id_negative_edges


def get_embs(ood_data: list, model='deberta', dataname='cora', device='cuda:4', n_labels=7, feat_shrink=64):

    if model == "deberta":
        base_path = "./saved_model/cora/deberta-base-seed0.ckpt"
        ckpt_path = './cora/plm'

        if os.path.exists('cora/ood_embs.pth'):
            return  torch.load(f'{dataname}/ood_embs.pth')

         # init model
        tokenizer = AutoTokenizer.from_pretrained(base_path)
        model = AutoModel.from_pretrained(base_path)
        model = BertClassifier(model, n_labels=n_labels, feat_shrink=feat_shrink).to(device)
        model.eval()

        # load state dict
        state_dict = torch.load(ckpt_path, map_location=device)
        model.load_state_dict(state_dict)
    else:
        raise ValueError("model not supported")

    X = tokenizer(ood_data, padding=True, truncation=True, max_length=512)
    dataset = Dataset(X)
    dataloader = DataLoader(dataset, batch_size=128, shuffle=False)
    ood_data_embs = []

    with torch.no_grad():  # 禁用梯度计算以节省内存
        for batch in tqdm(dataloader, desc="Processing batches"):
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            node_ids = batch['node_id']  # 原始索引

            # 获取模型的输出
            bert_outputs = model.bert_encoder(input_ids, attention_mask=attention_mask, return_dict=True, output_hidden_states=True)
            hidden = bert_outputs['hidden_states'][-1]  # 形状为 (batch_size, seq_len, hidden_dim)

            # 获取 [CLS] 的表示作为句子嵌入
            cls_token_emb = hidden.permute(1, 0, 2)[0]
            ood_data_embs.append(cls_token_emb.cpu())

    ood_embs = torch.cat(ood_data_embs)
    torch.save(ood_embs, f'{dataname}/ood_embs.pt')
    return ood_embs


def retrieve(embeddings, k=4):
    m, d = embeddings.shape
    dist_matrix = torch.cdist(embeddings, embeddings, p=2)  # 计算欧几里得距离
    dist_matrix.fill_diagonal_(float('inf'))  # 避免自匹配

    k_half = k // 2

    max_indices = torch.topk(dist_matrix, k_half, largest=True).indices  # (m, k/2)引
    min_indices = torch.topk(dist_matrix, k_half, largest=False).indices  # (m, k/2)

    def filter_mutual_pairs(indices):
        mutual_pairs = set()
        for i in range(m):
            for j in indices[i]:
                j = j.item()
                if i != j and (j,i) not in mutual_pairs:
                    mutual_pairs.add((i, j))
        return list(mutual_pairs)

    max_mutual_pairs = filter_mutual_pairs(max_indices)
    min_mutual_pairs = filter_mutual_pairs(min_indices)

    return max_mutual_pairs, min_mutual_pairs


def get_ood_ood_links(text, ood_embs, k=4):
    '''
    ood_embs: torch.Tensor, shape=(m, d), m为ood数据的数量，d为嵌入维度
    k: int, 用于k近邻搜索的参数
    '''
    # import ipdb; ipdb.set_trace()
    p_ood_ood_positive_pairs, p_ood_ood_negative_pairs = retrieve(ood_embs, k=k)
    prompt_template = '''
    Given the title and abstract of the two paper nodes, determine whether there is a potential citation relationship between them.
    If a citation relationship exists, respond with 'True'. Otherwise, respond with 'False'.

    Paper 1:
    {paper1}

    Paper 2:
    {paper2}
    '''

    ood_ood_positive = []
    ood_ood_negative = []

    for ood_idx1, ood_idx2 in tqdm(p_ood_ood_positive_pairs):
        prompt = prompt_template.format(paper1=text[ood_idx1], paper2=text[ood_idx2])
        response = get_response(prompt)
        response_txt = response.choices[0].message.content

        if response_txt.split()[0] == 'False':
            ood_ood_negative.append([ood_idx1, ood_idx2])
        else:
            ood_ood_positive.append([ood_idx1, ood_idx2])
    return ood_ood_positive, ood_ood_negative


if __name__ == '__main__':
    ood_data, ood_id_positive_edges, ood_id_negative_edges = get_ood_id_links(dataname='cora', threshold=3) # (539, 539)

    ood_embs = get_embs(ood_data, model='deberta', dataname='cora', device='cuda:7', n_labels=7, feat_shrink=64)
    ood_ood_positive, ood_ood_negative = get_ood_ood_links(ood_data, ood_embs, k=4)
    # save ood_id_positive_edges, ood_id_negative_edges, ood_ood_positive, ood_ood_negative in the form of Dataset, for following `pytorch` easily loading dataset
    torch.save(ood_id_positive_edges, 'cora/ood_id_positive_edges.pt')
    torch.save(ood_id_negative_edges, 'cora/ood_id_negative_edges.pt')
    torch.save(ood_ood_positive, 'cora/ood_ood_positive.pt')
    torch.save(ood_ood_negative, 'cora/ood_ood_negative.pt')