KBaba7/llama.cpp

1#!/bin/bash2 3set -euo pipefail4 5cd "$(dirname "$0")/.." || exit6 7if [[ -z "${PROMPT_CACHE_FILE+x}" || -z "${CHAT_SAVE_DIR+x}" ]]; then8    echo >&2 "error: PROMPT_CACHE_FILE and CHAT_SAVE_DIR must be provided"9    exit 110fi11 12MODEL="${MODEL:-./models/llama-13b/ggml-model-q4_0.gguf}"13PROMPT_TEMPLATE="${PROMPT_TEMPLATE:-./prompts/chat.txt}"14USER_NAME="${USER_NAME:-User}"15AI_NAME="${AI_NAME:-ChatLLaMa}"16DATE_TIME="$(date +%H:%M)"17DATE_YEAR="$(date +%Y)"18 19LOG="${CHAT_SAVE_DIR}/main.log"20LOG_BG="${CHAT_SAVE_DIR}/main-bg.log"21CUR_PROMPT_FILE="${CHAT_SAVE_DIR}/current-prompt.txt"22CUR_PROMPT_CACHE="${CHAT_SAVE_DIR}/current-cache.bin"23NEXT_PROMPT_FILE="${CHAT_SAVE_DIR}/next-prompt.txt"24NEXT_PROMPT_CACHE="${CHAT_SAVE_DIR}/next-cache.bin"25 26SESSION_AND_SAMPLE_PATTERN='main: session file matches [[:digit:]]+ / [[:digit:]]+'\27'|'\28'sampling time =[[:space:]]+[[:digit:]]+.[[:digit:]]+ ms /[[:space:]]+[[:digit:]]+'29SED_DELETE_MESSAGES="/^(${USER_NAME}:|${AI_NAME}:|\\.\\.\\.)/,\$d"30 31CTX_SIZE=204832CTX_ROTATE_POINT=$((CTX_SIZE * 3 / 5)) # REVIEW33OPTS=(--model "$MODEL" --ctx_size "$CTX_SIZE" --repeat_last_n 256 "$@")34 35# An unbuffered `tail -c+N`36skip_bytes() {37    LANG=C IFS= read -r -n "$1" -d '' c38    while LANG=C IFS= read -r -n 1 -d '' c; do39        printf '%s' "$c"40    done41}42 43mkdir -p "$CHAT_SAVE_DIR"44echo >"$LOG"45trap "tail -n100 ${LOG}" EXIT46 47if [[ ! -e "$CUR_PROMPT_FILE" ]]; then48    sed -e "s/\[\[USER_NAME\]\]/${USER_NAME}/g" \49        -e "s/\[\[AI_NAME\]\]/${AI_NAME}/g" \50        -e "s/\[\[DATE_TIME\]\]/${DATE_TIME}/g" \51        -e "s/\[\[DATE_YEAR\]\]/${DATE_YEAR}/g" \52        "$PROMPT_TEMPLATE" >"$CUR_PROMPT_FILE"53fi54 55if [[ ! -e "$NEXT_PROMPT_FILE" ]]; then56    sed -r "$SED_DELETE_MESSAGES" "$CUR_PROMPT_FILE" >"$NEXT_PROMPT_FILE"57fi58 59if [[ "$(tail -c4 "$NEXT_PROMPT_FILE")" != "..." ]]; then60    echo '...' >>"$NEXT_PROMPT_FILE"61fi62 63if [[ ! -e "$PROMPT_CACHE_FILE" ]]; then64    echo 'Prompt cache does not exist, building...'65    # Default batch_size to 64 here for better user feedback during initial prompt processing66    ./llama-cli 2>>"$LOG" \67        --batch_size 64 \68        "${OPTS[@]}" \69        --prompt-cache "$PROMPT_CACHE_FILE" \70        --file "$CUR_PROMPT_FILE" \71        --n_predict 172    echo73    echo 'Done!'74fi75 76if [[ ! -e "$CUR_PROMPT_CACHE" ]]; then77    cp "$PROMPT_CACHE_FILE" "$CUR_PROMPT_CACHE"78fi79if [[ ! -e "$NEXT_PROMPT_CACHE" ]]; then80    cp "$PROMPT_CACHE_FILE" "$NEXT_PROMPT_CACHE"81fi82 83printf '%s ' "$(< "$CUR_PROMPT_FILE")"84n_tokens=085 86while read -e line; do87    # Limit generation to remaining context, with a buffer and estimating 2 chars/token for input88    n_predict=$((CTX_SIZE - n_tokens - ${#line} / 2 - 32))89 90    # Swap prompts when we're about to run out of context91    if ((n_predict <= 0)); then92        wait # for background main (below) to finish with next prompt93        mv "$NEXT_PROMPT_FILE"  "$CUR_PROMPT_FILE"94        mv "$NEXT_PROMPT_CACHE" "$CUR_PROMPT_CACHE"95 96        sed -r "$SED_DELETE_MESSAGES" "$CUR_PROMPT_FILE" >"$NEXT_PROMPT_FILE"97        echo '...' >>"$NEXT_PROMPT_FILE"98        cp "$PROMPT_CACHE_FILE" "$NEXT_PROMPT_CACHE"99 100        n_tokens=0101        n_predict=$((CTX_SIZE / 2))102    fi103 104    echo " ${line}" >>"$CUR_PROMPT_FILE"105    if ((n_tokens > CTX_ROTATE_POINT)); then106        echo " ${line}" >>"$NEXT_PROMPT_FILE"107    fi108 109    n_prompt_len_pre=$(($(wc -c <"$CUR_PROMPT_FILE")))110 111    printf '%s: ' "$AI_NAME" >>"$CUR_PROMPT_FILE"112 113    ./llama-cli 2>>"$LOG" "${OPTS[@]}" \114            --prompt-cache "$CUR_PROMPT_CACHE" \115            --prompt-cache-all \116            --file "$CUR_PROMPT_FILE" \117            --reverse-prompt "${USER_NAME}:" \118            --n_predict "$n_predict" |119        skip_bytes 1 |                  # skip BOS token added by ./llama-cli120        tee "$CUR_PROMPT_FILE.tmp" |    # save prompt + generation to tmp file121        skip_bytes "$n_prompt_len_pre"  # print generation122 123    mv "$CUR_PROMPT_FILE.tmp" "$CUR_PROMPT_FILE"124 125    # if we hit n_predict instead of reverse-prompt, we need to add the prompt126    if [[ "$(tail -n1 "$CUR_PROMPT_FILE")" != "${USER_NAME}:" ]]; then127        printf '\n%s:' "$USER_NAME"128        printf '\n%s:' "$USER_NAME" >> "$CUR_PROMPT_FILE"129    fi130 131    printf ' '132 133    if ! session_and_sample_msg=$(tail -n30 "$LOG" | grep -oE "$SESSION_AND_SAMPLE_PATTERN"); then134        echo >&2 "Couldn't get number of tokens from ./llama-cli output!"135        exit 1136    fi137 138    n_tokens=$(awk '{sum+=$1} END {print sum}' <<< "$(cut -d/ -f2 <<< "$session_and_sample_msg")")139 140    if ((n_tokens > CTX_ROTATE_POINT)); then141        tail -c+$((n_prompt_len_pre + 1)) "$CUR_PROMPT_FILE" >>"$NEXT_PROMPT_FILE"142    fi143 144    # Update cache for next prompt in background, ideally during user input145    ./llama-cli >>"$LOG_BG" 2>&1 "${OPTS[@]}" \146          --prompt-cache "$NEXT_PROMPT_CACHE" \147          --file "$NEXT_PROMPT_FILE" \148          --n_predict 1 &149done150