bunch of stuff

1 files changed, 12 insertions, 3 deletions

diff --git a/colab_setup.py b/colab_setup.py
index 42d57d9..faf8f79 100644
--- a/colab_setup.py
+++ b/colab_setup.py
@@ -2,7 +2,7 @@
 """
 run this in a colab cell to stand up main.py + a public wss:// tunnel to it.
 session is ephemeral — restart the notebook, get a new url, re-paste it
-into your bot.py-side config/connection.txt.
+into your bot.py-side files/config/connection.txt (or rerun setup_wizard.py).
 """
 import os
 import re
@@ -14,8 +14,13 @@ from pathlib import Path
 SCRIPT_DIR = Path(__file__).resolve().parent
 CONFIG_DIR = SCRIPT_DIR / "config"
 
-CHATMODEL = "hf.co/Stetsoncody/Llama-3.2-3B-Instruct-abliterated-GGUF:Q4_K_M"
-EMBEDMODEL = "hf.co/CompendiumLabs/bge-base-en-v1.5-gguf"
+def model_from_main(name):
+    # read the model names straight out of main.py so they can't drift apart
+    text = (SCRIPT_DIR / "main.py").read_text(encoding="utf-8")
+    return re.search(rf'^{name}\s*=\s*["\'](.+?)["\']', text, re.M).group(1)
+
+CHATMODEL = model_from_main("chatmodel")
+EMBEDMODEL = model_from_main("embedmodel")
 
 
 def run(cmd, **kwargs):
@@ -65,9 +70,13 @@ def ensure_params():
 
 def start_main():
     print("\n-- starting main.py --")
+    # same deal as the tunnel below, main.py prints a lot and raises BrokenPipe once this cell's stdout is gone
+    log_file = open(SCRIPT_DIR / "main.log", "w")
     return subprocess.Popen(
         f"python3 {SCRIPT_DIR / 'main.py'}",
         shell=True,
+        stdout=log_file,
+        stderr=subprocess.STDOUT,
     )