Added saveToDisk method to NER

ankane · Aug 29, 2022 · 3abdc4e · 3abdc4e
1 parent 3bd7308
commit 3abdc4e
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 0.1.2 (unreleased)
 
 - Added support for text categorization
+- Added `saveToDisk` method to `NER`
 
 ## 0.1.1 (2022-08-28)
 

diff --git a/src/NER.php b/src/NER.php
@@ -40,6 +40,13 @@ public function entities($text)
         return $this->doc($text)->entities();
     }
 
+    public function saveToDisk($filename)
+    {
+        if ($this->ffi->mitie_save_named_entity_extractor($filename, $this->pointer) != 0) {
+            throw new Exception('Unable to save model');
+        }
+    }
+
     public function tokens($text)
     {
         return $this->doc($text)->tokens();

diff --git a/tests/NERTest.php b/tests/NERTest.php
@@ -43,4 +43,20 @@ public function testMissingFile()
 
         new Mitie\NER('missing.dat');
     }
+
+    public function testSaveToDisk()
+    {
+        $path = tempnam(sys_get_temp_dir(), 'model');
+        $this->model()->saveToDisk($path);
+        $this->assertFileExists($path);
+        unlink($path);
+    }
+
+    public function testSaveToDiskError()
+    {
+        $this->expectException(Mitie\Exception::class);
+        $this->expectExceptionMessage('Unable to save model');
+
+        $this->model()->saveToDisk('missing/ner_model.dat');
+    }
 }