import data_prep
import model_predict
import gradio as gr

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",
    "RoBERTa-Large": "research-dump/roberta-large_deletion_multiclass_complete_final"
}

def process_url(url, model_key):
    model_name = model_dict[model_key]
    processed_text = data_prep.process_data(url)
    final_scores = model_predict.predict_text(processed_text, model_name)
    highest_prob_label = max(final_scores, key=final_scores.get)
    highest_prob = final_scores[highest_prob_label]
    progress_bars = {label: score for label, score in final_scores.items()}
    
    return processed_text, highest_prob_label, highest_prob, progress_bars #,highlighted_text


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="Processed Text"),
    gr.Textbox(label="Label with Highest Probability"),
    gr.Textbox(label="Probability"),
    gr.JSON(label="All Labels and Probabilities"),
    #gr.HTML(label="Processed Text")
]

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