RK3588 NPU example, load prompt from text file

rkllm-update.cpp

// Copyright (c) 2024 by Rockchip Electronics Co., Ltd. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Modified by Pelochus
// also modified by av1d

/*
    this loads the system prompt from a file named prompt.txt
    which should be in the same folder as your model. if it
    doesn't find one, it loads a default prompt which is hard-coded.
    it will let you know which prompt it loaded.

    example contents of prompt.txt:

<|im_start|>system You are a helpful assistant. <|im_end|> <|im_start|>user
<|im_end|><|im_start|>assistant

    the text you input at the prompt will be placed between those two lines.

    run as you would normally, or add 'debug' argument:
    syntax: rkllm model.rkllm debug
    'debug' is optional and prints the prompt which is in use
*/

#include <string.h>
#include <unistd.h>
#include <string>
#include "rkllm.h"
#include <fstream>
#include <sstream>
#include <iostream>
#include <csignal>
#include <vector>

using namespace std;

std::string PROMPT_TEXT_PREFIX;
std::string PROMPT_TEXT_POSTFIX;

/* 
    load the prompt template from prompt.txt.
    if it doesn't exist we set a default template.
*/
void loadPromptText() {
    std::ifstream file("prompt.txt");
    std::string line1, line2;
    
    if (std::getline(file, line1) && std::getline(file, line2)) {
        PROMPT_TEXT_PREFIX = line1;
        PROMPT_TEXT_POSTFIX = line2;
        printf("loaded template from prompt.txt\n");
    } else {
        printf("prompt.txt not found, loading defaults.\n");
        // Handle the case when the file doesn't contain at least two lines
        PROMPT_TEXT_PREFIX = "<|im_start|>system You are a helpful assistant. <|im_end|> <|im_start|>user";
        PROMPT_TEXT_POSTFIX = "<|im_end|><|im_start|>assistant";
    }
}

LLMHandle llmHandle = nullptr;

void exit_handler(int signal)
{
    if (llmHandle != nullptr)
    {
        std::cout << "Caught exit signal. Exiting..." << std::endl;
        LLMHandle _tmp = llmHandle;
        llmHandle = nullptr;
        rkllm_destroy(_tmp);
        exit(signal);
    }
}

void callback(const char *text, void *userdata, LLMCallState state)
{
    if (state == LLM_RUN_FINISH)
    {
        printf("\n");
    }
    else if (state == LLM_RUN_ERROR)
    {
        printf("\\LLM run error\n");
    }
    else
    {
        printf("%s", text);
    }
}

int main(int argc, char **argv)
{
    /*
    if (argc != 2)
    {
        printf("Usage: %s [rkllm_model_path]\n", argv[0]);
        return -1;
    }
    */

    if (argc < 2 || argc > 3)
    {
        printf("Usage: %s [rkllm_model_path] [debug]\n", argv[0]);
        return -1;
    }
    
    const char *rkllm_model_path = argv[1];
    int debug = 0; // Default: Debug mode off

    if (argc == 3 && strcmp(argv[2], "debug") == 0)
    {
        debug = 1; // Enable debug mode
        printf("Debug mode is enabled\n");
    }

    loadPromptText();
    
    signal(SIGINT, exit_handler);
    std::string rkllm_model(argv[1]);
    printf("rkllm init start\n");

    RKLLMParam param = rkllm_createDefaultParam();
    param.modelPath = rkllm_model.c_str();
    param.target_platform = "rk3588";
    param.num_npu_core = 2;
    param.top_k = 1;
    param.max_new_tokens = 256;
    param.max_context_len = 512;
    
    int initResult = rkllm_init(&llmHandle, param, callback);
    if (initResult != 0) {
        printf("RKLLM init failed!\n");
        return -1;
    }
    printf("RKLLM init success!\n");
    
    std::vector<std::string> pre_input = {
        "Welcome to ezrkllm! This is an adaptation of Rockchip's ",
        "rknn-llm repo (see github.com/airockchip/rknn-llm) ",
        "for running LLMs on its SoCs' NPUs.\n",
        "You are currently running the runtime for ",
        param.target_platform,
        "\nTo exit the model, enter either exit or quit\n",
        "\nMore information here: https://github.com/Pelochus/ezrknpu",
        "\nDetailed information for devs here: https://github.com/Pelochus/ezrknn-llm"
    };
    
    std::cout << "\n*************************** Pelochus' ezrkllm runtime *************************\n" << std::endl;
    
    for (const auto& line : pre_input)
    {
        std::cout << line;
    }
    
    std::cout << "\n*******************************************************************************\n" << std::endl;

    while (true)
    {
        std::string input_str;
        printf("\n");
        printf("You: ");
        if (!std::getline(std::cin, input_str)) {
            std::cout << "Error reading input. Exiting..." << std::endl;
            break;
        }
        
        if (input_str == "exit" || input_str == "quit")
        {
            std::cout << "Quitting program..." << std::endl;
            break;
        }
        
        std::string text = PROMPT_TEXT_PREFIX + input_str + " " + PROMPT_TEXT_POSTFIX;
        /* print concatenated prompt for debugging ... */
        if (debug == 1)
        {
            printf("[debug]: %s\n", text.c_str());
        }
        printf("LLM: ");
        int runResult = rkllm_run(llmHandle, text.c_str(), nullptr);
        if (runResult != 0) {
            printf("Error running LLM. Exiting...\n");
            break;
        }
    }

    rkllm_destroy(llmHandle);

    return 0;
}