app.py

import gradio as gr
import joblib
import pandas as pd
import datasets
import json
import numpy as np

# Load the model
pipe = joblib.load("./model.pkl")

title = "Premium Amount Prediction"
description = "This model predicts the Premium Amount. Drag and drop any slice from the dataset or edit values as you wish in the dataframe component below."

# Load and prepare dataset
df = datasets.load_dataset("silvaKenpachi/mental_health")["train"].to_pandas()
df.dropna(axis=0, inplace=True)

# Load configuration
with open("./config.json") as f:
    config_dict = json.load(f)
all_headers = config_dict["sklearn"]["columns"]

# Filter headers to only include those present in the dataset
headers = [col for col in all_headers if col in df.columns]

# Define input and output interfaces
#inputs = [gr.Dataframe(headers=headers, row_count=(2, "dynamic"), col_count=(len(headers), "fixed"), label="Input Data", interactive=True)]
inputs = [gr.Dataframe(headers=all_headers, row_count=("dynamic"), col_count=(len(all_headers), "fixed"), label="Input Data", interactive=True)]
outputs = [gr.Dataframe(row_count=("dynamic"), col_count=(1, "fixed"), label="Predictions", headers=["Depression"])]

#def infer(inputs):
    #data = pd.DataFrame(inputs, columns=headers)
    #predictions = pipe.predict(data)
    #return pd.DataFrame(predictions, columns=["Depression"])

#code to fix missing columns with na
#def infer(inputs):
    #data = pd.DataFrame(inputs, columns=headers)
    # Add missing columns with default values (e.g., 0)
    #for col in all_headers:
        #if col not in data.columns:
            #data[col] = 0
    # Ensure the order of columns matches the training data
    #data = data[all_headers]
    #predictions = pipe.predict(data)
    #return pd.DataFrame(predictions, columns=["Depression"])


def infer(inputs):
    data = pd.DataFrame(inputs, columns=headers)
    
    # Replace empty strings with NaN
    data = data.replace('', np.nan)
    
    # Add missing columns with default values (e.g., 0)
    for col in all_headers:
        if col not in data.columns:
            data[col] = 0
    
    # Ensure the order of columns matches the training data
    data = data[all_headers]
    
    # Fill NaN values with default values (e.g., 0)
    data = data.fillna(0)
    
    # Convert all data to float
    data = data.astype(float)
    
    predictions = pipe.predict(data)
    return pd.DataFrame(predictions, columns=["Depression"])



gr.Interface(
    fn=infer,
    inputs=inputs,
    outputs=outputs,
    title=title,
    description=description,
    examples=[df[headers].head(3).values.tolist()],
    cache_examples=False
).launch(debug=True)