Update app.py

app.py

@@ -9,7 +9,6 @@ import joblib
 import os
 
 # Step 1: Load Hugging Face model for anomaly detection
-# Using the "huggingface-course/distilbert-base-uncased-finetuned-imdb" model
 tokenizer = AutoTokenizer.from_pretrained("huggingface-course/distilbert-base-uncased-finetuned-imdb")
 model = AutoModelForSequenceClassification.from_pretrained("huggingface-course/distilbert-base-uncased-finetuned-imdb")
 anomaly_detection = pipeline("text-classification", model=model, tokenizer=tokenizer)
@@ -55,6 +54,12 @@ def detect_anomaly(logs):
 
 # Step 5: Function to predict failures based on device metrics
 def predict_failure(device_metrics):
+    # Check if metrics are None or missing required fields
+    if device_metrics is None:
+        return "Device metrics are missing."
+    if 'cpu_usage' not in device_metrics or 'memory_usage' not in device_metrics or 'error_rate' not in device_metrics:
+        return "Invalid metrics format. Please provide 'cpu_usage', 'memory_usage', and 'error_rate'."
+
     # Convert device metrics into a numpy array for prediction
     metrics_array = np.array([device_metrics['cpu_usage'], device_metrics['memory_usage'], device_metrics['error_rate']]).reshape(1, -1)
     failure_prediction = failure_prediction_model.predict(metrics_array)  # Use the Random Forest model for failure prediction