Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
| 1890 | |
| 1891 | // Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am) |
| 1892 | static void llama_sampler_dry_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) { |
| 1893 | auto * ctx = (llama_sampler_dry *) smpl->ctx; |
| 1894 | |
| 1895 | if (ctx->dry_multiplier == 0.0f || ctx->dry_base < 1.0f || ctx->dry_penalty_last_n == 0) { |
| 1896 | return; |
| 1897 | } |
| 1898 | |
| 1899 | int32_t effective_dry_penalty_last_n = (ctx->dry_penalty_last_n == -1) ? ctx->total_context_size : std::max(ctx->dry_penalty_last_n, 0); |
| 1900 | int last_n_repeat = std::min(std::min((int)ctx->last_tokens.size(), effective_dry_penalty_last_n), ctx->total_context_size); |
| 1901 | |
| 1902 | if (last_n_repeat <= ctx->dry_allowed_length) { |
| 1903 | return; |
| 1904 | } |
| 1905 | |
| 1906 | ctx->dry_repeat_count.assign(last_n_repeat, 0); |
| 1907 | ctx->dry_max_token_repeat.clear(); |
| 1908 | |
| 1909 | // Step 1: Look for restart sequences to limit the maximum repetition length. |
| 1910 | // Work backwards through the context looking for any token that begins a restart sequence. |
| 1911 | // |
| 1912 | // The collection `restart_sequences` is a mapping from a "head" token to all "tail" |
| 1913 | // sequences that together comprise a restart sequence. This allows us to quickly check |
| 1914 | // whether each token is the head of a complete sequence. Most restart sequences are actually |
| 1915 | // a single token, and for these the "tail" is an empty vector. |
| 1916 | // |
| 1917 | // If the token is a "head", test all restart sequences that begin with this token |
| 1918 | // (there will often only be one sequence for each token, but if sequences like 'aaaq1' and |
| 1919 | // 'aaa1' are used as restart strings, both could start with 'aaa' when tokenized). The |
| 1920 | // longest matching sequence (if any) is used to limit the maximum repetition length. |
| 1921 | // |
| 1922 | // Note that in the case case of a short sequence contained in a longer one, this might fail to |
| 1923 | // find the smallest value for `rep_limit`. For example, if 'amniotic' and 'ni' are both used as |
| 1924 | // restart sequences, 'ni' will be found first, and since it's shorter it will fail to suppress |
| 1925 | // 'otic'. This is a minor issue since fully contained restart sequences are likely to be rare. |
| 1926 | // |
| 1927 | // This is theoretically worst-case O(N^2) for arbitrary restart sequences, which is why we |
| 1928 | // have already clamped the maximum tail sequence length when generating `restart_sequences`. |
| 1929 | // With clamping, this scan is O(N) in the context length. |
| 1930 | |
| 1931 | int rep_limit = last_n_repeat; |
| 1932 | for (int i = 0; i < last_n_repeat; ++i) { |
| 1933 | llama_token token = ctx->last_tokens.rat(i); |
| 1934 | auto its = ctx->dry_processed_breakers.equal_range(token); |
| 1935 | if (its.first == ctx->dry_processed_breakers.end()) { |
| 1936 | continue; |
| 1937 | } |
| 1938 | int longest_match = -1; |
| 1939 | for (auto it = its.first; it != its.second; ++it) { |
| 1940 | // Note that (*it) does not contain the head character, so seq_len will be |
| 1941 | // the restart sequence length minus 1. |
| 1942 | // In the common case of a single-token restart sequence, (*it) will be empty |
| 1943 | // and we will trivially match. |
| 1944 | int seq_len = (int)it->second.size(); |
| 1945 | if (seq_len > longest_match && seq_len <= (int)i) { |
| 1946 | bool match = true; |
| 1947 | for (int offset = 0; offset < seq_len; ++offset) { |
| 1948 | // The -1 when indexing `last_tokens` is because we already matched the head. |
| 1949 | if (it->second[offset] != ctx->last_tokens.rat(i - offset - 1)) { |