deduplicate

onnx · Aug 23, 2024 · daa04f7 · daa04f7
1 parent ba89e19
commit daa04f7
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 132 deletions.
diff --git a/src/turnkeyml/common/test_helpers.py b/src/turnkeyml/common/test_helpers.py
@@ -100,6 +100,72 @@ def forward(self, x):
 }
 
 
+extras_dot_py = {
+    "compiled.py": """
+# labels: name::linear author::selftest test_group::selftest task::test
+import torch
+
+torch.manual_seed(0)
+
+
+class LinearTestModel(torch.nn.Module):
+    def __init__(self, input_features, output_features):
+        super(LinearTestModel, self).__init__()
+        self.fc = torch.nn.Linear(input_features, output_features)
+
+    def forward(self, x):
+        output = self.fc(x)
+        return output
+
+
+input_features = 10
+output_features = 10
+
+# Compiled model
+model = LinearTestModel(input_features, output_features)
+model = torch.compile(model)
+inputs = {"x": torch.rand(input_features)}
+model(**inputs)
+
+# Non-compiled model
+model2 = LinearTestModel(input_features * 2, output_features)
+inputs2 = {"x": torch.rand(input_features * 2)}
+model2(**inputs2)
+    """,
+    "selected_models.txt": f"""
+    {os.path.join(corpus_dir,"linear.py")}
+    {os.path.join(corpus_dir,"linear2.py")}
+    """,
+    "timeout.py": """
+# labels: name::timeout author::turnkey license::mit test_group::a task::test
+import torch
+
+torch.manual_seed(0)
+
+
+class LinearTestModel(torch.nn.Module):
+    def __init__(self, input_features, output_features):
+        super(LinearTestModel, self).__init__()
+        self.fc = torch.nn.Linear(input_features, output_features)
+
+    def forward(self, x):
+        output = self.fc(x)
+        return output
+
+
+input_features = 500000
+output_features = 1000
+
+# Model and input configurations
+model = LinearTestModel(input_features, output_features)
+inputs = {"x": torch.rand(input_features)}
+
+output = model(**inputs)
+
+    """,
+}
+
+
 def create_test_dir(
     key: str,
     test_scripts: Dict = None,

diff --git a/test/benchmark.py b/test/benchmark.py
@@ -618,75 +618,10 @@ def test_011_cli_cache_move(self):
     cache_dir, corpus_dir = common.create_test_dir("benchmark")
     new_cache_dir = f"{cache_dir}2"
 
-    extras_dot_py = {
-        "compiled.py": """
-# labels: name::linear author::selftest test_group::selftest task::test
-import torch
-
-torch.manual_seed(0)
-
-
-class LinearTestModel(torch.nn.Module):
-    def __init__(self, input_features, output_features):
-        super(LinearTestModel, self).__init__()
-        self.fc = torch.nn.Linear(input_features, output_features)
-
-    def forward(self, x):
-        output = self.fc(x)
-        return output
-
-
-input_features = 10
-output_features = 10
-
-# Compiled model
-model = LinearTestModel(input_features, output_features)
-model = torch.compile(model)
-inputs = {"x": torch.rand(input_features)}
-model(**inputs)
-
-# Non-compiled model
-model2 = LinearTestModel(input_features * 2, output_features)
-inputs2 = {"x": torch.rand(input_features * 2)}
-model2(**inputs2)
-    """,
-        "selected_models.txt": f"""
-    {os.path.join(corpus_dir,"linear.py")}
-    {os.path.join(corpus_dir,"linear2.py")}
-    """,
-        "timeout.py": """
-# labels: name::timeout author::turnkey license::mit test_group::a task::test
-import torch
-
-torch.manual_seed(0)
-
-
-class LinearTestModel(torch.nn.Module):
-    def __init__(self, input_features, output_features):
-        super(LinearTestModel, self).__init__()
-        self.fc = torch.nn.Linear(input_features, output_features)
-
-    def forward(self, x):
-        output = self.fc(x)
-        return output
-
-
-input_features = 500000
-output_features = 1000
-
-# Model and input configurations
-model = LinearTestModel(input_features, output_features)
-inputs = {"x": torch.rand(input_features)}
-
-output = model(**inputs)
-
-    """,
-    }
-
     extras_dir = os.path.join(corpus_dir, "extras")
     os.makedirs(extras_dir, exist_ok=True)
 
-    for key, value in extras_dot_py.items():
+    for key, value in common.extras_dot_py.items():
         file_path = os.path.join(extras_dir, key)
 
         with open(file_path, "w", encoding="utf") as f:

diff --git a/test/cli.py b/test/cli.py
@@ -877,75 +877,10 @@ def test_022_cli_cache_move(self):
     cache_dir, corpus_dir = common.create_test_dir("cli")
     new_cache_dir = f"{cache_dir}2"
 
-    extras_dot_py = {
-        "compiled.py": """
-# labels: name::linear author::selftest test_group::selftest task::test
-import torch
-
-torch.manual_seed(0)
-
-
-class LinearTestModel(torch.nn.Module):
-    def __init__(self, input_features, output_features):
-        super(LinearTestModel, self).__init__()
-        self.fc = torch.nn.Linear(input_features, output_features)
-
-    def forward(self, x):
-        output = self.fc(x)
-        return output
-
-
-input_features = 10
-output_features = 10
-
-# Compiled model
-model = LinearTestModel(input_features, output_features)
-model = torch.compile(model)
-inputs = {"x": torch.rand(input_features)}
-model(**inputs)
-
-# Non-compiled model
-model2 = LinearTestModel(input_features * 2, output_features)
-inputs2 = {"x": torch.rand(input_features * 2)}
-model2(**inputs2)
-    """,
-        "selected_models.txt": f"""
-    {os.path.join(corpus_dir,"linear.py")}
-    {os.path.join(corpus_dir,"linear2.py")}
-    """,
-        "timeout.py": """
-# labels: name::timeout author::turnkey license::mit test_group::a task::test
-import torch
-
-torch.manual_seed(0)
-
-
-class LinearTestModel(torch.nn.Module):
-    def __init__(self, input_features, output_features):
-        super(LinearTestModel, self).__init__()
-        self.fc = torch.nn.Linear(input_features, output_features)
-
-    def forward(self, x):
-        output = self.fc(x)
-        return output
-
-
-input_features = 500000
-output_features = 1000
-
-# Model and input configurations
-model = LinearTestModel(input_features, output_features)
-inputs = {"x": torch.rand(input_features)}
-
-output = model(**inputs)
-
-    """,
-    }
-
     extras_dir = os.path.join(corpus_dir, "extras")
     os.makedirs(extras_dir, exist_ok=True)
 
-    for key, value in extras_dot_py.items():
+    for key, value in common.extras_dot_py.items():
         file_path = os.path.join(extras_dir, key)
 
         with open(file_path, "w", encoding="utf") as f: