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 ofnn.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")
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 ofnn.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)
self.add_module('model_layer', nn.Module())
, for each sub-module.