ck_cli_v5.c File Reference

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <time.h>
#include <dlfcn.h>
#include <pthread.h>
#include <unistd.h>
#include "tokenizer/tokenizer.h"

Go to the source code of this file.

Typedefs
typedef int(*	decode_t) (int32_t token, float *logits_out)
typedef int(*	embed_t) (const int32_t *tokens, int num_tokens)
typedef int(*	forward_t) (float *logits_out)
typedef int(*	get_int_t) (void)
typedef void *(*	get_ptr_t) (void)
typedef int(*	init_t) (const char *weights_path)
typedef int(*	kv_enable_t) (int capacity)
typedef int(*	sample_argmax_t) (void)
typedef void(*	void_func_t) (void)

Functions
void *	engine_thread_func (void *arg)
int	main (int argc, char **argv)

Typedef Documentation

decode_t

typedef int(* decode_t) (int32_t token, float *logits_out)

Definition at line 17 of file ck_cli_v5.c.

embed_t

typedef int(* embed_t) (const int32_t *tokens, int num_tokens)

Definition at line 14 of file ck_cli_v5.c.

forward_t

typedef int(* forward_t) (float *logits_out)

Definition at line 15 of file ck_cli_v5.c.

get_int_t

typedef int(* get_int_t) (void)

Definition at line 19 of file ck_cli_v5.c.

get_ptr_t

typedef void*(* get_ptr_t) (void)

Definition at line 20 of file ck_cli_v5.c.

init_t

typedef int(* init_t) (const char *weights_path)

Definition at line 13 of file ck_cli_v5.c.

kv_enable_t

typedef int(* kv_enable_t) (int capacity)

Definition at line 16 of file ck_cli_v5.c.

sample_argmax_t

typedef int(* sample_argmax_t) (void)

Definition at line 18 of file ck_cli_v5.c.

void_func_t

typedef void(* void_func_t) (void)

Definition at line 21 of file ck_cli_v5.c.

Function Documentation

engine_thread_func()

void* engine_thread_func

(

void *

arg

)

Definition at line 53 of file ck_cli_v5.c.

                                    {
    SharedState *s = (SharedState *)arg;
    
    printf("[Engine] Thread started. Initializing model...\n");
    if (s->init(s->weights_path) != 0) {
        fprintf(stderr, "[Engine] Failed to init model\n");
        return NULL;
    }
    s->kv_enable(s->get_context());
 
    while (1) {
        pthread_mutex_lock(&s->mutex);
        while (!s->task_ready && !s->quit) {
            pthread_cond_wait(&s->cond_task, &s->mutex);
        }
        if (s->quit) {
            pthread_mutex_unlock(&s->mutex);
            break;
        }
        
        // Start Task
        int n_prompt = s->n_prompt;
        int32_t prompt[1024];
        memcpy(prompt, s->prompt_tokens, n_prompt * sizeof(int32_t));
        int max_gen = s->max_gen;
        s->task_ready = false;
        pthread_mutex_unlock(&s->mutex);
 
        // 1. Prefill
        s->embed(prompt, n_prompt);
        s->forward(NULL);
        int32_t next_token = s->sample();
 
        // 2. Feedback first token
        pthread_mutex_lock(&s->mutex);
        s->last_token = next_token;
        s->token_ready = true;
        pthread_cond_signal(&s->cond_done);
        pthread_mutex_unlock(&s->mutex);
 
        // 3. Decode Loop
        for (int i = 0; i < max_gen; i++) {
            if (s->decode(next_token, NULL) != 0) break;
            next_token = s->sample();
 
            pthread_mutex_lock(&s->mutex);
            s->last_token = next_token;
            s->token_ready = true;
            pthread_cond_signal(&s->cond_done);
            pthread_mutex_unlock(&s->mutex);
 
            if (next_token == 151643 || next_token == 151645) break; 
        }
    }
    return NULL;
}

Referenced by main().

main()

int main	(	int	argc,
		char **	argv
	)

Definition at line 110 of file ck_cli_v5.c.

                                {
    if (argc < 3) {
        printf("Usage: %s <libmodel.so> <weights.bump>\n", argv[0]);
        return 1;
    }
 
    SharedState state = {0};
    pthread_mutex_init(&state.mutex, NULL);
    pthread_cond_init(&state.cond_task, NULL);
    pthread_cond_init(&state.cond_done, NULL);
    state.weights_path = argv[2];
 
    void *handle = dlopen(argv[1], RTLD_NOW);
    if (!handle) { fprintf(stderr, "%s\n", dlerror()); return 1; }
 
    state.init = dlsym(handle, "ck_model_init");
    state.embed = dlsym(handle, "ck_model_embed_tokens");
    state.forward = dlsym(handle, "ck_model_forward");
    state.kv_enable = dlsym(handle, "ck_model_kv_cache_enable");
    state.decode = dlsym(handle, "ck_model_decode");
    state.sample = dlsym(handle, "ck_model_sample_argmax");
    state.get_context = dlsym(handle, "ck_model_get_context_window");
    get_ptr_t get_offsets = dlsym(handle, "ck_model_get_vocab_offsets");
    get_ptr_t get_strings = dlsym(handle, "ck_model_get_vocab_strings");
    get_int_t get_vocab_size = dlsym(handle, "ck_model_get_vocab_size");
    get_int_t get_num_merges = dlsym(handle, "ck_model_get_num_merges");
 
    // Start Engine
    pthread_t engine_thread;
    pthread_create(&engine_thread, NULL, engine_thread_func, &state);
 
    // Wait for engine to init (simple wait for this demo)
    sleep(1);
 
    // Tokenizer setup
    CKTokenizer *tokenizer = ck_tokenizer_create(CK_TOKENIZER_BPE);
    ck_tokenizer_load_binary(tokenizer, get_vocab_size(), get_offsets(), get_strings(), get_num_merges(), NULL);
 
    char input[1024];
    while (1) {
        printf("\nYou: ");
        if (!fgets(input, sizeof(input), stdin)) break;
        if (strncmp(input, "/exit", 5) == 0) break;
 
        // Tokenize
        int32_t ids[1024];
        int n = ck_tokenizer_encode(tokenizer, input, -1, ids, 1024);
 
        // Submit Task
        pthread_mutex_lock(&state.mutex);
        memcpy(state.prompt_tokens, ids, n * sizeof(int32_t));
        state.n_prompt = n;
        state.max_gen = 100;
        state.task_ready = true;
        state.token_ready = false;
        pthread_cond_signal(&state.cond_task);
        pthread_mutex_unlock(&state.mutex);
 
        printf("Assistant: ");
        fflush(stdout);
 
        // UI Detokenizer Loop
        while (1) {
            pthread_mutex_lock(&state.mutex);
            while (!state.token_ready && !state.quit) {
                pthread_cond_wait(&state.cond_done, &state.mutex);
            }
            int32_t tok = state.last_token;
            state.token_ready = false;
            pthread_mutex_unlock(&state.mutex);
 
            if (tok == 151643 || tok == 151645) break;
 
            const char *word = ck_tokenizer_id_to_token(tokenizer, tok);
            if (word) {
                if ((unsigned char)word[0] == 0xC4 && (unsigned char)word[1] == 0xA0) {
                    printf(" %s", word + 2);
                } else {
                    printf("%s", word);
                }
                fflush(stdout);
            }
        }
        printf("\n");
    }
 
    pthread_mutex_lock(&state.mutex);
    state.quit = true;
    pthread_cond_signal(&state.cond_task);
    pthread_mutex_unlock(&state.mutex);
    pthread_join(engine_thread, NULL);
 
    return 0;
}

References CK_TOKENIZER_BPE, ck_tokenizer_create(), ck_tokenizer_encode(), ck_tokenizer_id_to_token(), ck_tokenizer_load_binary(), engine_thread_func(), and ids.