test registerModelPrefixToEncoding

danny50610 · Aug 20, 2023 · a3cafc6 · a3cafc6
1 parent 6affbe7
commit a3cafc6
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/src/EncodingFactory.php b/src/EncodingFactory.php
@@ -3,7 +3,6 @@
 namespace Danny50610\BpeTokeniser;
 
 use Closure;
-use Exception;
 use InvalidArgumentException;
 use SplFileObject;
 
@@ -79,7 +78,7 @@ public static function registerModelToEncoding(string $modelName, string $encodi
     public static function registerModelPrefixToEncoding(string $modelPrefix, string $encodingName)
     {
         if (array_key_exists($modelPrefix, self::$modelPrefixToEncoding)) {
-            throw new Exception("{$modelPrefix} already exists");
+            throw new InvalidArgumentException("Prefix \"{$modelPrefix}\" already exists in map");
         }
 
         self::$modelPrefixToEncoding[$modelPrefix] = $encodingName;

diff --git a/tests/EncodingFactoryTest.php b/tests/EncodingFactoryTest.php
@@ -62,6 +62,22 @@ public function testRegisterModelToEncodingAlreadyExists()
         EncodingFactory::registerModelToEncoding('gpt-4', 'cl100k_base');
     }
 
+    public function testRegisterModelPrefixToEncoding()
+    {
+        EncodingFactory::registerModelPrefixToEncoding('gpt-9000-', 'cl100k_base');
+
+        $enc = EncodingFactory::createByModelName('gpt-9000-danny');
+        $this->assertSame('cl100k_base', $enc->getName());
+    }
+
+    public function testRegisterModelPrefixToEncodingAlreadyExists()
+    {
+        $this->expectException(InvalidArgumentException::class);
+        $this->expectExceptionMessage('Prefix "gpt-4-" already exists in map');
+
+        EncodingFactory::registerModelPrefixToEncoding('gpt-4-', 'cl100k_base');
+    }
+
     public function testRegisterEncoding()
     {
         EncodingFactory::registerEncoding('danny', function () {