import gradio as gr
from transformers import pipeline

# Load the classification pipeline
classifier = pipeline(
    "sentiment-analysis", 
    # model="Karzan/user_profile_model", 
    model="Karzan/roles-based-on-skills-model",
    return_all_scores=True, 
    top_k=10
)

# Define the prediction function
def classify_text(text):
    # Perform classification
    results = classifier(text)
    # Format the output
    formatted_results = [
        {"label": item["label"], "score": round(item["score"], 4)}
        for result in results for item in result
    ]
    output = {}
    print(formatted_results)
    for i in range(len(formatted_results)):
        output[formatted_results[i]['label']] = formatted_results[i]['score']
    return output

demo = gr.Interface(fn=classify_text, inputs=[gr.Textbox(label="Input")], outputs=gr.Label(label="Classification"), title="Text Classification")
demo.launch()