upgini 1.2.97__tar.gz → 1.2.98a3922.dev2__tar.gz

{upgini-1.2.97 → upgini-1.2.98a3922.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.2.97
+Version: 1.2.98a3922.dev2
 Summary: Intelligent data search & enrichment for Machine Learning
 Project-URL: Bug Reports, https://github.com/upgini/upgini/issues
 Project-URL: Homepage, https://upgini.com/

upgini-1.2.98a3922.dev2/src/upgini/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.2.98a3922.dev2"

{upgini-1.2.97 → upgini-1.2.98a3922.dev2}/src/upgini/autofe/vector.py

@@ -38,12 +38,32 @@ class OnnxModel(PandasOperator, metaclass=OperatorRegistry):
     is_vector: bool = True
     output_type: Optional[str] = "float"
     model_name: str = ""
+    model_version: int = 1
 
     def get_params(self) -> Dict[str, Optional[str]]:
         res = super().get_params()
         res.update(
             {
                 "model_name": self.model_name,
+                "model_version": self.model_version,
+            }
+        )
+        return res
+
+
+class CatboostModel(PandasOperator, metaclass=OperatorRegistry):
+    name: str = "catboost"
+    is_vector: bool = True
+    output_type: Optional[str] = "float"
+    model_name: str = ""
+    model_version: int = 1
+
+    def get_params(self) -> Dict[str, Optional[str]]:
+        res = super().get_params()
+        res.update(
+            {
+                "model_name": self.model_name,
+                "model_version": self.model_version,
             }
         )
         return res

{upgini-1.2.97 → upgini-1.2.98a3922.dev2}/src/upgini/data_source/data_source_publisher.py

@@ -495,15 +495,24 @@ class DataSourcePublisher:
                 self.logger.exception("Failed to reannounce all ADS-es")
 
     def upload_autofe_model(
-        self, file_path: str, name: str, model_type: Optional[Literal["ONNX"]] = None, description: str = ""
+        self,
+        file_path: str,
+        name: str,
+        input_names: List[str],
+        search_id: str,
+        model_type: Optional[Literal["ONNX", "CATBOOST"]] = None,
+        description: str = "",
     ):
-        if model_type is not None and model_type not in ["ONNX"]:
+        if model_type is not None and model_type not in ["ONNX", "CATBOOST"]:
             raise ValueError(f"Invalid model type: {model_type}. Available values: ONNX")
         metadata = {
             "modelName": name,
+            "inputNames": input_names,
+            "searchTaskId": search_id,
             "modelType": model_type or "ONNX",
             "description": description,
         }
+
         trace_id = str(uuid.uuid4())
         with MDC(trace_id=trace_id):
             try:

upgini-1.2.97/src/upgini/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.2.97"