From b03e00c08f792559a84f3de2845eff6f50655d2f Mon Sep 17 00:00:00 2001
From: Richard Gong <richard@modal.com>
Date: Thu, 15 Feb 2024 19:12:56 +0000
Subject: [PATCH] Loosen Mixtral loss

---
 ci/check_inference.py | 2 ++
 ci/check_loss.py      | 5 ++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/ci/check_inference.py b/ci/check_inference.py
index ec25864..b76058d 100644
--- a/ci/check_inference.py
+++ b/ci/check_inference.py
@@ -12,4 +12,6 @@
 
     output = subprocess.check_output(["modal", "run", "src.inference", "--run-folder", f"/runs/{run_name}", "--prompt", prompt])
 
+    print("Asserting that output contains [SQ] SELECT ... [/SQL]")
+
     assert b"[SQL] SELECT" in output and b"[/SQL]" in output
diff --git a/ci/check_loss.py b/ci/check_loss.py
index aab58ac..90cca2e 100644
--- a/ci/check_loss.py
+++ b/ci/check_loss.py
@@ -27,5 +27,8 @@
     train_loss = float(results["TrainingLoss"].iloc[-1])
     val_loss = float(results["ValidationLoss"].iloc[-1])
 
+    # Arbitrary threshold
+    max_loss = 0.8 if b"Mixtral" in contents else 0.25
+
     print(f"Loss: {train_loss:.2f} (training), {val_loss:.2f} (validation)")
-    sys.exit(val_loss > 0.4)  # Arbitrary threshold
+    sys.exit(val_loss > max_loss)