compare result among 2 results

Author: yizhuoz004

tripy/examples/nanogpt/README.md

@@ -39,6 +39,9 @@ for expected accuracy.
     ```
     (?s).*?
     What is the answer to life, the universe, and everything\? How can we know what's important\? How can
+    ====
+    (?s).*?
+    What is the answer to life, the universe, and everything\? How can we get back at these questions\? And
     ```
      -->
     <!-- Tripy: TEST: EXPECTED_STDOUT End -->
@@ -63,6 +66,9 @@ 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
+    ====
+    (?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 = 200
 
     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