1#!/usr/bin/env bash
  2
  3API_URL="${API_URL:-http://127.0.0.1:8080}"
  4
  5CHAT=(
  6    "Hello, Assistant."
  7    "Hello. How may I help you today?"
  8)
  9
 10INSTRUCTION="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions."
 11
 12trim() {
 13    shopt -s extglob
 14    set -- "${1##+([[:space:]])}"
 15    printf "%s" "${1%%+([[:space:]])}"
 16}
 17
 18trim_trailing() {
 19    shopt -s extglob
 20    printf "%s" "${1%%+([[:space:]])}"
 21}
 22
 23format_prompt() {
 24    if [[ "${#CHAT[@]}" -eq 0 ]]; then
 25        echo -n "[INST] <<SYS>>\n${INSTRUCTION}\n<</SYS>>"
 26    else
 27        LAST_INDEX=$(( ${#CHAT[@]} - 1 ))
 28        echo -n "${CHAT[$LAST_INDEX]}\n[INST] $1 [/INST]"
 29    fi
 30}
 31
 32tokenize() {
 33    curl \
 34        --silent \
 35        --request POST \
 36        --url "${API_URL}/tokenize" \
 37        --header "Content-Type: application/json" \
 38        --data-raw "$(jq -ns --arg content "$1" '{content:$content}')" \
 39    | jq '.tokens[]'
 40}
 41
 42N_KEEP=$(tokenize "[INST] <<SYS>>\n${INSTRUCTION}\n<</SYS>>" | wc -l)
 43
 44chat_completion() {
 45    PROMPT="$(trim_trailing "$(format_prompt "$1")")"
 46    DATA="$(echo -n "$PROMPT" | jq -Rs --argjson n_keep $N_KEEP '{
 47        prompt: .,
 48        temperature: 0.2,
 49        top_k: 40,
 50        top_p: 0.9,
 51        n_keep: $n_keep,
 52        n_predict: 1024,
 53        stop: ["[INST]"],
 54        stream: true
 55    }')"
 56
 57    # Create a temporary file to hold the Python output
 58    TEMPFILE=$(mktemp)
 59
 60    exec 3< <(curl \
 61        --silent \
 62        --no-buffer \
 63        --request POST \
 64        --url "${API_URL}/completion" \
 65        --header "Content-Type: application/json" \
 66        --data-raw "${DATA}")
 67
 68    python -c "
 69import json
 70import sys
 71
 72answer = ''
 73while True:
 74    line = sys.stdin.readline()
 75    if not line:
 76        break
 77    if line.startswith('data: '):
 78        json_content = line[6:].strip()
 79        content = json.loads(json_content)['content']
 80        sys.stdout.write(content)
 81        sys.stdout.flush()
 82        answer += content
 83
 84answer = answer.rstrip('\n')
 85
 86# Write the answer to the temporary file
 87with open('$TEMPFILE', 'w') as f:
 88    f.write(answer)
 89    " <&3
 90
 91    exec 3<&-
 92
 93    # Read the answer from the temporary file
 94    ANSWER=$(cat $TEMPFILE)
 95
 96    # Clean up the temporary file
 97    rm $TEMPFILE
 98
 99    printf "\n"
100
101    CHAT+=("$1" "$(trim "$ANSWER")")
102}
103
104while true; do
105    echo -en "\033[0;32m"  # Green color
106    read -r -e -p "> " QUESTION
107    echo -en "\033[0m"  # Reset color
108    chat_completion "${QUESTION}"
109done