#!/usr/bin/env bash

################################################################################################################# SOF
#####################################################################################################################

BANNER="\
##################################################################################!\
#                                                                                #!\
#                                                                                #!\
#                                                                                #!\
#                                               ##             .###              #!\
#                       ##                      ##             ####              #!\
#                       ##                      ##             #:##              #!\
#  :####    ##    ##  #######    .####.    :###.##   ##.####     ##     ##.###:  #!\
#  ######   ##    ##  #######   .######.  :#######   #######     ##     #######: #!\
#  #:  :##  ##    ##    ##      ###  ###  ###  ###   ###.        ##     ###  ### #!\
#   :#####  ##    ##    ##      ##.  .##  ##.  .##   ##          ##     ##.  .## #!\
# .#######  ##    ##    ##      ##    ##  ##    ##   ##          ##     ##    ## #!\
# ## .  ##  ##    ##    ##      ##.  .##  ##.  .##   ##          ##     ##.  .## #!\
# ##:  ###  ##:  ###    ##.     ###  ###  ###  ###   ##          ##     ###  ### #!\
# ########   #######    #####   .######.  :#######   ##       ########  #######: #!\
#   ###.##    ###.##    .####    .####.    :###.##   ##       ########  ##.###:  #!\
#                                                                       ##       #!\
#                                                                       ##       #!\
#                                                                       ##       #!\
#                                                                                #!\
##################################################################################"

#####################################################################################################################
#####################################################################################################################

PARAMS="$1"

#####################################################################################################################
#####################################################################################################################

BRK="##################################################################################"

#####################################################################################################################
#####################################################################################################################

CODE="import sys\n\
import torch\n\
import model_training.models.reward_model\n\
from huggingface_hub import hf_hub_download\n\
from transformers import AutoModel, AutoTokenizer, AutoModelForSequenceClassification\n\
from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXConfig, GPTNeoXModel, GPTNeoXPreTrainedModel\n\
def c(s):\n\
        COLORS={'r':31,'g':32,'y':33,'b':34,'m':35,'c':36,\n\
                'R':91,'G':92,'Y':93,'B':94,'M':95,'C':96}\n\
        return f'\033[{COLORS[s]}m'\n\
if __name__=='__main__':\n\
        print('{C1}...[ {C2}LOADING {C3}]... {C4}- {C5}<<< {C6}dr1p {C7}- {C8}autodr1p {C9}- {C10}https://huggingface.co/OpenAssistant/oasst-rm-2.1-pythia-1.4b-epoch-2.5 {C11}>>>\666'.format(\
                                                                                C1=c('r'),C2=c('C'),C3=c('r'),C4=c('c'),C5=c('r'),C6=c('c'),C7=c('r'),C8=c('c'),C9=c('r'),C10=c('c'),C11=c('r')))
        EXAMPLE_PROMPT='<|prompter|>Hi how are you?<|endoftext|><|assistant|>Hi, I am Open-Assistant a large open-source language model trained by LAION AI. How can I help you today?<|endoftext|>'\n\
        model_name='OpenAssistant/oasst-rm-2.1-pythia-1.4b-epoch-2.5'\n\
        files=['special_tokens_map.json','config.json','tokenizer.json','tokenizer_config.json','pytorch_model.bin']\n\
        for _files in files: hf_hub_download(repo_id=model_name, filename=_files)\n\
        device='cuda' if torch.cuda.is_available() else 'cpu'\n\
        tokenizer=AutoTokenizer.from_pretrained(model_name)\n\
        model=AutoModelForSequenceClassification.from_pretrained(model_name)\n\
        print('{C1}EXAMPLE PROMPT {C2}\666\666{C3}{EXAMPLE}\666'.format(C1=c('r'),C2=':'*14,C3=c('R'),EXAMPLE=EXAMPLE_PROMPT))\n\
        try:\n\
                if sys.argv[1].lower().strip()=='decoded':\n\
                        while True:\n\
                                input_text=input('{C1}ASK {C2}? {C3}< {C4}'.format(C1=c('c'),C2=c('r'),C3=c('c'),C4=c('r')))\n\
                                if input_text.lower()=='quit' or input_text.lower()=='exit':\n\
                                        sys.exit(0)\n\
                                else:\n\
                                        inputs=tokenizer(input_text, return_tensors='pt')\n\
                                        score=model(**inputs).logits[0].cpu().detach()\n\
                                        print('\666{C1}{C2}\666{C3}{C4}\666{C5}{C6}\666'.format(C1=c('r'),C2=':'*14,C3=c('c'),C4=score,C5=c('r'),C6=':'*14,))\n\
        except:\n\
                input_text=EXAMPLE_PROMPT\n\
                inputs=tokenizer(input_text, return_tensors='pt')\n\
                score=model(**inputs).logits[0].cpu().detach()\n\
                print('\666{C1}{C2}\666{C3}{C4}\666{C5}{C6}\666'.format(C1=c('r'),C2=':'*14,C3=c('c'),C4=score,C5=c('r'),C6=':'*14,))"

# RGB ###############################################################################################################
################################################################################################################# RGB

function rgb () {

        #############################################################################################################
        #############################################################################################################

        R=$((RANDOM%256))
        G=$((RANDOM%256))
        B=$((RANDOM%256))

        #############################################################################################################
        #############################################################################################################

        echo -e "\e[38;2;"$R";"$G";"$B"m"

        #############################################################################################################
        #############################################################################################################

} # RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB RGB #

# num ###############################################################################################################
################################################################################################################# num

function num () {

        #############################################################################################################
        #############################################################################################################

        l=$(echo $1|wc -c)
        n=0

        #############################################################################################################
        #############################################################################################################

        while [ $n -lt $((l+3)) ]; do
                echo -ne "#"
                n=$((n+1))
        done

        #############################################################################################################
        #############################################################################################################

} # NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM NUM #

# LOG ###############################################################################################################
################################################################################################################# LOG

function log () {

        #############################################################################################################
        #############################################################################################################

        [ $3 -eq 0 ] && msg "$(num "$1 - $2")![ $1 - $2 ]!$(num "$1 - $2")"
        [ $3 -eq 1 ] && msg "+ SUCCESS > $1 - $2"
        [ $3 -eq 2 ] && msg "- FAILURE > $1 - $2"
        [ $3 -eq 3 ] && msg "!$(num "$2")!# $2 #!$(num "$2")!"
        [ $3 -eq 4 ] && msg "> $1 - $2"

        #############################################################################################################
        #############################################################################################################

} # LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG LOG #

# MSG ###############################################################################################################
################################################################################################################# MSG

function msg () {

        #############################################################################################################
        #############################################################################################################

        MSG="$1"

        #############################################################################################################
        #############################################################################################################

        for ((i=0; i < ${#MSG}; i++)); do
                        MSGCHAR=${MSG:i:1}
                        if [ "$MSGCHAR" == "!" ]; then
                                echo -e ""
                        else
                                echo -ne "$(rgb)$MSGCHAR"
                        fi
        done; echo ""

        #############################################################################################################
        #############################################################################################################

} # MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG MSG #

# PREREQ_PROGS ######################################################################################################
######################################################################################################## PREREQ_PROGS

function prereq_progs () {

        #############################################################################################################
        #############################################################################################################

        PREREQUISITES_PROGS=(sed git python3.9 docker)
        MISSING_PROGS=()

        #############################################################################################################
        #############################################################################################################

        for p in ${PREREQUISITES_PROGS[@]}; do

                #####################################################################################################
                #####################################################################################################

                which $p >> /dev/null 2>&1

                #####################################################################################################
                #####################################################################################################

                        if [ $? -eq 0 ]; then
                                log "PROGRAM" "LOCATED: $(echo $p|tr a-z A-Z)" 1
                        else
                                log "PROGRAM" "MISSING: $(echo $p|tr a-z A-Z)" 2
                                MISSING_PROGS+=("$p")
                        fi

                        #############################################################################################
                        #############################################################################################

                        if [ "docker" == $p ]; then
                                docker compose > /dev/null 2>&1
                                if [ $? -eq 0 ]; then
                                        log "PROGRAM" "LOCATED: $(echo 'DOCKER COMPOSE'|tr a-z A-Z)" 1
                                else
                                        log "PROGRAM" "MISSING: $(echo 'DOCKER COMPOSE'|tr a-z A-Z)" 2
                                        MISSING_PROGS+=("docker compose")
                                fi
                        fi

                #####################################################################################################
                #####################################################################################################

        done

        #############################################################################################################
        #############################################################################################################

        for p in ${MISSING_PROGS[@]}; do

                #####################################################################################################
                #####################################################################################################

                MESSAGE="INSTALLATION REQUIRED"

                #####################################################################################################
                #####################################################################################################

                        [ "sed" == $p ] && (log "$(echo $p|tr a-z A-Z)" "$MESSAGE" 2)
                        [ "git" == $p ] && (log "$(echo $p|tr a-z A-Z)" "$MESSAGE" 2)
                        [ "python3.9" == $p ] && (log "$(echo $p|tr a-z A-Z)" "$MESSAGE" 2)
                        [ "docker" == $p ] && ( log "$(echo $p|tr a-z A-Z)" "$MESSAGE" 2 )

                #####################################################################################################
                #####################################################################################################

        done

        #############################################################################################################
        #############################################################################################################

        TXT="PRE-REQUISITES CHECK - INSTALLED PROGRAMS"

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        if [ 0 -ne ${#MISSING_PROGS[@]} ]; then
                log "$TXT" "FAILED" 2
                msg "!$BRK!"
                log "!FAILURE" "TROUBLE LOCATING PRE-REQUISITE PROGRAMS: ABORTING INSTALLATION" 2
                exit 1
        else
                log "$TXT" "PASSED" 1
        fi

        #############################################################################################################
        #############################################################################################################

} # PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ #

# PREREQ_PMODS ######################################################################################################
######################################################################################################## PREREQ_PMODS

function prereq_pmods () {

        #############################################################################################################
        #############################################################################################################

        PREREQUISITES_PMODS=("pip sqlmodel torch transformers huggingface_hub uvicorn")
        MISSING_PMODS=()

        #############################################################################################################
        #############################################################################################################

        for p in ${PREREQUISITES_PMODS[@]}; do

                #####################################################################################################
                #####################################################################################################

                log "SHELLED" "CHECKING PYTHON MODULE: $(echo $p|tr a-z A-Z)" 4
                msg "!$BRK!"
                python3 -c "import $p" >> /dev/null 2>&1

                #####################################################################################################
                #####################################################################################################

                if [ $? -eq 0 ]; then
                        log "SHELLED" "PRE-LOCATED PYTHON MODULE: $(echo $p|tr a-z A-Z)" 1
                        msg "!$BRK!"
                        if [ "pip" == $p ]; then
                                log "SHELLED" "CHECKING PYTHON PIP MODULE UPGRADABILITY" 4
                                msg "!$BRK!"
                                python3 -m pip install --upgrade pip|grep Successfully > /dev/null 2>&1
                                if [ $? -eq 0 ]; then
                                        log "SHELLED" "UPGRADED PYTHON MODULE PIP" 4
                                        msg "!$BRK!"
                                else
                                        log "SHELLED" "SKIPPING: UPGRADE PYTHON MODULE PIP ( ALREADY NEWEST )" 4
                                        msg "!$BRK!"
                                fi
                        fi
                else
                        log "PRE-MISSING PYTHON MODULE:" "$(echo $p|tr a-z A-Z)" 2
                        msg "!$BRK!"
                        MISSING_PMODS+=("$p")
                fi

                #####################################################################################################
                #####################################################################################################

        done

        #############################################################################################################
        #############################################################################################################

        COUNT=0

        #############################################################################################################
        #############################################################################################################

        for ap in ${MISSING_PMODS[@]}; do

                #####################################################################################################
                #####################################################################################################

                if [ $COUNT -gt 0 ]; then
                        msg "!$BRK!"
                else
                        COUNT=$((COUNT+1))
                fi
                log "SHELLED" "PIP INSTALLING PYTHON MODULE: $(echo $ap|tr a-z A-Z)" 4

                #####################################################################################################
                #####################################################################################################

                msg "!$BRK!"
                [ "pip" == $ap ] && ( wget https://bootstrap.pypa.io/get-pip.py; python3 get-pip.py )
                [ "sqlmodel" == $ap ] && ( python3 -m pip install $ap )
                [ "torch" == $ap ] && ( python3 -m pip install $ap )
                [ "transformers" == $ap ] && ( python3 -m pip install $ap )
                [ "huggingface_hub" == $ap ] && ( python3 -m pip install $ap )
                [ "uvicorn" == $ap ] && ( python3 -m pip install $ap; msg "!$BRK!" )

                #####################################################################################################
                #####################################################################################################

        done

        #############################################################################################################
        #############################################################################################################

        MISSING_PMODS=()

        #############################################################################################################
        #############################################################################################################

        if [ ${#MISSING_PMODS[@]} -gt 0 ]; then
                msg "!$BRK!"
        fi

        #############################################################################################################
        #############################################################################################################

        for p in ${PREREQUISITES_PMODS[@]}; do

                #####################################################################################################
                #####################################################################################################

                python3 -c "import $p" >> /dev/null 2>&1

                #####################################################################################################
                #####################################################################################################

                if [ $? -eq 0 ]; then
                        log "SHELLED" "POST-LOCATED PYTHON MODULE: $(echo $p|tr a-z A-Z)" 1
                else
                        log "SHELLED" "POST-MISSING PYTHON MODULE: $(echo $p|tr a-z A-Z)" 2
                        MISSING_PMODS+=("$p")
                fi

                #####################################################################################################
                #####################################################################################################

        done

        #############################################################################################################
        #############################################################################################################

        TXT="PRE-REQUISITES CHECK - INSTALLED PYTHON3 MODULES"

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        if [ 0 -ne ${#MISSING_PMODS[@]} ]; then
                log "$TXT" "FAILED" 2
                msg "!$BRK!"
                log "!FAILURE" "TROUBLE LOCATING OR INSTALLING PRE-REQUISITE PYTHON MODULES: ABORTING INSTALLATION" 2
                exit 1
        else
                log "$TXT" "PASSED" 1
        fi
        msg "!$BRK!"

        #############################################################################################################
        #############################################################################################################

} # PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ PREREQ #

# INSTALL ###########################################################################################################
############################################################################################################# INSTALL

function install () {

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "CHANGING PWD ~/AI" 4
        msg "!$BRK!"
        [ ! -e ~/ai ] && mkdir ~/ai
        cd ~/ai

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "GIT CLONE REPOSITORY" 4
        msg "!$BRK!"
        if [ ! -e "Open-Assistant" ]; then
                git clone --recursive https://github.com/LAION-AI/Open-Assistant.git > /dev/null 2>&1
                if [ $? -eq 0 ]; then
                        log "SHELLED" "CLONED GITHUB REPOSITORY SUCCESS" 1
                else
                        log "SHELLED" "CLONED GITHUB REPOSITORY FAILURE" 2
                        MISSING_PMODS+=("$p")
                fi
        else
                log "SHELLED" "CLONED GITHUB REPOSITORY ( ALREADY EXISTS )" 4
        fi
        msg "!$BRK!"
        if [ $? -eq 1 ]; then
                log "!FAILURE" "TROUBLE CLONING GITHUB REPOSITORY: ABORTING INSTALLATION" 2
                exit 1
        fi

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "CHANGING PWD ~/AI/OPEN-ASSISTANT" 4
        msg "!$BRK!"
        cd Open-Assistant

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "CREATE PYTHON VIRTUAL ENVIRONMENT" 4
        msg "!$BRK!"
        if [ ! -e venv ]; then
                log "SHELLED" "CREATING PYTHON VIRTUAL ENVIRONMENT" 4
                msg "!$BRK!"
                python3.11 -m venv python3.11 venv > /dev/null 2>&1
                if [ $? -eq 0 ]; then
                        log "SUCCESS" "CREATED PYTHON VIRTUAL ENVIRONMENT" 1
                        msg "!$BRK!"
                        echo 'alias python3="/home/$USER/ai/Open-Assistant/venv/bin/python3 -q"' >> venv/bin/activate
                else
                        log "FAILURE" "TROUBLE SETTING UP THE PYTHON VIRTUAL ENVIRONMENT: ABORTING INSTALLATION" 2
                        exit 1
                fi
        else
                log "SHELLED" "SKIPPING: CREATE PYTHON VIRTUAL ENVIRONMENT ( ALREADY EXISTS )" 4
                msg "!$BRK!"
        fi

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "ACTIVATING PYTHON VIRTUAL ENVIRONMENT" 4
        source venv/bin/activate

        #############################################################################################################
        #############################################################################################################

        log "MAIN" "PREREQUISITES CHECK - PYTHON MODULES" 3; prereq_pmods

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "INSTALL PYTHON MODULES FROM CLONED GITHUB REPOSITORY DIRECTORY ~/AI/OPEN-ASSISTANT/MODEL" 4
        msg "!$BRK!"
        python3 -c "import model_training" > /dev/null 2>&1
        if [ $? -eq 0 ]; then
                log "SHELLED" "SKIPPING: INSTALL PYTHON MODULES FROM CLONED GITHUB REPOSITORY DIRECTORY ~/AI/OPEN-ASSISTANT/MODEL" 4
        else
                log "SHELLED" "INSTALLING PYTHON MODULES FROM CLONED GITHUB REPOSITORY DIRECTORY ~/AI/OPEN-ASSISTANT/MODEL" 4
                msg "!$BRK!"
                python3 -m pip install -e model
        fi

        #############################################################################################################
        #############################################################################################################

} # INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL INSTALL #

# RUNTIME ###########################################################################################################
############################################################################################################# RUNTIME

function runtime () {

        #############################################################################################################
        #############################################################################################################

        FILE="/home/$USER/ai/Open-Assistant/autodr1p.py"

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "CREATING FILE AUTODR1P.PY" 4
        msg "!$BRK!"
        echo -e "${CODE}" > "${FILE}"
        sed -i s,"666","n",g "$FILE"

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "DOWNLOADED MODEL PATH: /HOME/$USER/.CACHE/HUGGINGFACE/HUB/MODELS--OPENASSISTANT--OASST-RM-2.1-PYTHIA-1.4B-EPOCH-2.5" 4
        msg "!$BRK!"

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "LOADING PYTHON FILE: AUTODR1P.PY ( WILL TAKE A MOMENT TO LOAD )" 4
        msg "!$BRK!"

        #############################################################################################################
        #############################################################################################################

        python3 "$FILE" "$PARAMS"

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "CHANGING PWD ~/AI/OPEN-ASSISTANT/INFERENCE" 4
        msg "!$BRK!"
        cd inference

        #############################################################################################################
        #############################################################################################################

        log "SHELLED" "PIP INSTALLING ~/AI/OPEN-ASSISTANT/INFERENCE/WORKER/REQUIREMENTS.TXT" 4
        msg "!$BRK!"
        python -m pip install -r worker/requirements.txt

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        log "SHELLED" "PIP INSTALLING ~/AI/OPEN-ASSISTANT/INFERENCE/SERVER/REQUIREMENTS.TXT" 4
        msg "!$BRK!"
        python -m pip install -r server/requirements.txt

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        log "SHELLED" "PIP INSTALLING ~/AI/OPEN-ASSISTANT/INFERENCE/TEXT-CLIENT/REQUIREMENTS.TXT" 4
        msg "!$BRK!"
        python -m pip install -r text-client/requirements.txt

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        log "SHELLED" "PIP INSTALLING ~/AI/OPEN-ASSISTANT/INFERENCE/OASST-SHARED/REQUIREMENTS.TXT" 4
        msg "!$BRK!"
        python -m pip install -e ../oasst-shared

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        log "SHELLED" "DOCKER COMPOSE BUILD: INFERENCE" 4
        msg "!$BRK!"
        DOCKER_BUILDKIT=1 docker compose --profile inference build

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        log "SHELLED" "DOCKER COMPOSE DAEMON UP: INFERENCE" 4
        msg "!$BRK!"
        docker compose --profile inference up -d

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        log "SHELLED" "DOCKER COMPOSE BUILD: CI" 4
        msg "!$BRK!"
        DOCKER_BUILDKIT=1 docker compose --profile ci up --build --attach-dependencies

        #############################################################################################################
        #############################################################################################################

        msg "!$BRK!"
        log "SHELLED" "FINISHED: EXIT" 4

        #############################################################################################################

} # RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME RUNTIME #

# MAIN ##############################################################################################################
################################################################################################################ MAIN

msg "!$BANNER"
log "MAIN" "PRE-REQUISITES CHECK - INSTALLED PROGRAMS" 3; prereq_progs
log "MAIN" "INSTALLATION" 3; install
log "MAIN" "RUNTIME" 3; runtime

# EOF ###############################################################################################################
################################################################################################################# EOF
