import data_prep
import model_predict
import gradio as gr

# Dictionary of model names and corresponding display names
model_dict = {
    "bert-base": "research-dump/bert-base-uncased_deletion_multiclass_complete_Final",
    "bert-large": "research-dump/bert-large-uncased_deletion_multiclass_complete_final",
    "roberta-base": "research-dump/roberta-base_deletion_multiclass_complete_final"
}

def process_url(url, model_key):
    # Get the actual model path from the model_dict
    model_name = model_dict[model_key]

    # Process the text from the URL
    processed_text = data_prep.process_data(url)
    
    # Predict the labels and their probabilities
    final_scores = model_predict.predict_text(processed_text, model_name)
    
    # Find the label with the highest probability
    highest_prob_label = max(final_scores, key=final_scores.get)
    highest_prob = final_scores[highest_prob_label]
    
    # Create progress bar style output for all labels
    progress_bars = {label: score for label, score in final_scores.items()}
    
    return highest_prob_label, highest_prob, progress_bars

# Define the interface for the Gradio app
url_input = gr.Textbox(label="URL")
model_name_input = gr.Dropdown(label="Model Name", choices=list(model_dict.keys()), value=list(model_dict.keys())[0])
outputs = [
    gr.Textbox(label="Label with Highest Probability"),
    gr.Textbox(label="Probability"),
    gr.JSON(label="All Labels and Probabilities")
]

demo = gr.Interface(fn=process_url, inputs=[url_input, model_name_input], outputs=outputs)
demo.launch()