llama.vim : add processing info overlay

This commit is contained in:
Georgi Gerganov 2024-10-09 15:08:31 +03:00
parent 07e7dd47f2
commit 9d13e87b1b
No known key found for this signature in database
GPG key ID: BF970631944C16B7

View file

@ -11,12 +11,13 @@
" "
" color of the suggested text " color of the suggested text
highlight llama_hint guifg=#ff772f highlight llama_hl_hint guifg=#ff772f
highlight llama_hl_info guifg=#77ff2f
let s:default_config = { let s:default_config = {
\ 'endpoint': 'http://127.0.0.1:8012/infill', \ 'endpoint': 'http://127.0.0.1:8012/infill',
\ 'n_prefix': 32, \ 'n_prefix': 128,
\ 'n_suffix': 32, \ 'n_suffix': 128,
\ 'n_predict': 64, \ 'n_predict': 64,
\ 'n_probs': 3, \ 'n_probs': 3,
\ 'temperature': 0.1, \ 'temperature': 0.1,
@ -71,6 +72,16 @@ function! llama#fim() abort
\ ) \ )
let l:can_accept = v:true let l:can_accept = v:true
let l:has_timing = v:false
let l:n_prompt = 0
let l:t_prompt_ms = 1.0
let l:s_prompt = 0
let l:n_gen = 0
let l:t_gen_ms = 1.0
let l:s_gen = 0
let s:content = [] let s:content = []
let l:raw = system(l:curl_command) let l:raw = system(l:curl_command)
@ -96,6 +107,20 @@ function! llama#fim() abort
while len(s:content) > 0 && s:content[-1] == "" while len(s:content) > 0 && s:content[-1] == ""
call remove(s:content, -1) call remove(s:content, -1)
endwhile endwhile
" if response.timings
if len(get(l:response, 'timings', {})) > 0
let l:has_timing = v:true
let l:timings = get(l:response, 'timings', {})
let l:n_prompt = get(l:timings, 'prompt_n', 0)
let l:t_prompt_ms = get(l:timings, 'prompt_ms', 1)
let l:s_prompt = get(l:timings, 'prompt_per_second', 0)
let l:n_gen = get(l:timings, 'predicted_n', 0)
let l:t_gen_ms = get(l:timings, 'predicted_ms', 1)
let l:s_gen = get(l:timings, 'predicted_per_second', 0)
endif
endif endif
if len(s:content) == 0 if len(s:content) == 0
@ -108,48 +133,62 @@ function! llama#fim() abort
" display virtual text with the suggestion " display virtual text with the suggestion
let l:bufnr = bufnr('%') let l:bufnr = bufnr('%')
let s:ns_id = nvim_create_namespace('llama_virtual_text')
call nvim_buf_set_extmark(l:bufnr, s:ns_id, l:pos_y - 1, l:pos_x - 1, { let s:id_vt_fim = nvim_create_namespace('vt_fim')
\ 'virt_text': [[s:content[0], 'llama_hint']], let s:id_vt_info = nvim_create_namespace('vt_info')
call nvim_buf_set_extmark(l:bufnr, s:id_vt_fim, l:pos_y - 1, l:pos_x - 1, {
\ 'virt_text': [[s:content[0], 'llama_hl_hint']],
\ 'virt_text_win_col': l:pos_x == 1 ? 0 : virtcol('.') \ 'virt_text_win_col': l:pos_x == 1 ? 0 : virtcol('.')
\ }) \ })
call nvim_buf_set_extmark(l:bufnr, s:ns_id, l:pos_y - 1, 0, { call nvim_buf_set_extmark(l:bufnr, s:id_vt_fim, l:pos_y - 1, 0, {
\ 'virt_lines': map(s:content[1:], {idx, val -> [[val, 'llama_hint']]}), \ 'virt_lines': map(s:content[1:], {idx, val -> [[val, 'llama_hl_hint']]}),
\ 'virt_text_win_col': virtcol('.') \ 'virt_text_win_col': virtcol('.')
\ }) \ })
" construct the info message:
if l:has_timing
let l:info = printf("prompt: %d (%.2f ms, %.2f t/s) | predict: %d (%.2f ms, %.2f t/s)",
\ l:n_prompt, l:t_prompt_ms, l:s_prompt,
\ l:n_gen, l:t_gen_ms, l:s_gen)
call nvim_buf_set_extmark(l:bufnr, s:id_vt_info, l:pos_y - 1, l:pos_x - 1, {
\ 'virt_text': [[l:info, 'llama_hl_info']],
\ 'virt_text_pos': 'right_align',
\ })
endif
" accept suggestion with Tab and reject it with any other key " accept suggestion with Tab and reject it with any other key
if l:can_accept if l:can_accept
inoremap <buffer> <Tab> <C-O>:call llama#accept_virtual_text()<CR> inoremap <buffer> <Tab> <C-O>:call llama#accept_vt_fim()<CR>
else else
inoremap <buffer> <Tab> <C-O>:call llama#cancel_virtual_text()<CR> inoremap <buffer> <Tab> <C-O>:call llama#cancel_vt_fim()<CR>
endif endif
for l:key in range(32, 127) + [8, 27] for l:key in range(32, 127) + [8, 27]
if l:key != 0x7C if l:key != 0x7C
if l:key == 8 if l:key == 8
execute 'inoremap <buffer> <Bs> <C-O>:call llama#cancel_virtual_text()<CR><Bs>' execute 'inoremap <buffer> <Bs> <C-O>:call llama#cancel_vt_fim()<CR><Bs>'
elseif l:key == 27 elseif l:key == 27
execute 'inoremap <buffer> <Esc> <C-O>:call llama#cancel_virtual_text()<CR><Esc>' execute 'inoremap <buffer> <Esc> <C-O>:call llama#cancel_vt_fim()<CR><Esc>'
elseif l:key == 32 elseif l:key == 32
execute 'inoremap <buffer> <Space> <C-O>:call llama#cancel_virtual_text()<CR><Space>' execute 'inoremap <buffer> <Space> <C-O>:call llama#cancel_vt_fim()<CR><Space>'
elseif l:key == 127 elseif l:key == 127
execute 'inoremap <buffer> <Del> <C-O>:call llama#cancel_virtual_text()<CR><Del>' execute 'inoremap <buffer> <Del> <C-O>:call llama#cancel_vt_fim()<CR><Del>'
else else
execute 'inoremap <buffer> ' . nr2char(l:key) . ' <C-O>:call llama#cancel_virtual_text()<CR>' . nr2char(l:key) execute 'inoremap <buffer> ' . nr2char(l:key) . ' <C-O>:call llama#cancel_vt_fim()<CR>' . nr2char(l:key)
endif endif
endif endif
endfor endfor
inoremap <buffer> <Up> <C-O>:call llama#cancel_virtual_text()<CR><Up> inoremap <buffer> <Up> <C-O>:call llama#cancel_vt_fim()<CR><Up>
inoremap <buffer> <Down> <C-O>:call llama#cancel_virtual_text()<CR><Down> inoremap <buffer> <Down> <C-O>:call llama#cancel_vt_fim()<CR><Down>
inoremap <buffer> <Left> <C-O>:call llama#cancel_virtual_text()<CR><Left> inoremap <buffer> <Left> <C-O>:call llama#cancel_vt_fim()<CR><Left>
inoremap <buffer> <Right> <C-O>:call llama#cancel_virtual_text()<CR><Right> inoremap <buffer> <Right> <C-O>:call llama#cancel_vt_fim()<CR><Right>
endfunction endfunction
function! llama#accept_virtual_text() function! llama#accept_vt_fim()
let l:pos_x = col('.') let l:pos_x = col('.')
let l:pos_y = line('.') let l:pos_y = line('.')
@ -166,13 +205,15 @@ function! llama#accept_virtual_text()
" move the cursor to the end of the accepted text " move the cursor to the end of the accepted text
call cursor(l:pos_y + len(s:content) - 1, l:pos_x + s:pos_dx) call cursor(l:pos_y + len(s:content) - 1, l:pos_x + s:pos_dx)
call llama#cancel_virtual_text() call llama#cancel_vt_fim()
endfunction endfunction
function! llama#cancel_virtual_text() function! llama#cancel_vt_fim()
" clear the virtual text " clear the virtual text
let l:bufnr = bufnr('%') let l:bufnr = bufnr('%')
call nvim_buf_clear_namespace(l:bufnr, s:ns_id, 0, -1)
call nvim_buf_clear_namespace(l:bufnr, s:id_vt_fim, 0, -1)
call nvim_buf_clear_namespace(l:bufnr, s:id_vt_info, 0, -1)
" remove the mappings " remove the mappings
iunmap <buffer> <Tab> iunmap <buffer> <Tab>