Add test for Pipeline save and restore with ModelManager (#191)

aa753421 · Jan 12, 2018 · 7435bec · 7435bec
1 parent d953ef6
commit 7435bec
Showing 1 changed file with 37 additions and 0 deletions.
diff --git a/tests/Phpml/PipelineTest.php b/tests/Phpml/PipelineTest.php
@@ -7,6 +7,7 @@
 use Phpml\Classification\SVC;
 use Phpml\FeatureExtraction\TfIdfTransformer;
 use Phpml\FeatureExtraction\TokenCountVectorizer;
+use Phpml\ModelManager;
 use Phpml\Pipeline;
 use Phpml\Preprocessing\Imputer;
 use Phpml\Preprocessing\Imputer\Strategy\MostFrequentStrategy;
@@ -104,4 +105,40 @@ public function testPipelineTransformers(): void
 
         $this->assertEquals($expected, $predicted);
     }
+
+    public function testSaveAndRestore(): void
+    {
+        $pipeline = new Pipeline([
+            new TokenCountVectorizer(new WordTokenizer()),
+            new TfIdfTransformer(),
+        ], new SVC());
+
+        $pipeline->train([
+            'Hello Paul',
+            'Hello Martin',
+            'Goodbye Tom',
+            'Hello John',
+            'Goodbye Alex',
+            'Bye Tony',
+        ], [
+            'greetings',
+            'greetings',
+            'farewell',
+            'greetings',
+            'farewell',
+            'farewell',
+        ]);
+
+        $testSamples = ['Hello Max', 'Goodbye Mark'];
+        $predicted = $pipeline->predict($testSamples);
+
+        $filepath = tempnam(sys_get_temp_dir(), uniqid('pipeline-test', true));
+        $modelManager = new ModelManager();
+        $modelManager->saveToFile($pipeline, $filepath);
+
+        $restoredClassifier = $modelManager->restoreFromFile($filepath);
+        $this->assertEquals($pipeline, $restoredClassifier);
+        $this->assertEquals($predicted, $restoredClassifier->predict($testSamples));
+        unlink($filepath);
+    }
 }