llama.vocab source code

1 module llama.vocab;
2 
3 import llama.llama;
4 
5 // Drop const for C APIs that don't take const vocab pointer.
6 private llama_vocab* mutableVocab(const(llama_vocab)* vocab) @trusted @nogc nothrow
7 {
8     return cast(llama_vocab*) vocab;
9 }
10 
11 /// Split `text` into tokens. Returns a GC-allocated slice.
12 llama_token[] tokenize(const(llama_vocab)* vocab, const(char)[] text,
13                        bool addSpecial = true, bool parseSpecial = true) @trusted
14 {
15     if (text.length == 0) return null;
16     auto v = mutableVocab(vocab);
17     int n = -llama_tokenize(v, text.ptr, cast(int) text.length, null, 0, addSpecial, parseSpecial);
18     if (n <= 0) return null;
19     auto tokens = new llama_token[](n);
20     int result = llama_tokenize(v, text.ptr, cast(int) text.length,
21                                 tokens.ptr, cast(int) tokens.length, addSpecial, parseSpecial);
22     return result < 0 ? null : tokens[0 .. result];
23 }
24 
25 /// The string piece for a single token.
26 string tokenToString(const(llama_vocab)* vocab, llama_token token) @trusted
27 {
28     char[256] buf;
29     int n = llama_token_to_piece(mutableVocab(vocab), token, buf.ptr, cast(int) buf.length, 0, true);
30     return n < 0 ? null : buf[0 .. n].idup;
31 }
32 
33 /// Decode a token sequence back into text.
34 string detokenize(const(llama_vocab)* vocab, const(llama_token)[] tokens,
35                   bool removeSpecial = false, bool unparseSpecial = false) @trusted
36 {
37     if (tokens.length == 0) return "";
38     auto v = mutableVocab(vocab);
39     int bufSize = cast(int)(tokens.length * 8 + 64);
40     char[] buf = new char[](bufSize);
41     int n = llama_detokenize(v, tokens.ptr, cast(int) tokens.length,
42                              buf.ptr, cast(int) buf.length, removeSpecial, unparseSpecial);
43     if (n < 0)
44     {
45         buf = new char[](-n + 1);
46         n = llama_detokenize(v, tokens.ptr, cast(int) tokens.length,
47                              buf.ptr, cast(int) buf.length, removeSpecial, unparseSpecial);
48         if (n < 0) return "";
49     }
50     return buf[0 .. n].idup;
51 }
52 
53 llama_token bosToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_bos(mutableVocab(vocab)); } /// BOS token.
54 llama_token eosToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_eos(mutableVocab(vocab)); } /// EOS token.
55 llama_token eotToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_eot(mutableVocab(vocab)); } /// EOT (end-of-turn) token.
56 llama_token nlToken (const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_nl (mutableVocab(vocab)); } /// Newline token.
57 llama_token padToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_pad(mutableVocab(vocab)); } /// Padding token.
58 llama_token sepToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_sep(mutableVocab(vocab)); } /// Sentence separator token.
59 
60 // Fill-in-the-Middle special tokens.
61 llama_token fimPreToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_fim_pre(mutableVocab(vocab)); } /// FIM prefix token.
62 llama_token fimSufToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_fim_suf(mutableVocab(vocab)); } /// FIM suffix token.
63 llama_token fimMidToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_fim_mid(mutableVocab(vocab)); } /// FIM middle token.
64 llama_token fimPadToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_fim_pad(mutableVocab(vocab)); } /// FIM padding token.
65 llama_token fimRepToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_fim_rep(mutableVocab(vocab)); } /// FIM repo token.
66 llama_token fimSepToken(const(llama_vocab)* vocab) @trusted @nogc nothrow { return llama_vocab_fim_sep(mutableVocab(vocab)); } /// FIM separator token.
67 
68 /// Vocabulary type as int (compare to `LLAMA_VOCAB_TYPE_*` constants).
69 int vocabType(const(llama_vocab)* vocab) @trusted @nogc nothrow
70 {
71     return cast(int) llama_vocab_type(mutableVocab(vocab));
72 }
73 
74 /// Raw text piece for a token (pointer into model memory; do not free).
75 const(char)* tokenText(const(llama_vocab)* vocab, llama_token token) @trusted @nogc nothrow
76 {
77     return llama_vocab_get_text(mutableVocab(vocab), token);
78 }
79 
80 /// Log-probability score stored for a token in the vocab.
81 float tokenScore(const(llama_vocab)* vocab, llama_token token) @trusted @nogc nothrow
82 {
83     return llama_vocab_get_score(mutableVocab(vocab), token);
84 }
85 
86 /// Token attribute flags (control, normal, byte, etc.).
87 llama_token_attr tokenAttr(const(llama_vocab)* vocab, llama_token token) @trusted @nogc nothrow
88 {
89     return llama_vocab_get_attr(mutableVocab(vocab), token);
90 }
91 
92 /// True if the token is a control token (not renderable text).
93 bool isControl(const(llama_vocab)* vocab, llama_token token) @trusted @nogc nothrow
94 {
95     return llama_vocab_is_control(mutableVocab(vocab), token);
96 }
97 
98 /// True if the token signals end of generation.
99 bool isEog(const(llama_vocab)* vocab, llama_token token) @trusted @nogc nothrow
100 {
101     return llama_vocab_is_eog(mutableVocab(vocab), token);
102 }