compare result among 2 results

Author: yizhuoz004

tripy/examples/nanogpt/README.md

@@ -38,7 +38,10 @@ for expected accuracy.
     <!--
     ```
     (?s).*?
-    What is the answer to life, the universe, and everything\? How can we know what's important\? How can
+    What is the answer to life, the universe, and everything\? How can we know what's real\? How can
+    ====
+    (?s).*?
+    What is the answer to life, the universe, and everything\? The answer to the questions that are asked of us
     ```
      -->
     <!-- Tripy: TEST: EXPECTED_STDOUT End -->
@@ -62,7 +65,10 @@ To run with a quantization mode, pass `--quant-mode` to `example.py`. The suppor
     <!--
     ```
     (?s).*?
-    What is the answer to life, the universe, and everything\? How can we know what's important\? How can
+    What is the answer to life, the universe, and everything\? The answer to the questions that are asked of us
+    ====
+    (?s).*?
+    What is the answer to life, the universe, and everything\? How is life possible, what is the meaning of
     ```
      -->
     <!-- Tripy: TEST: EXPECTED_STDOUT End -->

tripy/examples/nanogpt/example.py

@@ -82,7 +82,7 @@ def main():
     input_ids = encoder.encode(args.input_text, allowed_special={"<|endoftext|>"})
 
     TEMPERATURE = 0.8
-    TOP_K = 10
+    TOP_K = 5
 
     padded_seq_len = len(input_ids) + args.max_new_tokens
 

tripy/tests/test_examples.py

@@ -103,7 +103,12 @@ def test_examples(example, sandboxed_install_run):
             if block.has_marker("test: expected_stdout"):
                 print("Checking command output against expected output: ", end="")
                 out = statuses[-1].stdout.strip()
-                matched = re.match(dedent(block_text).strip(), out)
+                matched = False
+                expected_outs = dedent(block_text).split("====")
+                for expected in expected_outs:
+                    if re.match(expected.strip(), out):
+                        matched = True
+                        break
                 print("matched!" if matched else "did not match!")
                 print(f"==== STDOUT ====\n{out}")
                 assert matched