Update: Add custom NPKPredictionModel implementation

config.json

@@ -2,7 +2,8 @@
   "model_type": "npk",
   "architectures": ["NPKPredictionModel"],
   "auto_map": {
-  "AutoConfig": "modeling_npk.NPKConfig",
-  "AutoModel": "modeling_npk.NPKPredictionModel"
-  }
-  }
+    "AutoConfig": "modeling_npk.NPKConfig",
+    "AutoModel": "modeling_npk.NPKPredictionModel"
+  },
+  "trust_remote_code": true
+}

modeling_npk.py

@@ -1,4 +1,3 @@
-
 import pickle
 import pandas as pd
 from transformers import PreTrainedModel, PretrainedConfig
@@ -17,6 +16,18 @@ class NPKPredictionModel(PreTrainedModel):
         super().__init__(config)
         self.xgb_model = None
         self.label_encoder = None
+        self._load_models()
+    
+    def _load_models(self):
+        # Load the XGBoost model and label encoder
+        xgb_path = hf_hub_download(repo_id=self.config._name_or_path, filename="npk_prediction_model.pkl")
+        le_path = hf_hub_download(repo_id=self.config._name_or_path, filename="label_encoder.pkl")
+        
+        with open(xgb_path, 'rb') as f:
+            self.xgb_model = pickle.load(f)
+        
+        with open(le_path, 'rb') as f:
+            self.label_encoder = pickle.load(f)
     
     def forward(self, inputs):
         # Preprocess inputs
@@ -48,15 +59,4 @@ class NPKPredictionModel(PreTrainedModel):
             config = NPKConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
         
         model = cls(config)
-        
-        # Load the XGBoost model and label encoder
-        xgb_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="npk_prediction_model.pkl")
-        le_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="label_encoder.pkl")
-        
-        with open(xgb_path, 'rb') as f:
-            model.xgb_model = pickle.load(f)
-        
-        with open(le_path, 'rb') as f:
-            model.label_encoder = pickle.load(f)
-        
-        return model
+        return model