Spaces:
Runtime error
Runtime error
File size: 3,158 Bytes
5b111f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import pandas as pd
import numpy as np
import joblib
import os
from huggingface_hub import hf_hub_download, login
class EngineConditionModel:
def __init__(self, model_name):
self.model_name = model_name
self.model = None
self.scaler = None
self.features = [
'Engine rpm', 'Lub oil pressure', 'Fuel pressure',
'Coolant pressure', 'lub oil temp', 'Coolant temp'
]
self.load_model()
def load_model(self):
try:
hf_token = os.getenv("HF_TOKEN")
if hf_token:
try:
login(token=hf_token)
except Exception as e:
print(f"HF login warning: {e}")
model_path = hf_hub_download(
repo_id=self.model_name,
filename="model/best_engine_model.joblib", # adjust if different
repo_type="model",
token=hf_token
)
self.model = joblib.load(model_path)
print("β
Loaded model artifact.")
# Optional scaler
try:
scaler_path = hf_hub_download(
repo_id=self.model_name,
filename="scaler.joblib",
repo_type="model",
token=hf_token
)
self.scaler = joblib.load(scaler_path)
print("β
Loaded scaler.")
except Exception:
self.scaler = None
except Exception as e:
print(f"β HF Hub load error: {e}
β οΈ Falling back to dummy model.")
from sklearn.ensemble import RandomForestClassifier
self.model = RandomForestClassifier(n_estimators=10, random_state=42)
X_dummy = np.random.rand(10, len(self.features))
y_dummy = np.random.randint(0, 2, 10)
self.model.fit(X_dummy, y_dummy)
self.scaler = None
def preprocess(self, data):
df = pd.DataFrame([data]) if isinstance(data, dict) else data.copy()
X = df[self.features]
if self.scaler is not None:
return self.scaler.transform(X)
return X.values
def predict(self, data):
try:
X = self.preprocess(data)
pred = int(self.model.predict(X)[0])
if hasattr(self.model, "predict_proba"):
proba = self.model.predict_proba(X)
conf = float(proba[0, 1]) if proba.shape[1] > 1 else float(proba[0, 0])
else:
conf = 0.5
condition = 'Maintenance Required' if pred == 1 else 'Normal'
return {"prediction": pred, "probability": conf, "condition": condition}
except Exception as e:
print(f"β Prediction error: {e}")
return {"prediction": -1, "probability": 0.0, "condition": "Error"}
_engine_model_instance = None
def load_engine_model(model_name="dhani10/engine-condition-model"):
global _engine_model_instance
if _engine_model_instance is None:
_engine_model_instance = EngineConditionModel(model_name)
return _engine_model_instance
|