Fix device mapping and pipeline creation for production use

production_inference.py

@@ -82,20 +82,36 @@ class ProthomAloModel:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
             
             # Load model
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
-                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-                device_map=device,
-                trust_remote_code=True
-            )
-            
-            # Create pipeline for easier use
-            self.pipeline = pipeline(
-                "text-generation",
-                model=self.model,
-                tokenizer=self.tokenizer,
-                device=0 if device == "cuda" else -1
-            )
+            if device == "auto":
+                # Use device_map="auto" for automatic device placement
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_name,
+                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+                    device_map="auto",
+                    trust_remote_code=True
+                )
+                # Create pipeline without device specification when using device_map
+                self.pipeline = pipeline(
+                    "text-generation",
+                    model=self.model,
+                    tokenizer=self.tokenizer
+                )
+            else:
+                # Specific device handling
+                device_obj = torch.device("cuda" if device == "cuda" and torch.cuda.is_available() else "cpu")
+                self.model = AutoModelForCausalLM.from_pretrained(
+                    self.model_name,
+                    torch_dtype=torch.float16 if device_obj.type == "cuda" else torch.float32,
+                    trust_remote_code=True
+                ).to(device_obj)
+                
+                # Create pipeline with device specification
+                self.pipeline = pipeline(
+                    "text-generation",
+                    model=self.model,
+                    tokenizer=self.tokenizer,
+                    device=0 if device_obj.type == "cuda" else -1
+                )
             
             logger.info("Model loaded successfully with production optimizations")
             return True