Update model_loader.py

model_loader.py

@@ -5,10 +5,10 @@ from config import HF_TOKEN, MODEL_ID
 def load_model():
     try:
         print(f"🔄 Loading tokenizer and model: {MODEL_ID}")
-
+        
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_ID,
-            token=HF_TOKEN or None,  # Handles public and private
+            token=HF_TOKEN or None,
             trust_remote_code=True
         )
 
@@ -16,19 +16,21 @@ def load_model():
             MODEL_ID,
             token=HF_TOKEN or None,
             trust_remote_code=True,
-            device_map="auto" if torch.cuda.is_available() else None,
+            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
         )
 
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"✅ Model loaded successfully.\nDevice set to: {device}")
+        print("✅ Model loaded successfully.")
 
         return pipeline(
             "text-generation",
             model=model,
             tokenizer=tokenizer,
-            device=0 if device == "cuda" else -1
+            max_new_tokens=2048,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
         )
 
     except Exception as e: