import gradio as gr
import pickle
import numpy as np

# Load all models and the vectorizer
with open("random_forest_model.pkl", "rb") as f:
    random_forest_model = pickle.load(f)

with open("logistic_model.pkl", "rb") as f:
    logistic_model = pickle.load(f)

with open("knn_yelp_model.pkl", "rb") as f:
    knn_yelp_model = pickle.load(f)

with open("svm_linear.pkl", "rb") as f:
    svm_linear = pickle.load(f)

with open("svm_poly.pkl", "rb") as f:
    svm_poly = pickle.load(f)

with open("svm_rbf.pkl", "rb") as f:
    svm_rbf = pickle.load(f)

with open("vectorizer.pkl", "rb") as f:
    vectorizer = pickle.load(f)

# Function to predict class probabilities
def predict_sentiment(review, model_name):
    # Vectorize the input review
    review_vec = vectorizer.transform([review])
    
    # Select the model based on user input
    if model_name == 'Random Forest':
        model = random_forest_model
    elif model_name == 'Logistic Regression':
        model = logistic_model
    elif model_name == 'KNN':
        model = knn_yelp_model
    elif model_name == 'SVM Linear':
        model = svm_linear
    elif model_name == 'SVM Poly':
        model = svm_poly
    elif model_name == 'SVM RBF':
        model = svm_rbf
    
    # Predict probabilities for each class (assuming binary classification)
    probabilities = model.predict_proba(review_vec)[0]
    return {"Negative": probabilities[0], "Positive": probabilities[1]}

# Set up the Gradio interface
iface = gr.Interface(
    fn=predict_sentiment, 
    inputs=[
        gr.Textbox(label="Enter Your Review"),
        gr.Dropdown(
            choices=["Random Forest", "Logistic Regression", "KNN", "SVM Linear", "SVM Poly", "SVM RBF"],
            label="Select Model"
        ),
    ], 
    outputs=gr.Label(num_top_classes=2), 
    live=True
)

# Launch the interface
iface.launch()