Update model_loader.py

model_loader.py

@@ -1,47 +1,38 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from config import HF_TOKEN, MODEL_ID
+import os
 
-def load_model():
-    try:
-        print(f"🔄 Loading tokenizer and model: {MODEL_ID}")
+# Hugging Face Token
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    raise ValueError("❌ HF_TOKEN not set in environment variables.")
 
-        # Load tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(
-            MODEL_ID,
-            token=HF_TOKEN,
-            trust_remote_code=True,
-            padding_side="left"  # For chat-style models
-        )
+# Model to Load
+MODEL_ID = os.getenv("MODEL_ID", "meta-llama/Meta-Llama-3-8B-Instruct")  # 👈 updated here
 
-        # Set max length (MedGemma supports up to 8192 tokens)
-        tokenizer.model_max_length = 8192
+# Paths and Settings
+TEMP_DIR = os.getenv("TEMP_DIR", "/tmp")  # For screenshots, temporary files, etc.
+ENABLE_EXECUTION = os.getenv("ENABLE_EXECUTION", "true").lower() == "true"  # toggle command execution
 
-        # Load model
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            token=HF_TOKEN,
-            trust_remote_code=True,
-            device_map="auto" if torch.cuda.is_available() else "cpu",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            low_cpu_mem_usage=True
-        )
-
-        print("✅ Model loaded successfully.")
-
-        # Return generation pipeline with large max token output and context window
-        return pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_new_tokens=1024,         # Max response length
-            truncation=True,             # Safely truncate input if needed
-            pad_token_id=tokenizer.eos_token_id,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
-
-    except Exception as e:
-        print(f"❌ Failed to load model: {e}")
-        raise RuntimeError(f"Model loading failed: {e}")
+# OS Commands (unchanged)
+OS_COMMANDS = {
+    "open": {
+        "chrome": "google-chrome",
+        "firefox": "firefox",
+        "notepad": "notepad",
+        "calculator": "gnome-calculator",
+        "terminal": "gnome-terminal",
+        "vscode": "code"
+    },
+    "close": {
+        "chrome": "pkill chrome",
+        "firefox": "pkill firefox",
+        "notepad": "pkill notepad",
+        "calculator": "pkill gnome-calculator",
+        "vscode": "pkill code"
+    },
+    "screenshot": {
+        "default": "import -window root screenshot.png"
+    },
+    "shutdown": "shutdown now",
+    "restart": "reboot",
+    "lock": "gnome-screensaver-command -l"
+}