Update handler.py

handler.py

@@ -8,19 +8,27 @@ class CustomModelHandler:
         self.load_model()
 
     def load_model(self):
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, trust_remote_code=True)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_name_or_path, 
-            trust_remote_code=True,
-            torch_dtype="auto"
-        )
-        self.model.to(self.device)
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, trust_remote_code=True)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name_or_path, 
+                trust_remote_code=True,
+                torch_dtype="auto"
+            )
+            self.model.to(self.device)
+        except Exception as e:
+            print(f"An error occurred while loading the model: {e}")
+            raise
 
     def predict(self, inputs):
-        inputs = self.tokenizer(inputs, return_tensors="pt").to(self.device)
-        outputs = self.model.generate(**inputs)
-        predictions = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-        return predictions
+        try:
+            inputs = self.tokenizer(inputs, return_tensors="pt").to(self.device)
+            outputs = self.model.generate(**inputs)
+            predictions = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+            return predictions
+        except Exception as e:
+            print(f"An error occurred during prediction: {e}")
+            raise
 
 # Initialize the handler with the model path
 handler = CustomModelHandler("microsoft/Phi-3-vision-128k-instruct")