Spaces:
Sleeping
Sleeping
File size: 5,811 Bytes
559e904 | 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 | import gzip
import pickle
import zlib
from pathlib import Path
import numpy as np
import torch
from joblib import Parallel, delayed
from tqdm import tqdm
def load_features_chunk(filepath):
with gzip.open(filepath, "rb") as f:
data = pickle.load(f) # type: ignore
features_list, metadata_list = [], []
for i, item in enumerate(data):
features_list.append(item["features"].flatten())
skel_hash = zlib.crc32(item["expr_str"].encode("utf-8"))
metadata_list.append(
(
filepath.name,
i,
len(item["tokens"]),
skel_hash,
item["expr_str"],
item["expr_instantiated_str"],
)
)
return np.array(features_list, dtype=np.float32), metadata_list
def run_collision_removal(
input_dir: Path, output_dir: Path, ckpt_path: Path, num_cores=4
):
output_dir.mkdir(parents=True, exist_ok=True)
files = sorted(input_dir.glob("chunk_*.pkl.gz"))
print("Загрузка данных для поиска коллизий...")
results = Parallel(n_jobs=num_cores)(
delayed(load_features_chunk)(f) for f in tqdm(files)
)
all_features, all_metadata = [], []
for feat, meta in results: # type: ignore
all_features.append(feat)
all_metadata.extend(meta)
features_matrix = np.vstack(all_features)
N = len(features_matrix)
print("Сортировка по Бритве Оккама...")
sorted_indices = sorted(
range(N), key=lambda i: (all_metadata[i][2], all_metadata[i][4])
)
features_matrix = features_matrix[sorted_indices]
all_metadata = [all_metadata[i] for i in sorted_indices]
hash_array = np.array([m[3] for m in all_metadata], dtype=np.int64)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
X_tensor = torch.tensor(features_matrix, device=device)
Hash_tensor = torch.tensor(hash_array, device=device)
active_mask = torch.ones(N, dtype=torch.bool, device=device)
X_means = X_tensor.mean(dim=1)
ALL_INDICES = torch.arange(N, device=device)
Q_BATCH, T_BATCH, CAND_BATCH = 1024, 65536, 10000
MAE_THRESHOLD, MAX_ERROR_THRESHOLD, MASK_MAE_TOLERANCE = 0.03, 0.03, 0.012
for start_idx in tqdm(range(0, N, Q_BATCH), desc="GPU Очистка коллизий"):
end_idx = min(start_idx + Q_BATCH, N)
q_indices = ALL_INDICES[start_idx:end_idx]
q_indices = q_indices[active_mask[start_idx:end_idx]]
if q_indices.numel() == 0:
continue
min_q = q_indices[0].item()
target_pool = ALL_INDICES[min_q + 1 :]
target_pool = target_pool[active_mask[min_q + 1 :]]
if target_pool.numel() == 0:
continue
Hash_q, Mean_q = Hash_tensor[q_indices], X_means[q_indices]
for t_chunk_start in range(0, target_pool.numel(), T_BATCH):
t_chunk_indices = target_pool[t_chunk_start : t_chunk_start + T_BATCH]
Hash_t, Mean_t = Hash_tensor[t_chunk_indices], X_means[t_chunk_indices]
mean_diffs = torch.abs(Mean_q.unsqueeze(1) - Mean_t.unsqueeze(0))
diff_skel_mask = Hash_q.unsqueeze(1) != Hash_t.unsqueeze(0)
candidate_mask = (mean_diffs <= MAE_THRESHOLD) & diff_skel_mask
rows, cols = candidate_mask.nonzero(as_tuple=True)
if rows.numel() > 0:
cand_global_q, cand_global_t = q_indices[rows], t_chunk_indices[cols]
for c_start in range(0, cand_global_q.numel(), CAND_BATCH):
c_end = c_start + CAND_BATCH
q_sub, t_sub = (
cand_global_q[c_start:c_end],
cand_global_t[c_start:c_end],
)
active_sub_mask = active_mask[q_sub] & active_mask[t_sub]
q_sub, t_sub = q_sub[active_sub_mask], t_sub[active_sub_mask]
if q_sub.numel() == 0:
continue
X_q_cand, X_t_cand = X_tensor[q_sub], X_tensor[t_sub]
y_q, mask_q = X_q_cand[:, :256], X_q_cand[:, 256:]
y_t, mask_t = X_t_cand[:, :256], X_t_cand[:, 256:]
mask_mae = torch.abs(mask_q - mask_t).mean(dim=1)
y_diffs = torch.abs(y_q - y_t)
y_mae, y_max_error = y_diffs.mean(dim=1), y_diffs.max(dim=1).values
hits = (
(mask_mae <= MASK_MAE_TOLERANCE)
& (y_mae <= MAE_THRESHOLD)
& (y_max_error <= MAX_ERROR_THRESHOLD)
)
targets_to_delete = t_sub[hits]
if targets_to_delete.numel() > 0:
active_mask[targets_to_delete] = False
indices_to_keep = active_mask.cpu().nonzero(as_tuple=True)[0].numpy()
print(f"Коллизий удалено: {N - len(indices_to_keep)}. Сохраняем результат...")
# Сохранение (воссоздаем чанки только с выжившими индексами)
from collections import defaultdict
keep_dict = defaultdict(set)
for global_idx in indices_to_keep:
file_name, local_idx = all_metadata[global_idx][0], all_metadata[global_idx][1]
keep_dict[file_name].add(local_idx)
for d in tqdm(files, desc="Запись чанков"):
if d.name not in keep_dict:
continue
with gzip.open(d, "rb") as file:
chunk = pickle.load(file)
clean_chunk = [item for i, item in enumerate(chunk) if i in keep_dict[d.name]]
with gzip.open(output_dir / d.name, "wb") as out_file:
pickle.dump(clean_chunk, out_file)
|