Update app.py

app.py

@@ -60,43 +60,34 @@ def test_llm_generation():
 
 
 def initialize_model():
-    """Separate model initialization with better error handling"""
     global model, tokenizer, generation_config
-    
+
     try:
         add_log("🔄 Initializing model...")
-        
+
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_ID,
             trust_remote_code=True,
-            use_fast=False  # Sometimes fast tokenizers cause issues
+            use_fast=False
         )
-        
-        # Ensure proper padding token
+
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
             add_log("✅ Set pad_token to eos_token")
-        
-        # Load model with proper device management
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-        
+
+        # Force GPU settings
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
-            torch_dtype=torch_dtype,
-            device_map="auto" if torch.cuda.is_available() else None,
+            torch_dtype=torch.float16,
             trust_remote_code=True,
+            device_map={"": 0},  # <- force GPU:0
             low_cpu_mem_usage=True
         )
-        
-        if not torch.cuda.is_available():
-            model = model.to(device)
-        
+
         model.eval()
-        
-        # Configure generation parameters
+
         generation_config = GenerationConfig(
-            max_new_tokens=4095,  # Reduced for stability
+            max_new_tokens=4096,
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
@@ -105,9 +96,16 @@ def initialize_model():
             repetition_penalty=1.1,
             length_penalty=1.0
         )
-        
+
         add_log(f"✅ Model loaded successfully on device: {model.device}")
         return True
+
+    except Exception as e:
+        error_msg = f"❌ Model initialization failed: {str(e)}"
+        add_log(error_msg)
+        add_log(f"Traceback: {traceback.format_exc()}")
+        return False
+
         
     except Exception as e:
         error_msg = f"❌ Model initialization failed: {str(e)}"