Spaces:
Sleeping
Sleeping
File size: 2,138 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 | import gzip
import pickle
from pathlib import Path
import numpy as np
import pandas as pd
import torch
from joblib import Parallel, delayed
from tqdm import tqdm
def process_and_save_chunk(filepath, output_dir, max_seq_len, min_mask_ratio):
with gzip.open(filepath, "rb") as f:
data = pickle.load(f) # type: ignore
x_list, y_list, meta_rows = [], [], []
chunk_name = filepath.name.replace(".pkl.gz", "")
for local_idx, item in enumerate(data):
features = item["features"]
mask = features[1]
if mask.mean() < min_mask_ratio or len(item["tokens"]) > max_seq_len:
continue
padded_tokens = np.zeros(max_seq_len, dtype=np.int64)
padded_tokens[: len(item["tokens"])] = item["tokens"]
x_list.append(features)
y_list.append(padded_tokens)
meta_rows.append(
{
"chunk_name": chunk_name,
"local_idx": local_idx,
"expr_str": item["expr_str"],
"expr_instantiated_str": item["expr_instantiated_str"],
}
)
if x_list:
x_tensor = torch.tensor(np.array(x_list), dtype=torch.float32)
y_tensor = torch.tensor(np.array(y_list), dtype=torch.long)
torch.save(
{"points": x_tensor, "tokens": y_tensor}, output_dir / f"{chunk_name}.pt"
)
return meta_rows
def run_compilation(
input_dir: Path, output_dir: Path, max_seq_len=128, min_mask_ratio=0.1, num_cores=4
):
output_dir.mkdir(parents=True, exist_ok=True)
files = sorted(input_dir.glob("chunk_*.pkl.gz"))
print("Компиляция PyTorch тензоров...")
results = Parallel(n_jobs=num_cores)(
delayed(process_and_save_chunk)(f, output_dir, max_seq_len, min_mask_ratio)
for f in tqdm(files)
)
all_meta_rows = [row for meta_chunk in results for row in meta_chunk] # type: ignore
df_meta = pd.DataFrame(all_meta_rows)
df_meta.to_csv(output_dir / "metadata.csv", index=False)
print(f"Готово! Сохранено графиков: {len(df_meta)}. Тензоры в {output_dir}")
|