Upload folder using huggingface_hub

MLBaseModelDriver.py

@@ -1,8 +1,6 @@
 import torch
 import sys
 import pandas as pd
-import numpy as np
-from sklearn.preprocessing import StandardScaler, OneHotEncoder
 from typing import TypedDict, Optional, Tuple
 import datetime
 import math
@@ -10,57 +8,25 @@ import importlib.util
 from huggingface_hub import hf_hub_download
 import pickle
 
-# Класс предобработки
+# --------- Класс DataPreprocessor ---------
 class DataPreprocessor:
-    def __init__(self):
-        self.feature_scaler = StandardScaler()
-        self.target_scaler = StandardScaler()
-        self.encoder = OneHotEncoder(sparse=False, handle_unknown='ignore')
-
-    def fit_transform(self, df):
-        df['listing_date'] = pd.to_datetime(df['listing_date'])
-        df['sale_date'] = pd.to_datetime(df['sale_date'])
-        df['days_on_market'] = (df['sale_date'] - df['listing_date']).dt.days
-        df['age'] = df['listing_date'].dt.year - df['year_built']
-        df = df[df['days_on_market'] >= 0].dropna(subset=['days_on_market'])
-
-        df = df.fillna({
-            'beds': df['beds'].median(),
-            'baths': df['baths'].median(),
-            'sqft': df['sqft'].median(),
-            'year_built': df['year_built'].median(),
-            'listing_price': df['listing_price'].median(),
-            'age': df['age'].median()
-        })
-
-        df = df[(df['sale_price'] > 50000) & (df['sale_price'] < 2000000)]
-
-        cat_feature = self.encoder.fit_transform(df[['property_type']])
-        cat_df = pd.DataFrame(cat_feature, columns=self.encoder.get_feature_names_out(['property_type']))
-        df = df.reset_index(drop=True).join(cat_df)
-
-        for col in ['sale_price', 'listing_price', 'sqft']:
-            df[col] = np.log1p(df[col])
-
-        features = ['beds', 'baths', 'sqft', 'listing_price', 'days_on_market', 'age'] + list(cat_df.columns)
-        targets = ['sale_price']
-
-        X = df[features]
-        y = df[['sale_price']]
-
-        X_scaled = self.feature_scaler.fit_transform(X)
-        y_scaled = self.target_scaler.fit_transform(y)
-
-        self.features = features
-        return pd.DataFrame(X_scaled, columns=features), pd.DataFrame(y_scaled, columns=targets)
-
-    def inverse_transform_target(self, y_scaled):
-        return np.expm1(self.target_scaler.inverse_transform(y_scaled.reshape(-1, 1)).flatten())
+    def __init__(self, label_encoders, scaler):
+        self.label_encoders = label_encoders
+        self.scaler = scaler
 
+    def preprocess(self, df: pd.DataFrame) -> torch.Tensor:
+        default_beds = 3
+        default_sqft = 1500.0
+        default_property_type = '6'
+        df['beds'] = df['beds'].fillna(default_beds)
+        df['sqft'] = pd.to_numeric(df['sqft'], errors='coerce').fillna(default_sqft)
+        df['property_type'] = df['property_type'].fillna(default_property_type)
+        df['property_type'] = df['property_type'].astype(int)
+        df[['sqft', 'price']] = self.scaler.transform(df[['sqft', 'price']])
+        X = df[['beds', 'sqft', 'property_type', 'price']]
+        return torch.tensor(X.values, dtype=torch.float32)
 
-"""
-Data container class representing the data shape of the synapse coming into `run_inference`
-"""
+# --------- Класс ProcessedSynapse ---------
 class ProcessedSynapse(TypedDict):
     id: Optional[str]
     nextplace_id: Optional[str]
@@ -84,11 +50,12 @@ class ProcessedSynapse(TypedDict):
     hoa_dues: Optional[float]
     query_date: Optional[str]
 
-
+# --------- Класс MLBaseModelDriver ---------
 class MLBaseModelDriver:
 
     def __init__(self):
-        self.model, self.label_encoder, self.scaler = self.load_model()
+        self.model, self.label_encoders, self.scaler = self.load_model()
+        self.preprocessor = DataPreprocessor(self.label_encoders, self.scaler)
 
     def load_model(self) -> Tuple[any, any, any]:
         print(f"Loading model...")
@@ -111,12 +78,10 @@ class MLBaseModelDriver:
 
     def _download_model_files(self) -> Tuple[str, str, str, str]:
         model_path = "ckoozzzu/NextPlace"
-
         model_file = hf_hub_download(repo_id=model_path, filename="model_files/real_estate_model.pth")
         scaler_file = hf_hub_download(repo_id=model_path, filename="model_files/scaler.pkl")
         label_encoders_file = hf_hub_download(repo_id=model_path, filename="model_files/label_encoder.pkl")
         model_class_file = hf_hub_download(repo_id=model_path, filename="MLBaseModel.py")
-
         return model_file, scaler_file, label_encoders_file, model_class_file
 
     def _import_model_class(self, model_class_file):
@@ -132,7 +97,8 @@ class MLBaseModelDriver:
             raise AttributeError(f"The module does not contain a class named 'MLBaseModel'")
 
     def run_inference(self, input_data: ProcessedSynapse) -> Tuple[float, str]:
-        input_tensor = self._preprocess_input(input_data)
+        df = pd.DataFrame([input_data])
+        input_tensor = self.preprocessor.preprocess(df)
 
         with torch.no_grad():
             prediction = self.model(input_tensor)
@@ -149,17 +115,3 @@ class MLBaseModelDriver:
             return sale_date
         else:
             return datetime.date.today() + datetime.timedelta(days=1)
-
-    def _preprocess_input(self, data: ProcessedSynapse) -> torch.tensor:
-        df = pd.DataFrame([data])
-        default_beds = 3
-        default_sqft = 1500.0
-        default_property_type = '6'
-        df['beds'] = df['beds'].fillna(default_beds)
-        df['sqft'] = pd.to_numeric(df['sqft'], errors='coerce').fillna(default_sqft)
-        df['property_type'] = df['property_type'].fillna(default_property_type)
-        df['property_type'] = df['property_type'].astype(int)
-        df[['sqft', 'price']] = self.scaler.transform(df[['sqft', 'price']])
-        X = df[['beds', 'sqft', 'property_type', 'price']]
-        input_tensor = torch.tensor(X.values, dtype=torch.float32)
-        return input_tensor