Update app.py

app.py

@@ -1,56 +1,69 @@
 import gradio as gr
-import joblib
+import pickle
+import numpy as np'
+from scipy.sparse import hstack
 from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.calibration import CalibratedClassifierCV
+from sklearn.linear_model import LogisticRegression
 
 # Load model and vectorizer
 try:
-    model = joblib.load("helix-psa.pkl")
-    vectorizer = joblib.load("helix-psa-vectorizer.pkl")
-    print("Model and vectorizer loaded successfully.")
-except Exception as e:
-    print("Error loading model or vectorizer:", e)
-    model, vectorizer = None, None  # Assign None so we can check later
-
-def predictPasswordStrength(password_input):
-    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]]
+    # --- Load and inference code ---
+    with open('password_model.pkl', 'rb') as f:
+      model = pickle.load(f)
+
+    with open('password_vectorizer.pkl', 'rb') as f:
+      vectorizer = pickle.load(f)
+
+def extract_features(password):
+    features = {}
+    features['length'] = len(password)
+    features['uppercase'] = sum(1 for c in password if c.isupper())
+    features['lowercase'] = sum(1 for c in password if c.islower())
+    features['digits'] = sum(1 for c in password if c.isdigit())
+    features['special'] = sum(1 for c in password if not c.isalnum())
+    return features
+
+def predict_password_strength(password, vectorizer, model): # Add vectorizer and model as arguments
+    # Extract features from the input password
+    features = extract_features(password)
+
+    # Transform the input password using the trained vectorizer
+    password_vectorized = vectorizer.transform([password])
+    password_vectorized = hstack((password_vectorized, np.array(list(features.values())).reshape(1, -1)))
+
+    text="No password analyzed."
+
+    # Make a prediction using the trained model
+    prediction = model.predict(password_vectorized)[0]
+    if prediction === 0:
+        text = "Password is very weak."
+    elif prediction === 1:
+        text = "Password is weak."
+    elif prediction === 2:
+        text = "Password is average."
+    elif prediction === 3:
+        text = "Password is strong."
+    elif prediction === 4:
+        text = "Password is very strong."
+
+    return password, prediction, text
 
     except Exception as e:
         return [[f"Error during prediction: {e}", "", ""]]
 
 demo = gr.Interface(
-    fn=predictPasswordStrength,
+    fn=predict_password_strength,
     inputs=[gr.Textbox('Hello123', label='Password', info='The password to check the strength of', max_lines=1)],
     outputs=[
         gr.Dataframe(
             row_count=(1, "fixed"),
             col_count=(3, "fixed"),
-            headers=["Password", "Prediction", "Confidence"],
+            headers=["Password", "Prediction", "Strength_Text"],
             label="Password Strength Analysis"
         )
     ],
     title='Helix - Password Strength Analyzer',
-    description='A password strength analyzer, trained on over 10 million different passwords.'
+    description='A password strength analyzer, trained on over 5 million different passwords.'
 )
 
 demo.launch()