Spaces:
Sleeping
Sleeping
| 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}.") | |