import gradio as gr
from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="WarTitan2077/Number-Classifier",
    tokenizer="WarTitan2077/Number-Classifier",
    top_k=None
)

labels = ["Symbolic", "Numeric", "Natural", "Integer", "Rational", "Irrational", "Real", "Prime", "Composite"]
label_map = {f"LABEL_{i}": label for i, label in enumerate(labels)}

def classify_numbers(input1, input2, input3):
    inputs = {"Input1": input1, "Input2": input2, "Input3": input3}
    results = {}
    for name, value in inputs.items():
        if value.strip():
            output = classifier(value)[0]
            result = {label_map[item["label"]]: round(item["score"], 3) for item in output}
            results[name] = result
        else:
            results[name] = {}
    return results

demo = gr.Interface(
    fn=classify_numbers,
    inputs=[
        gr.Textbox(label="Input 1"),
        gr.Textbox(label="Input 2"),
        gr.Textbox(label="Input 3")
    ],
    outputs=gr.JSON(label="Predictions"),
    api_name="/predict"
)

if __name__ == "__main__":
    demo.launch()