import gradio as gr
from transformers import pipeline

# Load your model
classifier = pipeline("text-classification", model="sachin6624/bert-trainer")

# Map labels to custom messages
label_mapping = {
    "LABEL_1": "The sentences are paraphrases (similar in meaning).",
    "LABEL_0": "The sentences are not paraphrases (different in meaning)."
}

# Custom prediction function
def predict(sentence1, sentence2):
    input_text = f"{sentence1} [SEP] {sentence2}"
    result = classifier(input_text)
    label = result[0]["label"]
    # Return the custom message for the predicted label
    return label_mapping.get(label, "Unknown label")

# Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=["text", "text"],
    outputs="text",  # Output is text for custom messages
    title="Paraphrase Classification"
)

# Launch the app
interface.launch()