Spaces:
Sleeping
Sleeping
File size: 3,453 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 | import gzip
import pickle
from pathlib import Path
import sympy as sp
from sympy.core.cache import clear_cache
from tqdm import tqdm
from plot2eq.core.tokenizer import InvalidExpressionError, Tokenizer
def get_c_node():
return sp.Symbol("C", real=True)
def has_top_level_shift(expr: sp.Expr) -> bool:
if expr.is_Symbol and str(expr).startswith("C"):
return True
if expr.is_Add:
for arg in expr.args:
if has_top_level_shift(arg): # type: ignore
return True
return False
def replace_cos_unevaluated(expr):
if not expr.args:
return expr
new_args = [replace_cos_unevaluated(arg) for arg in expr.args]
current_func = expr.func
new_func = sp.sin if current_func == sp.cos else current_func
return new_func(*new_args, evaluate=False)
def clean_item(item, tokenizer: Tokenizer):
try:
tokens = item["tokens"]
expr = tokenizer.token_seq_to_expr(tokens)
if expr.has(sp.sin) and expr.has(sp.cos):
return None
processed_expr = replace_cos_unevaluated(expr) if expr.has(sp.cos) else expr
if not has_top_level_shift(processed_expr):
processed_expr = sp.Add(processed_expr, get_c_node(), evaluate=False)
clear_expr = tokenizer.canonicalize_tree_structure(processed_expr)
new_tokens = tokenizer.expr_to_token_seq(clear_expr)
return {
"expr_str": str(clear_expr),
"expr_instantiated_str": item["expr_instantiated_str"],
"tokens": new_tokens,
"points": item["points"],
}
except InvalidExpressionError:
return None
def run_canonicalization(input_dir: Path, output_dir: Path, chunk_size: int = 50000):
output_dir.mkdir(parents=True, exist_ok=True)
tokenizer = Tokenizer()
input_files = sorted(input_dir.glob("chunk_*.pkl.gz"))
if not input_files:
print("В папке не найдено чанков для обработки.")
return
cleaned_buffer = []
new_chunk_index = 0
total_processed, total_skipped = 0, 0
operations_count = 0
print(f"Найдено {len(input_files)} чанков. Начинаю канонизацию формул...")
for filepath in tqdm(input_files, desc="Канонизация"):
with gzip.open(filepath, "rb") as f:
data = pickle.load(f)
for item in data:
operations_count += 1
if operations_count % 10000 == 0:
clear_cache()
cleaned = clean_item(item, tokenizer)
if cleaned:
cleaned_buffer.append(cleaned)
else:
total_skipped += 1
if len(cleaned_buffer) >= chunk_size:
savepath = output_dir / f"chunk_{new_chunk_index}.pkl.gz"
with gzip.open(savepath, "wb") as f_out:
pickle.dump(cleaned_buffer[:chunk_size], f_out)
total_processed += chunk_size
cleaned_buffer = cleaned_buffer[chunk_size:]
new_chunk_index += 1
if cleaned_buffer:
savepath = output_dir / f"chunk_{new_chunk_index}.pkl.gz"
with gzip.open(savepath, "wb") as f_out:
pickle.dump(cleaned_buffer, f_out)
total_processed += len(cleaned_buffer)
print(f"Успешно очищено: {total_processed}. Пропущено ошибок: {total_skipped}.")
|