Update app.py

app.py

@@ -3,41 +3,52 @@ import joblib
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.calibration import CalibratedClassifierCV
 
-# Attempt to load the model and vectorizer
+# Load model and vectorizer
 try:
     model = joblib.load("helix-psa.pkl")
-    print("Model loaded successfully:", type(model))
-except Exception as e:
-    print("Error loading model:", e)
-
-try:
     vectorizer = joblib.load("helix-psa-vectorizer.pkl")
-    print("Vectorizer loaded successfully:", type(vectorizer))
+    print("Model and vectorizer loaded successfully.")
 except Exception as e:
-    print("Error loading vectorizer:", e)
+    print("Error loading model or vectorizer:", e)
+    model, vectorizer = None, None  # Assign None so we can check later
 
 def predictPasswordStrength(password_input):
-    password_tfidf = vectorizer.transform([password_input])
-    
-    # Make predictions
-    predicted_proba = model.predict_proba(password_tfidf)
-    predicted_class = int(model.predict(password_tfidf)[0])  # Convert to Python integer
-    output = ''
-    if predicted_class == 0:
-        output = "The password is very weak..."
-    elif predicted_class == 1:
-        output = "The password is average."
-    else:
-        output = "The password is strong. But alas, it is not unbreakable."
-
-    confidence = float(predicted_proba.max())
-
-    return password_input, output, confidence
+    if model is None or vectorizer is None:
+        return [["Error: Model or vectorizer not loaded correctly.", "", ""]]
+
+    try:
+        password_tfidf = vectorizer.transform([password_input])
+
+        # Make predictions
+        predicted_proba = model.predict_proba(password_tfidf)
+        predicted_class = int(model.predict(password_tfidf)[0])  # Convert to Python integer
+        output = ''
+        if predicted_class == 0:
+            output = "The password is very weak..."
+        elif predicted_class == 1:
+            output = "The password is average."
+        else:
+            output = "The password is strong. But alas, it is not unbreakable."
+
+        confidence = float(predicted_proba.max())
+
+        # Return as a list of lists for DataFrame
+        return [[password_input, output, confidence]]
+
+    except Exception as e:
+        return [[f"Error during prediction: {e}", "", ""]]
 
 demo = gr.Interface(
     fn=predictPasswordStrength,
     inputs=[gr.Textbox('Hello123', label='Password', info='The password to check the strength of', max_lines=1)],
-    outputs=[gr.Dataframe(row_count = (2, "dynamic"), col_count=(3, "fixed"), label="Generated Names", headers=["Password", "Prediction", "Confidence"])],
+    outputs=[
+        gr.Dataframe(
+            row_count=(1, "fixed"),
+            col_count=(3, "fixed"),
+            headers=["Password", "Prediction", "Confidence"],
+            label="Password Strength Analysis"
+        )
+    ],
     title='Helix - Password Strength Analyzer',
     description='A password strength analyzer, trained on over 10 million different passwords.'
 )