Update print

modal-labs · Feb 15, 2024 · d8c8c3e · d8c8c3e
1 parent 0c9b41d
commit d8c8c3e
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/ci/check_inference.py b/ci/check_inference.py
@@ -10,8 +10,13 @@
 CREATE TABLE head (age INTEGER)
 How many heads of the departments are older than 56 ? [/INST] """
 
-    output = subprocess.check_output(["modal", "run", "src.inference", "--run-folder", f"/runs/{run_name}", "--prompt", prompt])
+    p = subprocess.Popen(["modal", "run", "src.inference", "--run-folder", f"/runs/{run_name}", "--prompt", prompt], stdout=subprocess.PIPE)
+    output = ""
 
-    print("Asserting that output contains [SQ] SELECT ... [/SQL]")
+    for line in iter(p.stdout.readline, b''):
+        output += line.decode()
+
+    print(output)
+    print("Asserting that the output contains the expected SQL query")
 
     assert b"[SQL] SELECT" in output and b"[/SQL]" in output