heatmap_advance / modules /ai_engine.py
Sanjayraju30's picture
Update modules/ai_engine.py
115c399 verified
import pandas as pd
import torch
from transformers import TimeSeriesTransformerModel, TimeSeriesTransformerConfig
class AIEngine:
def __init__(self):
config = TimeSeriesTransformerConfig(
input_size=4,
output_size=1,
context_length=10,
prediction_length=1
)
self.model = TimeSeriesTransformerModel(config)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model.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_tensor = self.preprocess(df)
with torch.no_grad():
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