Karzan's picture
Update app.py
9d786ad verified
raw
history blame contribute delete
909 Bytes
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()