parallelize preprocessing across files via ProcessPoolExecutor
Replace serial file-read + tree-sitter parse loop with _preprocess_files() using ProcessPoolExecutor. Module-level _preprocess_file() for pickling. Covers all three callers: analyze_clusters, analyze_by_package, infer.
This commit is contained in:
parent
710da56916
commit
5906adfc95
1 changed files with 29 additions and 28 deletions
|
|
@ -204,6 +204,32 @@ def frequency_filter(sequences, min_coverage=0.2):
|
||||||
return filtered
|
return filtered
|
||||||
|
|
||||||
|
|
||||||
|
def _preprocess_file(fp):
|
||||||
|
"""Preprocess one file. Module-level for ProcessPoolExecutor."""
|
||||||
|
with open(fp) as f:
|
||||||
|
code = f.read()
|
||||||
|
sequences = []
|
||||||
|
for method_seq in preprocess_by_method(fp, code):
|
||||||
|
if method_seq:
|
||||||
|
sequences.append(method_seq)
|
||||||
|
return (fp, sequences)
|
||||||
|
|
||||||
|
|
||||||
|
def _preprocess_files(file_paths):
|
||||||
|
"""Preprocess multiple files in parallel."""
|
||||||
|
sequences = []
|
||||||
|
seq_files = []
|
||||||
|
n_workers = os.cpu_count()
|
||||||
|
with ProcessPoolExecutor(max_workers=n_workers) as ex:
|
||||||
|
futures = {ex.submit(_preprocess_file, fp): fp for fp in file_paths}
|
||||||
|
for f in as_completed(futures):
|
||||||
|
fp, method_seqs = f.result()
|
||||||
|
for seq in method_seqs:
|
||||||
|
sequences.append(seq)
|
||||||
|
seq_files.append(fp)
|
||||||
|
return sequences, seq_files
|
||||||
|
|
||||||
|
|
||||||
def analyze_clusters(file_paths, extension, project_root="", min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, include_kore=False):
|
def analyze_clusters(file_paths, extension, project_root="", min_coverage=DEFAULT_COVERAGE, prefer=None, kmax=2, N=3, include_kore=False):
|
||||||
"""Run full pipeline: preprocess → frequency filter → ensemble infer.
|
"""Run full pipeline: preprocess → frequency filter → ensemble infer.
|
||||||
|
|
||||||
|
|
@ -211,16 +237,7 @@ def analyze_clusters(file_paths, extension, project_root="", min_coverage=DEFAUL
|
||||||
list of (label, ensemble_result_dict, method_count, meta) tuples.
|
list of (label, ensemble_result_dict, method_count, meta) tuples.
|
||||||
meta = {"files": set(paths), "imports": [lines], "arg_patterns": {...}, "packages": [...]}.
|
meta = {"files": set(paths), "imports": [lines], "arg_patterns": {...}, "packages": [...]}.
|
||||||
"""
|
"""
|
||||||
sequences = []
|
sequences, seq_files = _preprocess_files(file_paths)
|
||||||
seq_files = []
|
|
||||||
for fp in file_paths:
|
|
||||||
with open(fp) as f:
|
|
||||||
code = f.read()
|
|
||||||
for method_seq in preprocess_by_method(fp, code):
|
|
||||||
if method_seq:
|
|
||||||
sequences.append(method_seq)
|
|
||||||
seq_files.append(fp)
|
|
||||||
|
|
||||||
if not sequences:
|
if not sequences:
|
||||||
return []
|
return []
|
||||||
|
|
||||||
|
|
@ -259,17 +276,8 @@ def analyze_by_package(file_paths, extension, project_root="", min_coverage=DEFA
|
||||||
Returns:
|
Returns:
|
||||||
list of (package_label, ensemble_result_dict, method_count, meta).
|
list of (package_label, ensemble_result_dict, method_count, meta).
|
||||||
"""
|
"""
|
||||||
sequences = []
|
|
||||||
seq_files = []
|
|
||||||
t0 = time.time()
|
t0 = time.time()
|
||||||
for fp in file_paths:
|
sequences, seq_files = _preprocess_files(file_paths)
|
||||||
with open(fp) as f:
|
|
||||||
code = f.read()
|
|
||||||
for method_seq in preprocess_by_method(fp, code):
|
|
||||||
if method_seq:
|
|
||||||
sequences.append(method_seq)
|
|
||||||
seq_files.append(fp)
|
|
||||||
|
|
||||||
if not sequences:
|
if not sequences:
|
||||||
return []
|
return []
|
||||||
_vprint(f"Preprocess: {len(sequences)} methods from {len(file_paths)} {extension} files ({time.time()-t0:.1f}s)")
|
_vprint(f"Preprocess: {len(sequences)} methods from {len(file_paths)} {extension} files ({time.time()-t0:.1f}s)")
|
||||||
|
|
@ -326,14 +334,7 @@ def infer(file_paths, extension, min_coverage=DEFAULT_COVERAGE, prefer=None, kma
|
||||||
Returns:
|
Returns:
|
||||||
Ensemble result dict from infer_ensemble.
|
Ensemble result dict from infer_ensemble.
|
||||||
"""
|
"""
|
||||||
sequences = []
|
sequences, _ = _preprocess_files(file_paths)
|
||||||
for fp in file_paths:
|
|
||||||
with open(fp) as f:
|
|
||||||
code = f.read()
|
|
||||||
for method_seq in preprocess_by_method(fp, code):
|
|
||||||
if method_seq:
|
|
||||||
sequences.append(method_seq)
|
|
||||||
|
|
||||||
sequences = frequency_filter(sequences, min_coverage=0.2)
|
sequences = frequency_filter(sequences, min_coverage=0.2)
|
||||||
|
|
||||||
symbol_seqs = [[text for _, text, _ in seq] for seq in sequences]
|
symbol_seqs = [[text for _, text, _ in seq] for seq in sequences]
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue