From 130afbdd9247a3a4fe79d34e1c508f16a20f2fad Mon Sep 17 00:00:00 2001
From: Niels Geens <niels.geens@gmail.com>
Date: Tue, 1 Apr 2025 11:49:21 +0000
Subject: [PATCH] Updates after testing

---
 scripts/install.sh   | 28 ++++++++++++++++++++++------
 setup.py             |  3 ++-
 sia/config.py        |  6 +++---
 tools/train/setup.py |  2 +-
 4 files changed, 28 insertions(+), 11 deletions(-)

diff --git a/scripts/install.sh b/scripts/install.sh
index 7b73f4b..dc25a33 100755
--- a/scripts/install.sh
+++ b/scripts/install.sh
@@ -1,20 +1,36 @@
 #!/bin/bash
 
+add_to_environment() {
+    local path="$1"
+    # Check if the path is already in the environment file
+    if ! grep -q "$path" /etc/environment; then
+        # Get current PATH from environment file
+        current_path=$(grep "^PATH=" /etc/environment | cut -d'"' -f2)
+        
+        if [ -z "$current_path" ]; then
+            # If PATH doesn't exist in the file, create it
+            echo "PATH=\"\$PATH:$path\"" >> /etc/environment
+        else
+            # Replace existing PATH line with updated one
+            sed -i "s|^PATH=.*|PATH=\"$path:$current_path\"|" /etc/environment
+        fi
+    fi
+}
+
 echo "Installing ITB tool..."
 python3 -m venv "/root/venvs/itb"
 /root/venvs/itb/bin/pip install -e /root/sia/tools/itb/
-echo "PATH=\"/root/venvs/itb/bin/:\$PATH\"" > "/etc/profile.d/venv_itb.sh"
-source "/etc/profile.d/venv_itb.sh"
+add_to_environment "/root/venvs/itb/bin"
 
 echo "Installing Train tool..."
 python3 -m venv "/root/venvs/train"
 /root/venvs/train/bin/pip install -e /root/sia/tools/train/
 /root/venvs/train/bin/ipython kernel install --name=train
-echo "PATH=\"/root/venvs/train/bin/:\$PATH\"" > "/etc/profile.d/venv_train.sh"
-source "/etc/profile.d/venv_train.sh"
+add_to_environment "/root/venvs/train/bin"
 
 echo "Installing SIA core..."
 python3 -m venv "/root/venvs/sia"
 /root/venvs/sia/bin/pip install -e /root/sia/
-echo "PATH=\"/root/venvs/sia/bin/:\$PATH\"" > "/etc/profile.d/venv_sia.sh"
-source "/etc/profile.d/venv_sia.sh"
\ No newline at end of file
+add_to_environment "/root/venvs/sia/bin"
+
+source "/etc/environment
\ No newline at end of file
diff --git a/setup.py b/setup.py
index d19fb95..0efb7d7 100644
--- a/setup.py
+++ b/setup.py
@@ -9,10 +9,11 @@ setup(
         'aiohttp>=3.8.0',
         'bitsandbytes>=0.41.0',
         'dotenv-python>=0.0.1',
+        'flash-attn>=2.0.0',
         'huggingface_hub>=0.16.0',
         'lxml>=4.9.0',
-        'mistralai>=0.0.7',
         'mistral-common>=1.0.0',
+        'mistralai>=0.0.7',
         'openai>=1.0.0',
         'psutil>=5.9.0',
         'python-dotenv>=1.0.0',
diff --git a/sia/config.py b/sia/config.py
index 1409b04..4c96187 100644
--- a/sia/config.py
+++ b/sia/config.py
@@ -188,7 +188,7 @@ class Config:
         parser.add_argument(
             '--qwq-enable',
             action='store_true',
-            default=self._parse_bool_env('SIA_QWQ_ENABLED', True),  # Enable by default
+            default=self._parse_bool_env('SIA_QWQ_ENABLED', True),
             help='Enable QwQ LLM engine (env: SIA_QWQ_ENABLED)'
         )
         parser.add_argument(
@@ -200,13 +200,13 @@ class Config:
         parser.add_argument(
             '--qwq-temperature',
             type=float,
-            default=float(os.getenv('SIA_QWQ_TEMPERATURE', '0.1')),
+            default=float(os.getenv('SIA_QWQ_TEMPERATURE', '0.6')),
             help='QwQ temperature (default: 0.1, env: SIA_QWQ_TEMPERATURE)'
         )
         parser.add_argument(
             '--qwq-token-limit',
             type=int,
-            default=int(os.getenv('SIA_QWQ_TOKEN_LIMIT', '0')),
+            default=int(os.getenv('SIA_QWQ_TOKEN_LIMIT', '4096')),
             help='QwQ token limit (0 for model default, env: SIA_QWQ_TOKEN_LIMIT)'
         )
 
diff --git a/tools/train/setup.py b/tools/train/setup.py
index f633cf9..e79b8ee 100644
--- a/tools/train/setup.py
+++ b/tools/train/setup.py
@@ -10,7 +10,7 @@ setup(
 
     install_requires=[
         'accelerate>=0.25.0',
-        'bitsandbytes>=0.41.1',
+        'bitsandbytes>=0.45.0',
         'black>=22.0.0',
         'datasets>=2.14.6',
         'einops>=0.7.0',