echodict/llama.cpp

1#!/usr/bin/env bash2 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 extglob14    set -- "${1##+([[:space:]])}"15    printf "%s" "${1%%+([[:space:]])}"16}17 18trim_trailing() {19    shopt -s extglob20    printf "%s" "${1%%+([[:space:]])}"21}22 23format_prompt() {24    if [[ "${#CHAT[@]}" -eq 0 ]]; then25        echo -n "[INST] <<SYS>>\n${INSTRUCTION}\n<</SYS>>"26    else27        LAST_INDEX=$(( ${#CHAT[@]} - 1 ))28        echo -n "${CHAT[$LAST_INDEX]}\n[INST] $1 [/INST]"29    fi30}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: true55    }')"56 57    # Create a temporary file to hold the Python output58    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 json70import sys71 72answer = ''73while True:74    line = sys.stdin.readline()75    if not line:76        break77    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 += content83 84answer = answer.rstrip('\n')85 86# Write the answer to the temporary file87with open('$TEMPFILE', 'w') as f:88    f.write(answer)89    " <&390 91    exec 3<&-92 93    # Read the answer from the temporary file94    ANSWER=$(cat $TEMPFILE)95 96    # Clean up the temporary file97    rm $TEMPFILE98 99    printf "\n"100 101    CHAT+=("$1" "$(trim "$ANSWER")")102}103 104while true; do105    echo -en "\033[0;32m"  # Green color106    read -r -e -p "> " QUESTION107    echo -en "\033[0m"  # Reset color108    chat_completion "${QUESTION}"109done110