Skip to content

generate_model_file

Combined Documentation

get_imports_from_modules

Functionality

This function analyzes a given PyTorch model by iterating through all its modules, extracting the module class and package names, and generating necessary import statements. It filters out modules from __main__ and builtins and returns a sorted list of import statements.

Parameters

  • sequential_model: A PyTorch model (instance of nn.Module) whose modules will be inspected to derive import statements.

Usage

  • Purpose: Automatically determine and generate import statements required to reconstruct a model, useful for tasks such as model export or serving with Triton Inference Server.

Example

import torch
from torch import nn
from embedding_studio.inference_management.triton.utils.generate_model_file import get_imports_from_modules

# Assume you have a PyTorch model instance
model = nn.Sequential(
    nn.Linear(10, 5),
    nn.ReLU(),
    nn.Linear(5, 2)
)

imports = get_imports_from_modules(model)
for imp in imports:
    print(imp)

generate_model_py

Functionality

Generates a standalone Python script that can reconstruct a given PyTorch model for use with the Triton Inference Server. It writes the necessary imports and model structure code.

Parameters

  • model: The PyTorch model to generate code for.
  • filename: Path where the generated script will be saved.
  • embedding_studio_path: (Optional) Base path for package imports.

Usage

  • Purpose: Recreate the provided PyTorch model in a new Python script file suitable for Triton deployment.

Example

import torch
from torch import nn

model = nn.Sequential(
    nn.Linear(10, 5),
    nn.ReLU(),
    nn.Linear(5, 1)
)

generate_model_py(model, "generated_model.py")

add_modules

Functionality

This function recursively builds code to reconstruct a PyTorch model layer tree with proper indentation. It generates code that adds each layer as a module to a parent module.

Parameters

  • layer: The current layer or module to process.
  • layer_name: Name assigned to the layer in the generated code.
  • indent_level: Indentation level (default is 2).

Usage

  • Purpose: Generate code to reconstruct model layers in a parent module.

Example

For a container module with children, calling add_modules(container, 'layer', 2) generates code that recursively adds the container and its children.


generate_sequential_model_py

Functionality

Generates a Python script that can rebuild a PyTorch Sequential model for use with Triton Inference Server. It analyzes a given nn.Sequential instance to extract necessary imports and layers, then writes a standalone script that recreates the model.

Parameters

  • sequential_model: A PyTorch Sequential model instance. It is analyzed to determine module imports and structure.
  • filename: Path where the generated Python script is saved.
  • embedding_studio_path: Optional string for the package path. Default is "/embedding_studio".

Usage

This function automatically generates a script to recreate a neural network model. It is specialized for Sequential models.

Example

Assuming you have a Sequential model named model, run:

generate_sequential_model_py(model, "model.py")
This creates a file "model.py" that can reconstruct the model when loaded by Triton.


add_modules (Additional Information)

Functionality

Recursively builds code to add and register PyTorch modules in the generated model file. This function handles both container modules and leaf modules by producing code for the reconstruction process.

Parameters

  • layer: The current PyTorch module or layer to process. Expected to be an instance of nn.Module.
  • layer_name: A string representing the unique name assigned to the module being added in generated code.
  • indent_level: An integer indicating the current indentation level for formatting the generated Python code.

Usage

  • Purpose: Traverse and generate code-line instructions for model assembly. Used internally to transform model layers into recreatable code snippets.

Example

For a simple model, call:

add_modules(model, "model", indent_level=2)
This call recursively generates module additions such as: self.add_module('model_layer', nn.Module()), for each sub-module.