Vocab.php
145 lines
| 1 | <?php |
| 2 | |
| 3 | declare(strict_types=1); |
| 4 | |
| 5 | namespace Yethee\Tiktoken\Vocab; |
| 6 | |
| 7 | use Countable; |
| 8 | use InvalidArgumentException; |
| 9 | use OutOfBoundsException; |
| 10 | use RuntimeException; |
| 11 | use Yethee\Tiktoken\Exception\ParseError; |
| 12 | use Yethee\Tiktoken\Util\EncodeUtil; |
| 13 | |
| 14 | use function array_flip; |
| 15 | use function array_map; |
| 16 | use function assert; |
| 17 | use function base64_decode; |
| 18 | use function count; |
| 19 | use function explode; |
| 20 | use function fclose; |
| 21 | use function fgets; |
| 22 | use function file_exists; |
| 23 | use function fopen; |
| 24 | use function implode; |
| 25 | use function rewind; |
| 26 | use function sprintf; |
| 27 | use function stream_get_meta_data; |
| 28 | use function strval; |
| 29 | |
| 30 | /** @psalm-import-type NonEmptyByteVector from EncodeUtil */ |
| 31 | final class Vocab implements Countable |
| 32 | { |
| 33 | /** @var array<non-empty-string, int> */ |
| 34 | private array $tokenToRankMap; |
| 35 | |
| 36 | /** @var array<int, non-empty-string> */ |
| 37 | private array $rankToTokenMap; |
| 38 | |
| 39 | /** @param array<non-empty-string, int> $tokenRankMap */ |
| 40 | private function __construct(array $tokenRankMap) |
| 41 | { |
| 42 | $this->tokenToRankMap = $tokenRankMap; |
| 43 | /** @psalm-suppress PropertyTypeCoercion */ |
| 44 | $this->rankToTokenMap = array_map(strval(...), array_flip($tokenRankMap)); |
| 45 | |
| 46 | if (count($this->tokenToRankMap) !== count($this->rankToTokenMap)) { |
| 47 | throw new InvalidArgumentException('The map of tokens and ranks has duplicates of rank'); |
| 48 | } |
| 49 | } |
| 50 | |
| 51 | /** @param non-empty-string $bpeFile */ |
| 52 | public static function fromFile(string $bpeFile): self |
| 53 | { |
| 54 | if (! file_exists($bpeFile)) { |
| 55 | throw new RuntimeException(sprintf('File "%s" does not exist', $bpeFile)); |
| 56 | } |
| 57 | |
| 58 | $stream = fopen($bpeFile, 'rb'); |
| 59 | |
| 60 | if ($stream === false) { |
| 61 | throw new RuntimeException(sprintf('Could not open file: %s', $bpeFile)); |
| 62 | } |
| 63 | |
| 64 | try { |
| 65 | return self::fromStream($stream); |
| 66 | } finally { |
| 67 | fclose($stream); |
| 68 | } |
| 69 | } |
| 70 | |
| 71 | /** |
| 72 | * @param resource $stream |
| 73 | * |
| 74 | * @return static |
| 75 | */ |
| 76 | public static function fromStream($stream): self |
| 77 | { |
| 78 | $meta = stream_get_meta_data($stream); |
| 79 | |
| 80 | if ($meta['seekable']) { |
| 81 | rewind($stream); |
| 82 | } |
| 83 | |
| 84 | $line = fgets($stream); |
| 85 | $lineNo = 1; |
| 86 | $map = []; |
| 87 | |
| 88 | while ($line !== false) { |
| 89 | [$encodedToken, $rank] = explode(' ', $line); |
| 90 | $token = base64_decode($encodedToken, true); |
| 91 | |
| 92 | if ($token === false) { |
| 93 | throw new ParseError(sprintf('Could not decode token "%s" at line %d', $encodedToken, $lineNo)); |
| 94 | } |
| 95 | |
| 96 | assert($token !== ''); |
| 97 | |
| 98 | $map[$token] = (int) $rank; |
| 99 | |
| 100 | $line = fgets($stream); |
| 101 | $lineNo++; |
| 102 | } |
| 103 | |
| 104 | return new self($map); |
| 105 | } |
| 106 | |
| 107 | public function tryGetRank(string $binary): int|null |
| 108 | { |
| 109 | if ($binary === '') { |
| 110 | throw new InvalidArgumentException('Argument $binary cannot be an empty string'); |
| 111 | } |
| 112 | |
| 113 | return $this->tokenToRankMap[$binary] ?? null; |
| 114 | } |
| 115 | |
| 116 | /** @throws OutOfBoundsException */ |
| 117 | public function getRank(string $binary): int |
| 118 | { |
| 119 | if ($binary === '') { |
| 120 | throw new InvalidArgumentException('Argument $binary cannot be an empty string'); |
| 121 | } |
| 122 | |
| 123 | return $this->tokenToRankMap[$binary] ?? throw new OutOfBoundsException(sprintf( |
| 124 | 'No rank for bytes vector: [%s]', |
| 125 | implode(', ', EncodeUtil::toBytes($binary)), |
| 126 | )); |
| 127 | } |
| 128 | |
| 129 | /** |
| 130 | * @return non-empty-string |
| 131 | * |
| 132 | * @throws OutOfBoundsException |
| 133 | */ |
| 134 | public function getToken(int $rank): string |
| 135 | { |
| 136 | return $this->rankToTokenMap[$rank] ?? throw new OutOfBoundsException(sprintf('No token for rank: %d', $rank)); |
| 137 | } |
| 138 | |
| 139 | /** @psalm-api */ |
| 140 | public function count(): int |
| 141 | { |
| 142 | return count($this->tokenToRankMap); |
| 143 | } |
| 144 | } |
| 145 |