app.py

import gradio as gr
import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.calibration import CalibratedClassifierCV

# 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]]

    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=(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.'
)

demo.launch()