from pathlib import Path

import numpy as np

from preprocessor.feature_extractor import MotokoFeatureExtractor


def main() -> None:
    extractor = MotokoFeatureExtractor.from_config(
        Path("preprocessor/preprocessor_config.json")
    )

    sample = {
        "force": np.random.randn(320, 3).astype(np.float32),
        "torque": np.random.randn(320, 3).astype(np.float32),
        "pressure": np.random.randn(320, 16).astype(np.float32),
        "vibration": np.random.randn(320, 6).astype(np.float32),
    }

    features = extractor(sample)
    print("input_values:", features["input_values"].shape)
    print("attention_mask:", features["attention_mask"].shape)


if __name__ == "__main__":
    main()