#!/bin/bash
set -xe

set_proxy() {
    export http_proxy=http://gzbh-aip-paddlecloud140.gzbh:8128
    export https_proxy=http://gzbh-aip-paddlecloud140.gzbh:8128
}
unset_proxy() {
    unset http_proxy https_proxy
}

source /home/pt201/bin/activate
pip config set global.disable-pip-version-check true

#set_proxy
# Install xtrt, xtrt-llm, xpytorch
if [ -n "$(python3 -m pip list| grep tvm)" ]; then
    python3 -m pip uninstall -y tvm
    rm -rf /home/pt201/lib/python3.8/site-packages/tvm*
fi
python3 -m pip install tvm*.whl

if [ -n "$(python3 -m pip list| grep xtrt-llm)" ]; then
    python3 -m pip uninstall -y xtrt-llm
    rm -rf /home/pt201/lib/python3.8/site-packages/xtrt*
fi
python3 -m pip install xtrt_llm*.whl

if [ -f xpytorch*.run ]; then
    chmod +x xpytorch*.run
    bash xpytorch*.run
    if [ ! -f /home/pt201/lib/python3.8/site-packages/torch_xmlir/libtcmalloc_minimal.so ]; then
        ln -s /home/workspace/dep_libs/tcmalloc/output/lib/libtcmalloc_minimal.so /home/pt201/lib/python3.8/site-packages/torch_xmlir/libtcmalloc_minimal.so.4
    fi
fi
pip install hyperparameter
pip install numpy==1.23.5

# Install other dependency
pip install transformers_stream_generator
pip install nltk rouge_score
pip install tiktoken
#unset_proxy

# Install cuda
if [ ! -d /usr/local/cuda-11.7/ ]; then
    mkdir -p downloads
    pushd downloads
    wget https://bd.bcebos.com/klx-pytorch-work-bd/docker_conda_dep/torch202_cuda117/cuda_11.7.1_515.65.01_linux.run
    bash cuda_11.7.1_515.65.01_linux.run --silent --toolkit
    popd
fi

pip install https://su.bcebos.com/v1/klx-sdk-release-public/xtrt_llm/depends/thop/thop-0.0.0-cp38-cp38-linux_x86_64.whl