Rename Model.run() to Model.predict()

Signed-off-by: Ben Firshman <[email protected]>

Author: bfirsh

README.md

@@ -25,7 +25,7 @@ class ColorizationModel(cog.Model):
         self.model = torch.load("./weights.pth")
 
     @cog.input("input", type=Path, help="Grayscale input image")
-    def run(self, input):
+    def predict(self, input):
         # ... pre-processing ...
         output = self.model(processed_input)
         # ... post-processing ...

docs/python.md

@@ -10,7 +10,7 @@ class HelloWorldModel(cog.Model):
         self.prefix = "hello "
 
     @cog.input("text", type=str, help="Text that will get prefixed by 'hello '")
-    def run(self, text):
+    def predict(self, text):
         return self.prefix + text
 ```
 

end-to-end-test/end_to_end_test/conftest.py

@@ -70,7 +70,7 @@ def setup(self):
     @cog.input("text", type=str)
     @cog.input("path", type=Path)
     @cog.input("output_file", type=bool, default=False)
-    def run(self, text, path, output_file):
+    def predict(self, text, path, output_file):
         time.sleep(1)
         with open(path) as f:
             output = self.foo + text + f.read()

end-to-end-test/end_to_end_test/test_server.py

@@ -101,7 +101,7 @@ def test_server_end_to_end(cog_server, project_dir, tmpdir_factory):
     ).communicate()
     paths = sorted(glob(str(download_dir / "*.*")))
     filenames = [os.path.basename(f) for f in paths]
-    assert filenames == ["cog.yaml", "predict.py", "myfile.txt"]
+    assert filenames == ["cog.yaml", "myfile.txt", "predict.py"]
 
     with open(download_dir / "cog-example-output/output.02.txt") as f:
         assert f.read() == "fooquxbaz"

pkg/docker/cog.py

@@ -39,7 +39,7 @@ def setup(self):
         pass
 
     @abstractmethod
-    def run(self, **kwargs):
+    def predict(self, **kwargs):
         pass
 
 
@@ -70,7 +70,7 @@ def handle_request():
                         )
                     raw_inputs[key] = val
 
-                if hasattr(self.model.run, "_inputs"):
+                if hasattr(self.model.predict, "_inputs"):
                     try:
                         inputs = validate_and_convert_inputs(
                             self.model, raw_inputs, cleanup_functions
@@ -97,8 +97,8 @@ def ping():
         @app.route("/help")
         def help():
             args = {}
-            if hasattr(self.model.run, "_inputs"):
-                input_specs = self.model.run._inputs
+            if hasattr(self.model.predict, "_inputs"):
+                input_specs = self.model.predict._inputs
                 for name, spec in input_specs.items():
                     arg: Dict[str, Any] = {
                         "type": _type_name(spec.type),
@@ -186,8 +186,8 @@ def ping():
         @app.route("/help")
         def help():
             args = {}
-            if hasattr(self.model.run, "_inputs"):
-                input_specs = self.model.run._inputs
+            if hasattr(self.model.predict, "_inputs"):
+                input_specs = self.model.predict._inputs
                 for name, spec in input_specs.items():
                     arg = {
                         "type": _type_name(spec.type),
@@ -420,7 +420,7 @@ def upload_to_temp(self, path: Path) -> str:
 def validate_and_convert_inputs(
     model: Model, raw_inputs: Dict[str, Any], cleanup_functions: List[Callable]
 ) -> Dict[str, Any]:
-    input_specs = model.run._inputs
+    input_specs = model.predict._inputs
     inputs = {}
 
     for name, input_spec in input_specs.items():
@@ -516,7 +516,7 @@ def run_model(model, inputs, cleanup_functions):
     Run the model on the inputs, and append resulting paths
     to cleanup functions for removal.
     """
-    result = model.run(**inputs)
+    result = model.predict(**inputs)
     if isinstance(result, Path):
         cleanup_functions.append(result.unlink)
     return result

pkg/docker/cog_test.py

@@ -26,7 +26,7 @@ class Model(cog.Model):
         def setup(self):
             self.foo = "foo"
 
-        def run(self):
+        def predict(self):
             return self.foo + "bar"
 
     client = make_client(Model())
@@ -41,7 +41,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("text", type=str)
-        def run(self, text):
+        def predict(self, text):
             return self.foo + text
 
     client = make_client(Model())
@@ -56,7 +56,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("text", type=str)
-        def run(self, text):
+        def predict(self, text):
             return self.foo + text
 
     client = make_client(Model())
@@ -73,7 +73,7 @@ def setup(self):
                 self.foo = "foo"
 
             @cog.input("text", type=str)
-            def run(self, bad):
+            def predict(self, bad):
                 return self.foo + "bar"
 
 
@@ -83,7 +83,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=int)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 3
             return self.foo + str(num2)
 
@@ -102,7 +102,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=int)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 2
             return self.foo + str(num2)
 
@@ -117,7 +117,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=int, default=5)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 2
             return self.foo + str(num2)
 
@@ -136,7 +136,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=float)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 3
             return self.foo + str(num2)
 
@@ -158,7 +158,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("num", type=float)
-        def run(self, num):
+        def predict(self, num):
             num2 = num ** 2
             return self.foo + str(num2)
 
@@ -173,7 +173,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("flag", type=bool)
-        def run(self, flag):
+        def predict(self, flag):
             if flag:
                 return self.foo + "yes"
             else:
@@ -194,7 +194,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("flag", type=bool)
-        def run(self, flag):
+        def predict(self, flag):
             if flag:
                 return self.foo + "yes"
             else:
@@ -213,7 +213,7 @@ def setup(self):
         @cog.input("num1", type=float, min=3, max=10.5)
         @cog.input("num2", type=float, min=-4)
         @cog.input("num3", type=int, max=-4)
-        def run(self, num1, num2, num3):
+        def predict(self, num1, num2, num3):
             return num1 + num2 + num3
 
     client = make_client(Model())
@@ -235,7 +235,7 @@ def setup(self):
 
         @cog.input("text", type=str, options=["foo", "bar"])
         @cog.input("num", type=int, options=[1, 2, 3])
-        def run(self, text, num):
+        def predict(self, text, num):
             return text + ("a" * num)
 
     client = make_client(Model())
@@ -252,7 +252,7 @@ def setup(self):
                 pass
 
             @cog.input("text", type=str, options=[])
-            def run(self, text):
+            def predict(self, text):
                 return text
 
     with pytest.raises(ValueError):
@@ -262,7 +262,7 @@ def setup(self):
                 pass
 
             @cog.input("text", type=str, options=["foo"])
-            def run(self, text):
+            def predict(self, text):
                 return text
 
     with pytest.raises(ValueError):
@@ -272,7 +272,7 @@ def setup(self):
                 pass
 
             @cog.input("text", type=Path, options=["foo"])
-            def run(self, text):
+            def predict(self, text):
                 return text
 
 
@@ -283,7 +283,7 @@ def setup(self):
 
         @cog.input("text", type=str, options=["foo", "bar"])
         @cog.input("num", type=int, options=[1, 2, 3])
-        def run(self, text, num):
+        def predict(self, text, num):
             return text + ("a" * num)
 
     client = make_client(Model())
@@ -299,7 +299,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("path", type=Path)
-        def run(self, path):
+        def predict(self, path):
             with open(path) as f:
                 return self.foo + " " + f.read() + " " + os.path.basename(path)
 
@@ -318,7 +318,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("path", type=Path)
-        def run(self, path):
+        def predict(self, path):
             with open(path) as f:
                 return self.foo + " " + f.read() + " " + os.path.basename(path)
 
@@ -333,7 +333,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("path", type=Path, default=None)
-        def run(self, path):
+        def predict(self, path):
             if path is None:
                 return "noneee"
             with open(path) as f:
@@ -357,7 +357,7 @@ def setup(self):
             self.foo = "foo"
 
         @cog.input("text", type=str)
-        def run(self, text):
+        def predict(self, text):
             temp_dir = tempfile.mkdtemp()
             temp_path = os.path.join(temp_dir, "my_file.txt")
             with open(temp_path, "w") as f:
@@ -376,7 +376,7 @@ class Model(cog.Model):
         def setup(self):
             pass
 
-        def run(self):
+        def predict(self):
             temp_dir = tempfile.mkdtemp()
             temp_path = os.path.join(temp_dir, "my_file.bmp")
             img = Image.new("RGB", (255, 255), "red")
@@ -400,7 +400,7 @@ def setup(self):
         @cog.input("num1", type=int)
         @cog.input("num2", type=int, default=10)
         @cog.input("path", type=Path)
-        def run(self, text, num1, num2, path):
+        def predict(self, text, num1, num2, path):
             with open(path) as f:
                 path_contents = f.read()
             return self.foo + " " + text + " " + str(num1 * num2) + " " + path_contents
@@ -425,7 +425,7 @@ def setup(self):
         @cog.input("num1", type=int, help="First number")
         @cog.input("num2", type=int, default=10, help="Second number")
         @cog.input("path", type=Path, help="A file path")
-        def run(self, text, num1, num2, path):
+        def predict(self, text, num1, num2, path):
             with open(path) as f:
                 path_contents = f.read()
             return self.foo + " " + text + " " + str(num1 * num2) + " " + path_contents
@@ -461,15 +461,15 @@ class ModelSlow(cog.Model):
         def setup(self):
             time.sleep(0.5)
 
-        def run(self):
+        def predict(self):
             time.sleep(0.5)
             return ""
 
     class ModelFast(cog.Model):
         def setup(self):
             pass
 
-        def run(self):
+        def predict(self):
             return ""
 
     client = make_client(ModelSlow())