Exception
10 months ago
Util
10 months ago
Vocab
10 months ago
Encoder.php
10 months ago
EncoderProvider.php
10 months ago
Encoder.php
201 lines
| 1 | <?php |
| 2 | |
| 3 | declare(strict_types=1); |
| 4 | |
| 5 | namespace Yethee\Tiktoken; |
| 6 | |
| 7 | use Stringable; |
| 8 | use Yethee\Tiktoken\Exception\RegexError; |
| 9 | use Yethee\Tiktoken\Util\EncodeUtil; |
| 10 | use Yethee\Tiktoken\Vocab\Vocab; |
| 11 | |
| 12 | use function array_map; |
| 13 | use function array_merge; |
| 14 | use function array_values; |
| 15 | use function count; |
| 16 | use function implode; |
| 17 | use function preg_last_error_msg; |
| 18 | use function preg_match_all; |
| 19 | use function sprintf; |
| 20 | use function strlen; |
| 21 | use function substr; |
| 22 | |
| 23 | use const PHP_INT_MAX; |
| 24 | |
| 25 | /** @psalm-import-type NonEmptyByteVector from EncodeUtil */ |
| 26 | final class Encoder implements Stringable |
| 27 | { |
| 28 | /** |
| 29 | * @param non-empty-string $name |
| 30 | * @param non-empty-string $pattern |
| 31 | */ |
| 32 | public function __construct(public readonly string $name, private Vocab $vocab, private string $pattern) |
| 33 | { |
| 34 | } |
| 35 | |
| 36 | public function __toString(): string |
| 37 | { |
| 38 | return sprintf('Encoder(name="%s", vocab=%d)', $this->name, count($this->vocab)); |
| 39 | } |
| 40 | |
| 41 | /** @return list<int> */ |
| 42 | public function encode(string $text): array |
| 43 | { |
| 44 | if ($text === '') { |
| 45 | return []; |
| 46 | } |
| 47 | |
| 48 | if (preg_match_all($this->pattern, $text, $matches) === false) { |
| 49 | throw new RegexError(sprintf('Matching failed with error: %s', preg_last_error_msg())); |
| 50 | } |
| 51 | |
| 52 | $tokens = []; |
| 53 | |
| 54 | foreach ($matches[0] as $match) { |
| 55 | if ($match === '') { |
| 56 | continue; |
| 57 | } |
| 58 | |
| 59 | $rank = $this->vocab->tryGetRank($match); |
| 60 | |
| 61 | if ($rank !== null) { |
| 62 | $tokens[] = $rank; |
| 63 | |
| 64 | continue; |
| 65 | } |
| 66 | |
| 67 | foreach ($this->mergeBytePairs($match) as $rank) { |
| 68 | $tokens[] = $rank; |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | return $tokens; |
| 73 | } |
| 74 | |
| 75 | /** |
| 76 | * Encodes a given text into chunks of Byte-Pair Encoded (BPE) tokens, with each chunk containing a specified |
| 77 | * maximum number of tokens. |
| 78 | * |
| 79 | * @param string $text The input text to be encoded. |
| 80 | * @param positive-int $maxTokensPerChunk The maximum number of tokens allowed per chunk. |
| 81 | * |
| 82 | * @return list<list<int>> An array of arrays containing BPE token chunks. |
| 83 | */ |
| 84 | public function encodeInChunks(string $text, int $maxTokensPerChunk): array |
| 85 | { |
| 86 | if ($text === '') { |
| 87 | return []; |
| 88 | } |
| 89 | |
| 90 | if (preg_match_all($this->pattern, $text, $matches) === false) { |
| 91 | throw new RegexError(sprintf('Matching failed with error: %s', preg_last_error_msg())); |
| 92 | } |
| 93 | |
| 94 | $chunks = []; |
| 95 | $tokensInCurrentChunk = []; |
| 96 | |
| 97 | foreach ($matches[0] as $match) { |
| 98 | if ($match === '') { |
| 99 | continue; |
| 100 | } |
| 101 | |
| 102 | $rank = $this->vocab->tryGetRank($match); |
| 103 | $tokens = $rank !== null ? [$rank] : $this->mergeBytePairs($match); |
| 104 | |
| 105 | if (count($tokensInCurrentChunk) + count($tokens) > $maxTokensPerChunk) { |
| 106 | $chunks[] = $tokensInCurrentChunk; |
| 107 | $tokensInCurrentChunk = []; |
| 108 | } |
| 109 | |
| 110 | $tokensInCurrentChunk = array_merge($tokensInCurrentChunk, $tokens); |
| 111 | } |
| 112 | |
| 113 | if (count($tokensInCurrentChunk) > 0) { |
| 114 | $chunks[] = $tokensInCurrentChunk; |
| 115 | } |
| 116 | |
| 117 | return $chunks; |
| 118 | } |
| 119 | |
| 120 | /** @param array<int> $tokens */ |
| 121 | public function decode(array $tokens): string |
| 122 | { |
| 123 | if ($tokens === []) { |
| 124 | return ''; |
| 125 | } |
| 126 | |
| 127 | return implode(array_map($this->vocab->getToken(...), $tokens)); |
| 128 | } |
| 129 | |
| 130 | /** |
| 131 | * @param non-empty-string $piece |
| 132 | * |
| 133 | * @return list<int> |
| 134 | */ |
| 135 | private function mergeBytePairs(string $piece): array |
| 136 | { |
| 137 | $parts = []; |
| 138 | |
| 139 | for ($i = 0; $i <= strlen($piece); $i++) { |
| 140 | $parts[] = [$i, PHP_INT_MAX]; |
| 141 | } |
| 142 | |
| 143 | $getRank = function (array $parts, int $startIndex, int $skip = 0) use (&$piece): int { |
| 144 | if (($startIndex + $skip + 2) >= count($parts)) { |
| 145 | return PHP_INT_MAX; |
| 146 | } |
| 147 | |
| 148 | $offset = $parts[$startIndex][0]; |
| 149 | $length = $parts[$startIndex + $skip + 2][0] - $offset; |
| 150 | |
| 151 | return $this->vocab->tryGetRank(substr($piece, $offset, $length)) ?? PHP_INT_MAX; |
| 152 | }; |
| 153 | |
| 154 | for ($i = 0; $i < count($parts) - 2; $i++) { |
| 155 | $parts[$i][1] = $getRank($parts, $i); |
| 156 | } |
| 157 | |
| 158 | while (count($parts) > 1) { |
| 159 | $minRank = PHP_INT_MAX; |
| 160 | $partIndex = 0; |
| 161 | $stop = count($parts) - 1; |
| 162 | |
| 163 | for ($i = 0; $i < $stop; $i++) { |
| 164 | if ($minRank <= $parts[$i][1]) { |
| 165 | continue; |
| 166 | } |
| 167 | |
| 168 | $minRank = $parts[$i][1]; |
| 169 | $partIndex = $i; |
| 170 | } |
| 171 | |
| 172 | if ($minRank === PHP_INT_MAX) { |
| 173 | break; |
| 174 | } |
| 175 | |
| 176 | unset($parts[$partIndex + 1]); |
| 177 | $parts = array_values($parts); |
| 178 | |
| 179 | $parts[$partIndex][1] = $getRank($parts, $partIndex); |
| 180 | |
| 181 | if ($partIndex <= 0) { |
| 182 | continue; |
| 183 | } |
| 184 | |
| 185 | $parts[$partIndex - 1][1] = $getRank($parts, $partIndex - 1); |
| 186 | } |
| 187 | |
| 188 | $stop = count($parts) - 1; |
| 189 | $res = []; |
| 190 | |
| 191 | for ($i = 0; $i < $stop; $i++) { |
| 192 | $offset = $parts[$i][0]; |
| 193 | $length = $parts[$i + 1][0] - $offset; |
| 194 | |
| 195 | $res[] = $this->vocab->getRank(substr($piece, $offset, $length)); |
| 196 | } |
| 197 | |
| 198 | return $res; |
| 199 | } |
| 200 | } |
| 201 |