import gradio as gr
from bertopic import BERTopic

topic_model = BERTopic.load("TarekAi/ArBERTopic/my_model")

def predict_topic(document):
    topics, probabilities = topic_model.transform([document])
    return {"Predicted Topic": topics[0], "Probabilities": probabilities[0].tolist()}

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_topic,
    inputs=gr.inputs.Textbox(lines=5, placeholder="Enter your idea here..."),
    outputs=[
        gr.outputs.Textbox(label="Predicted Topic"),
        gr.outputs.Label(label="Topic Probabilities")
    ],
    title="BERTopic Inference",
    description="Enter a single idea to get the predicted topic and its probabilities."
)

# Launch the interface
iface.launch()