Utilizing Custom ONNX Models Stored in Hugging Face within HSSM

This guide will walk you through the process of using custom ONNX models stored in Hugging Face within HSSM (Hierarchical State Space Model) framework.

Prerequisites

1. Python 3.8 or later.
2. HSSM library installed in your Python environment.
3. A pre-trained ONNX model stored on Hugging Face model hub.

Step-by-step guide

Step 1: Import necessary libraries

import pandas as pd
import hssm
import ssms.basic_simulators

pytensor.config.floatX = "float32"

Step 2: Define HSSM Configuration

You will have to define the configuration of your model. Make sure you are defining the log-likelihood kind as "approx_differentiable" and providing the Hugging Face model name in the loglik field.

my_hssm = hssm.HSSM(
        data=dataset_lan,
        loglik_kind = "approx_differentiable",
        loglik = "levy.onnx",
        model="custom",
        model_config= {
            "backend": "jax",
            "list_params": ["v", "a", "z", "alpha", "t"],
            "bounds": {
                "v": (-3.0, 3.0),
                "a": (0.3, 3.0),
                "z": (0.1, 0.9),
                "alpha": (1.0, 2.0),
                "t": (1e-3, 2.0),
            },
        }                  
    )

This creates an HSSM object my_hssm using the custom ONNX model levy.onnx from the Hugging Face repository.

my_hssm.sample(cores=2, draws=500, tune=500, mp_ctx="forkserver")

Uploading ONNX Files to a Hugging Face Repository

If your ONNX file is not currently housed in your Hugging Face repository, you can include it by adhering to the steps delineated below:

1. Import the HfApi module from huggingface_hub:

from huggingface_hub import HfApi

2. Upload the ONNX file using the upload_file method:

api = HfApi()
api.upload_file(
    path_or_fileobj="test.onnx",
    path_in_repo="test.onnx",
    repo_id="franklab/HSSM",
    repo_type="model",
    create_pr=True,
)

The execution of these steps will generate a Pull Request (PR) on Hugging Face, which will subsequently be evaluated by a member of our team.

Creating a Pull Request and a New ONNX Model

1. Creating a Pull Request on Hugging Face

   Navigate to the following link: Hugging Face PR

   By doing so, you will generate a Pull Request on Hugging Face, which will be reviewed by our team members.

2. Creating a Custom ONNX Model

   Establish a Network Config and State Dictionary Files in PyTorch

   To construct a custom model and save it as an ONNX file, you must create a network configuration file and a state dictionary file in PyTorch. Refer to the instructions outlined in the README of the LANFactory package.

   Convert Network Config and State Dictionary Files to ONNX

   Once you've generated the network configuration and state dictionary files, you will need to convert these files into an ONNX format.