Update model_loader.py

model_loader.py

@@ -1,38 +1,32 @@
-import os
+def load_model():
+    try:
+        print(f"🔄 Loading tokenizer and model: {MODEL_ID}")
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID,
+            token=HF_TOKEN,
+            trust_remote_code=True
+        )
 
-# Hugging Face Token
-HF_TOKEN = os.getenv("HF_TOKEN")
-if not HF_TOKEN:
-    raise ValueError("❌ HF_TOKEN not set in environment variables.")
+        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
+        )
 
-# Model to Load
-MODEL_ID = os.getenv("MODEL_ID", "meta-llama/Meta-Llama-3-8B-Instruct")  # 👈 updated here
+        print("✅ Model loaded successfully.")
+        return pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=150,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
+        )
 
-# 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
-
-# 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"
-}
+    except Exception as e:
+        print(f"❌ Failed to load model: {e}")
+        raise RuntimeError(f"Model loading failed: {e}")