diff --git a/dsperse/src/analyzers/onnx_analyzer.py b/dsperse/src/analyzers/onnx_analyzer.py index 6a21dd9..77f409f 100644 --- a/dsperse/src/analyzers/onnx_analyzer.py +++ b/dsperse/src/analyzers/onnx_analyzer.py @@ -32,14 +32,13 @@ def analyze(self, save_path:str = None) -> Dict[str, Any]: """ # Extract model metadata graph = self.onnx_model.graph - # Create maps for initializers and value info initializer_map = {init.name: init for init in graph.initializer} # Build a comprehensive value_info map from the original full model - full_model_value_info_map = {vi.name: vi for vi in graph.value_info} - full_model_value_info_map.update({vi.name: vi for vi in graph.input}) - full_model_value_info_map.update({vi.name: vi for vi in graph.output}) + # full_model_value_info_map = {vi.name: vi for vi in graph.value_info} + # full_model_value_info_map.update({vi.name: vi for vi in graph.input}) + # full_model_value_info_map.update({vi.name: vi for vi in graph.output}) model_input_shape = self._get_model_input_shapes(graph, initializer_map) model_output_shape = self._get_model_output_shapes(graph) @@ -368,9 +367,9 @@ def generate_slices_metadata(self, model_metadata, slice_points, slices_paths, o # Get segment metadata segment_metadata = self._get_segment_metadata( - model_metadata, - segment_idx, - start_idx, + model_metadata, + segment_idx, + start_idx, end_idx, slice_path, output_dir @@ -464,7 +463,9 @@ def _get_segment_metadata(self, model_metadata, segment_idx, start_idx, end_idx, segment_shape = self._get_segment_shape(end_idx, model_metadata, start_idx, slice_path) - output_dir = os.path.join(os.path.dirname(output_dir), "slices", "segment_{}".format(segment_idx)) if output_dir else os.path.join(os.path.dirname(self.onnx_path), "slices", "segment_{}".format(segment_idx)) + output_dir = output_dir or os.path.join(os.path.dirname(self.onnx_path), "slices") + output_dir = os.path.join(output_dir, "segment_{}".format(segment_idx)) + os.makedirs(output_dir, exist_ok=True) segment_path = os.path.abspath(os.path.join(output_dir, f"segment_{segment_idx}.onnx")) @@ -500,7 +501,7 @@ def _get_segment_dependencies(self, model_metadata, start_idx, end_idx): for output in node_info['dependencies']['output']: output_map[output] = True - # Check inputs and add any missing to dependencies + # Check inputs and add any missing to dependencies for input_name in node_info['dependencies']['input']: if input_name not in output_map: if input_name not in segment_dependencies['input']: @@ -511,7 +512,7 @@ def _get_segment_dependencies(self, model_metadata, start_idx, end_idx): for output in output_map: if output not in segment_dependencies['input']: segment_dependencies['output'].append(output) - + # Filter input names to exclude weights and biases filtered_inputs = [] for input_name in segment_dependencies['input']: @@ -521,18 +522,18 @@ def _get_segment_dependencies(self, model_metadata, start_idx, end_idx): # Include model inputs and intermediate tensors if input_name in [inp.name for inp in self.onnx_model.graph.input] or input_name.startswith('/'): filtered_inputs.append(input_name) - + # If there are no inputs after filtering, include the first non-weight/bias input if not filtered_inputs: for input_name in segment_dependencies['input']: if not any(pattern in input_name.lower() for pattern in ["weight", "bias"]): filtered_inputs.append(input_name) break - + # If still no inputs, use the first input as a fallback if not filtered_inputs and segment_dependencies['input']: filtered_inputs.append(segment_dependencies['input'][0]) - + segment_dependencies['filtered_inputs'] = filtered_inputs return segment_dependencies diff --git a/dsperse/src/analyzers/runner_analyzer.py b/dsperse/src/analyzers/runner_analyzer.py index 8ed8e8a..5aa0f45 100644 --- a/dsperse/src/analyzers/runner_analyzer.py +++ b/dsperse/src/analyzers/runner_analyzer.py @@ -12,13 +12,13 @@ logger = logging.getLogger(__name__) class RunnerAnalyzer: - def __init__(self, model_directory): + def __init__(self, model_directory, slices_dir=None): """ Args: model_directory: Path to the model directory. """ self.model_directory = model_directory - self.slices_dir = Path(os.path.join(model_directory, "slices")).resolve() + self.slices_dir = Path(slices_dir or os.path.join(model_directory, "slices")).resolve() self.slices_metadata_path = self.slices_dir / "metadata.json" self.size_limit = 1000 * 1024 * 1024 # 1000MB diff --git a/dsperse/src/cli/base.py b/dsperse/src/cli/base.py index aaa5c8f..a13cd37 100644 --- a/dsperse/src/cli/base.py +++ b/dsperse/src/cli/base.py @@ -36,7 +36,7 @@ def _param_name_suggests_path(name: str) -> bool: if not name: return False name = name.lower() - for token in ("path", "dir", "file", "model", "slices", "output", "input", "run"): + for token in ("path", "dir", "file", "model", "slices", "run"): if token in name: return True return False @@ -57,11 +57,9 @@ def _looks_like_path(value: str) -> bool: return False -def _maybe_normalize_from_prompt(param_name: str, prompt_message: str, value: str) -> str: +def _maybe_normalize_from_prompt(param_name: str, value: str) -> str: try: - if _param_name_suggests_path(param_name) or _looks_like_path(value) or ( - prompt_message and any(t in prompt_message.lower() for t in ["path", "directory", "dir", "file"]) # heuristic - ): + if _param_name_suggests_path(param_name) or _looks_like_path(value): return normalize_path(value) except Exception: pass @@ -111,7 +109,10 @@ def configure_logging(log_level='WARNING'): "The answer to life, the universe, and everything is... 42 (but you need a neural network to understand why).", "Neural networks don't actually think. They just do math really fast.", "If you're reading this, you're awesome! Keep up the great work!", - "Dsperse: Making neural networks more transparent, one slice at a time." + "Dsperse: Making neural networks more transparent, one slice at a time.", + "Remember: With great power comes great responsibility (and large models).", + "Keep calm and slice on!", + "Why did the neural network go to school? To improve its 'weights'!", ] def print_header(): @@ -262,7 +263,7 @@ def prompt_for_value(param_name, prompt_message, default=None, required=True): logger.debug(f"Using default run name for {param_name}: {default}") return str(default) else: - normalized_default = _maybe_normalize_from_prompt(param_name, prompt_message, str(default)) + normalized_default = _maybe_normalize_from_prompt(param_name, str(default)) logger.debug(f"Using default value for {param_name}: {normalized_default}") return normalized_default value = user_input.strip().strip('\'"') # Strip surrounding quotes @@ -271,7 +272,7 @@ def prompt_for_value(param_name, prompt_message, default=None, required=True): logger.debug(f"User provided run name for {param_name}: {value}") return value else: - value = _maybe_normalize_from_prompt(param_name, prompt_message, value) + value = _maybe_normalize_from_prompt(param_name, value) logger.debug(f"User provided value for {param_name}: {value}") return value else: @@ -280,7 +281,7 @@ def prompt_for_value(param_name, prompt_message, default=None, required=True): if user_input.strip() or not required: if user_input.strip(): value = user_input.strip().strip('\'"') # Strip surrounding quotes - value = _maybe_normalize_from_prompt(param_name, prompt_message, value) + value = _maybe_normalize_from_prompt(param_name, value) logger.debug(f"User provided value for {param_name}: {value}") return value else: diff --git a/dsperse/src/cli/full_run.py b/dsperse/src/cli/full_run.py index ed840f0..487c6a3 100644 --- a/dsperse/src/cli/full_run.py +++ b/dsperse/src/cli/full_run.py @@ -33,7 +33,7 @@ def setup_parser(subparsers): help='Path to the model file (.onnx) or directory containing the model') full_run_parser.add_argument('--input-file', '--input', '--if', '-i', dest='input_file', help='Path to input file for inference and compilation calibration (e.g., input.json)') - full_run_parser.add_argument('--slices-dir', '--slices-directory', '--slices-directroy', '--sd', '-s', dest='slices_dir', + full_run_parser.add_argument('--slices-dir', '--slices-path', '--slices-directory', '--slices-directory', '--sd', '-s', dest='slices_dir', help='Optional: Pre-existing slices directory to reuse (skips slicing step)') full_run_parser.add_argument('--layers', '-l', help='Optional: Layers to compile (e.g., "3, 20-22") passed through to compile') # Optional: allow non-interactive mode later if desired; kept interactive by default @@ -72,9 +72,10 @@ def full_run(args): using_builtin = False builtin_name = None + canonical_model_dir = None # 1) Resolve inputs interactively - if (not hasattr(args, 'model_dir') or not args.model_dir) and (not hasattr(args, 'input_file') or not args.input_file): + if not getattr(args, 'model_dir', None) and not getattr(args, 'input_file', None): # Special prompt that accepts either a filesystem location or a built-in token choice = prompt_for_value( 'selection', @@ -96,22 +97,22 @@ def full_run(args): # Set args to point to sources args.model_dir = model_onnx args.input_file = input_json - # Output root under user's home - output_root = os.path.expanduser(os.path.join('~', 'dsperse', builtin_name)) - os.makedirs(output_root, exist_ok=True) - # For downstream steps, canonical model dir should be the output root - canonical_model_dir = normalize_path(output_root) + # For downstream steps, canonical model dir should be the output root. Make one under user's home + canonical_model_dir = normalize_path(os.path.join('~', 'dsperse', builtin_name)) + os.makedirs(canonical_model_dir, exist_ok=True) print(f"{Fore.CYAN}Using built-in model '{builtin_name}'. Outputs will be saved under {canonical_model_dir}{Style.RESET_ALL}") else: # Treat as a user-provided file or directory path args.model_dir = normalize_path(choice) - canonical_model_dir = _determine_model_dir(args.model_dir) + elif getattr(args, 'model_dir', None): + # If model_dir provided - normalize provided values + args.model_dir = normalize_path(args.model_dir) else: - # Normalize provided values - if hasattr(args, 'model_dir') and args.model_dir: - args.model_dir = normalize_path(args.model_dir) - # Determine canonical model directory for downstream steps - canonical_model_dir = _determine_model_dir(args.model_dir) + # If only input_file provided + args.model_dir = os.path.dirname(normalize_path(args.input_file)) + + # Determine canonical model directory for downstream steps + canonical_model_dir = canonical_model_dir or _determine_model_dir(args.model_dir) # Input file resolution if hasattr(args, 'input_file') and args.input_file: @@ -122,30 +123,24 @@ def full_run(args): args.input_file = prompt_for_value('input-file', 'Enter the input file', default=default_input, required=True) args.input_file = normalize_path(args.input_file) if args.input_file else args.input_file - # If user provided an existing slices directory, skip slicing step + # 2) Slice (unless slices-dir provided) slices_dir = None if hasattr(args, 'slices_dir') and args.slices_dir: + # If user provided an existing slices directory, skip slicing step slices_dir = normalize_path(args.slices_dir) - - # 2) Slice (unless slices-dir provided) - if not slices_dir: + print(f"{Fore.YELLOW}Skipping slicing step, using existing slices at: {slices_dir}{Style.RESET_ALL}") + else: # Default slices dir depends on whether we're using a built-in selection - default_slices_dir = os.path.join(canonical_model_dir, 'slices') + slices_dir = os.path.join(canonical_model_dir, 'slices') analysis_dir = os.path.join(canonical_model_dir, 'analysis') - try: - os.makedirs(default_slices_dir, exist_ok=True) - os.makedirs(analysis_dir, exist_ok=True) - except Exception: - pass + os.makedirs(slices_dir, exist_ok=True) + os.makedirs(analysis_dir, exist_ok=True) # Call existing slice command; keep its logic and interactivity. # For built-ins, we point the slicer to the built-in model file but output to ~/dsperse/{name}/slices model_metadata_path = os.path.join(analysis_dir, 'model_metadata.json') - slice_args = Namespace(model_dir=args.model_dir, output_dir=default_slices_dir, save_file=model_metadata_path) + slice_args = Namespace(model_dir=args.model_dir, output_dir=slices_dir, save_file=model_metadata_path) print(f"{Fore.CYAN}Step 1/5: Slicing model...{Style.RESET_ALL}") slice_model(slice_args) - slices_dir = default_slices_dir - else: - print(f"{Fore.YELLOW}Skipping slicing step, using existing slices at: {slices_dir}{Style.RESET_ALL}") # 3) Compile (circuitize) with calibration input compile_args = Namespace(slices_path=slices_dir, input_file=args.input_file, layers=getattr(args, 'layers', None)) @@ -154,10 +149,8 @@ def full_run(args): # 4) Run inference run_root_dir = os.path.join(canonical_model_dir, 'run') - try: - os.makedirs(run_root_dir, exist_ok=True) - except Exception: - pass + os.makedirs(run_root_dir, exist_ok=True) + inference_output_path = os.path.join(run_root_dir, 'inference_results.json') run_args = Namespace(slices_dir=slices_dir, run_metadata_path=None, input_file=args.input_file, output_file=inference_output_path) print(f"{Fore.CYAN}Step 3/5: Running inference over slices...{Style.RESET_ALL}") diff --git a/dsperse/src/cli/prove.py b/dsperse/src/cli/prove.py index aba2c72..c47a079 100644 --- a/dsperse/src/cli/prove.py +++ b/dsperse/src/cli/prove.py @@ -126,18 +126,20 @@ def is_run_root_dir(p): default_model_path = os.path.dirname(os.path.dirname(latest_run_path)) # Prompt with default if found - if default_run: - candidate = prompt_for_value('run-or-run-id-dir', 'Enter run directory (runs root or a run_* directory)', default=default_run) - else: - candidate = prompt_for_value('run-or-run-id-dir', 'Enter run directory (runs root or a run_* directory)') + candidate = prompt_for_value( + 'run-or-run-id-dir', + 'Enter run directory (runs root or a run_* directory)', + default=default_run or None + ) + # Handle run names (starts with "run_") - prepend run/ directory BEFORE normalization - if candidate and candidate.startswith('run_') and not candidate.startswith('/') and not candidate.startswith('./') and not candidate.startswith('../'): + if candidate and candidate.startswith('run_'): # Always try current directory's run/ first (for when running from model directory) current_run_dir = os.path.join(os.getcwd(), "run") if os.path.exists(current_run_dir): candidate = os.path.join(current_run_dir, candidate) - elif 'default_model_path' in locals() and default_model_path and default_model_path != os.getcwd(): + elif default_model_path and default_model_path != os.getcwd(): # Use stored default model path if different from current directory model_run_dir = os.path.join(default_model_path, "run") candidate = os.path.join(model_run_dir, candidate) @@ -149,6 +151,8 @@ def is_run_root_dir(p): model_path = os.path.join(models_dir, model_name) if os.path.isdir(model_path): model_run_dir = os.path.join(model_path, "run") + # XXX: so we just find any model dir with `run` in it and just use that? Seems brittle. + # What if multiple models have runs? That would lead to really confusing behavior. if os.path.exists(model_run_dir) and os.path.exists(os.path.join(model_run_dir, candidate)): candidate = os.path.join(model_run_dir, candidate) break @@ -156,7 +160,6 @@ def is_run_root_dir(p): elif candidate and candidate.startswith('/') and os.path.basename(candidate).startswith('run_'): # Check if this is a run name that was normalized to the wrong directory basename = os.path.basename(candidate) - dirname = os.path.dirname(candidate) # If the directory doesn't exist but we have model directories, look there if not os.path.exists(candidate): diff --git a/dsperse/src/cli/run.py b/dsperse/src/cli/run.py index b769930..da30f6e 100644 --- a/dsperse/src/cli/run.py +++ b/dsperse/src/cli/run.py @@ -27,7 +27,7 @@ def setup_parser(subparsers): run_parser.set_defaults(command='run') # Arguments with aliases/shorthands - run_parser.add_argument('--slices-dir', '--slices-directory', '--slices', '--sd', '-s', dest='slices_dir', + run_parser.add_argument('--slices-path', '--slices-dir', '--slices-directory', '--slices', '--sd', '-s', dest='slices_dir', help='Directory containing the slices') run_parser.add_argument('--run-metadata-path', help='Path to run metadata.json (auto-generated if not provided)') run_parser.add_argument('--input-file', '--input', '--if', '-i', dest='input_file', @@ -40,7 +40,7 @@ def setup_parser(subparsers): def run_inference(args): """ Run inference on a model based on the provided arguments. - + This command requires a slices directory. The parent directory of the slices is treated as the model directory, which is used for defaults like input/output paths. @@ -60,28 +60,24 @@ def run_inference(args): return # Validate slices directory has metadata and normalize to the actual slices directory - meta_in_dir = os.path.exists(os.path.join(args.slices_dir, 'metadata.json')) - meta_in_sub = os.path.exists(os.path.join(args.slices_dir, 'slices', 'metadata.json')) - - if not (meta_in_dir or meta_in_sub): - print(f"{Fore.YELLOW}Warning: No slices metadata found at the provided path. Please slice the model first.{Style.RESET_ALL}") - logger.error("Run requires a valid slices directory with metadata.json") - return - - if meta_in_dir: + if os.path.exists(os.path.join(args.slices_dir, 'metadata.json')): slices_dir_effective = args.slices_dir model_dir = os.path.dirname(args.slices_dir.rstrip('/')) or '.' - else: + elif os.path.exists(os.path.join(args.slices_dir, 'slices', 'metadata.json')): # metadata inside a 'slices' subfolder; treat provided path as model_dir slices_dir_effective = os.path.join(args.slices_dir, 'slices') model_dir = args.slices_dir + else: + print(f"{Fore.YELLOW}Warning: No slices metadata found at the provided path. Please slice the model first.{Style.RESET_ALL}") + logger.error("Run requires a valid slices directory with metadata.json") + return # Normalize derived paths slices_dir_effective = normalize_path(slices_dir_effective) model_dir = normalize_path(model_dir) # Get run metadata path if provided, otherwise None (Runner will auto-generate) - run_metadata_path = args.run_metadata_path if hasattr(args, 'run_metadata_path') and args.run_metadata_path else None + run_metadata_path = getattr(args, 'run_metadata_path', None) if run_metadata_path: run_metadata_path = normalize_path(run_metadata_path) @@ -123,7 +119,7 @@ def run_inference(args): # Use the Runner class for inference logger.info("Using Runner class for model inference") logger.info(f"Model path: {model_dir}, Slices path: {slices_dir_effective}") - + start_time = time.time() runner = Runner( model_path=model_dir, @@ -132,7 +128,7 @@ def run_inference(args): ) result = runner.run(args.input_file) elapsed_time = time.time() - start_time - + print(f"{Fore.GREEN}✓ Inference completed in {elapsed_time:.2f} seconds!{Style.RESET_ALL}") logger.info(f"Inference completed in {elapsed_time:.2f} seconds") diff --git a/dsperse/src/runner.py b/dsperse/src/runner.py index 6b75ede..c68e557 100644 --- a/dsperse/src/runner.py +++ b/dsperse/src/runner.py @@ -30,9 +30,9 @@ def __init__(self, model_path, slices_path=None, metadata_path=None, run_metadat if self.run_metadata_path is None or not self.run_metadata_path.exists(): logger.info("run metadata not found. Generating...") print(f"Generating run metadata at {self.run_metadata_path}") - runner_metadata = RunnerAnalyzer(self.model_path) + runner_metadata = RunnerAnalyzer(self.model_path, self.slices_path) self.run_metadata_path = runner_metadata.generate_metadata(save_path=self.run_metadata_path) - + with open(self.run_metadata_path, 'r') as f: self.metadata = json.load(f) @@ -61,7 +61,7 @@ def run(self, input_json_path) -> dict: seg_run_dir = run_dir / current_slice_id seg_run_dir.mkdir(parents=True, exist_ok=True) - + # Write input for this segment input_file = seg_run_dir / "input.json" output_file = seg_run_dir / "output.json" @@ -99,21 +99,21 @@ def run(self, input_json_path) -> dict: # filter tensor and make tensor next input.json file current_tensor = self._filter_tensor(current_slice_metadata, tensor) current_slice_id = slice_node.get("next") - + # Final processing probabilities = F.softmax(current_tensor, dim=1) prediction = torch.argmax(probabilities, dim=1).item() - + results = { "prediction": prediction, "probabilities": probabilities.tolist(), "tensor_shape": list(current_tensor.shape), "slice_results": slice_results } - + # Save inference output self._save_inference_output(results, run_dir / "run_result.json" ) - + return results @staticmethod @@ -168,17 +168,17 @@ def _run_ezkl_segment(self, slice_info: dict, input_tensor_path, output_witness_ exec_info['error'] = output_tensor if isinstance(output_tensor, str) else "Unknown EZKL error" return success, output_tensor, exec_info - + def _save_inference_output(self, results, output_path): """Save inference_output.json with execution details.""" model_path = self.metadata.get("model_path", "unknown") slice_results = results.get("slice_results", {}) - + # Count execution methods - ezkl_complete = sum(1 for r in slice_results.values() + ezkl_complete = sum(1 for r in slice_results.values() if r.get("method") == "ezkl_gen_witness") total_slices = len(slice_results) - + # Build execution results execution_results = [] for slice_id, exec_info in slice_results.items(): @@ -194,18 +194,18 @@ def _save_inference_output(self, results, output_path): # Propagate error message if present (e.g., EZKL failure reason before fallback) if "error" in exec_info and exec_info["error"]: witness_execution["error"] = exec_info["error"] - + # Create result_entry with segment_id and witness_execution result_entry = { "segment_id": slice_id, "witness_execution": witness_execution } - + execution_results.append(result_entry) - + # Calculate security percentage security_percent = (ezkl_complete / total_slices * 100) if total_slices > 0 else 0 - + # Build output structure inference_output = { "model_path": model_path, @@ -221,7 +221,7 @@ def _save_inference_output(self, results, output_path): "note": "Full ONNX vs verified chain comparison would require separate pure ONNX run" } } - + # Save to file with open(output_path, 'w') as f: json.dump(inference_output, f, indent=2) diff --git a/dsperse/src/utils/slicer_utils/onnx_slicer.py b/dsperse/src/utils/slicer_utils/onnx_slicer.py index a3e09e5..a455b32 100644 --- a/dsperse/src/utils/slicer_utils/onnx_slicer.py +++ b/dsperse/src/utils/slicer_utils/onnx_slicer.py @@ -1,13 +1,15 @@ +import logging import os.path +from typing import List + import onnx from onnx import shape_inference -import logging -from dsperse.src.analyzers.onnx_analyzer import OnnxAnalyzer -from typing import List, Dict -from dsperse.src.utils.utils import Utils from onnx.utils import extract_model from onnxruntime.tools import symbolic_shape_infer +from dsperse.src.analyzers.onnx_analyzer import OnnxAnalyzer +from dsperse.src.utils.utils import Utils + # Configure logger logger = logging.getLogger(__name__) diff --git a/pyproject.toml b/pyproject.toml index 7afdaed..8b65932 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "dsperse" version = "1.0.1" description = "Distributed zkML Toolkit" readme = "README.md" -requires-python = ">=3.8" +requires-python = ">=3.10" license = { file = "LICENSE" } authors = [{ name = "Inference Labs", email = "info@inferencelabs.com" }] dependencies = [