Update modules/ai_engine.py

modules/ai_engine.py

@@ -5,8 +5,8 @@ from transformers import TimeSeriesTransformerModel, TimeSeriesTransformerConfig
 class AIEngine:
     def __init__(self):
         config = TimeSeriesTransformerConfig(
-            input_size=4,  # SolarGen, WindGen, Tilt, Vibration
-            output_size=1,  # Health Score
+            input_size=4,
+            output_size=1,
             context_length=10,
             prediction_length=1
         )
@@ -14,14 +14,14 @@ class AIEngine:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
 
-    def preprocess_data(self, df: pd.DataFrame) -> torch.Tensor:
-        features = df[["SolarGen(kWh)", "WindGen(kWh)", "Tilt(°)", "Vibration(g)"]].values
-        return torch.tensor(features, dtype=torch.float32).unsqueeze(0).to(self.device)
+    def preprocess(self, df: pd.DataFrame) -> torch.Tensor:
+        data = df[["SolarGen(kWh)", "WindGen(kWh)", "Tilt(°)", "Vibration(g)"]].values
+        return torch.tensor(data, dtype=torch.float32).unsqueeze(0).to(self.device)
 
     def predict_health(self, df: pd.DataFrame) -> pd.DataFrame:
-        input_data = self.preprocess_data(df)
+        input_tensor = self.preprocess(df)
         with torch.no_grad():
-            predictions = self.model(input_data).logits.squeeze().cpu().numpy()
+            predictions = self.model(input_tensor).logits.squeeze().cpu().numpy()
         df["HealthScore"] = predictions
         df["ML_Anomaly"] = df["HealthScore"].apply(lambda x: "Risk" if x < 0.5 else "Normal")
         return df