Update model_loader.py

model_loader.py

@@ -3,20 +3,34 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from config import HF_TOKEN, MODEL_ID
 
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        token=HF_TOKEN,
-        trust_remote_code=True,
-        device_map="cpu",
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True
-    )
-    return pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=100,
-        do_sample=True,
-        temperature=0.7,
-    )
+    try:
+        print(f"🔄 Loading tokenizer and model: {MODEL_ID}")
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID,
+            token=HF_TOKEN,
+            trust_remote_code=True
+        )
+
+        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 pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=150,
+            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}")