git clone https://github.com/seonglae/treenformer && cd treenformer
pip3 install torch==2.1.2 --index-url https://download.pytorch.org/whl/cu121
pip install .
# Optional for flash attention or undefined symbol error
pip uninstall transformer-engine
pip install flash-attn>=2.5 --no-build-isolation
# Generate test text
python test.py treensformer