import gradio as gr
from transformers import pipeline

def check_duplicate(sentence1, sentence2):
    model_name = "AventIQ-AI/albert-duplicate-sentence-detection"
    classifier = pipeline("text-classification", model=model_name)
    
    input_text = f"{sentence1} [SEP] {sentence2}"
    result = classifier(input_text)[0]
    
    label_map = {"LABEL_0": "Not Duplicate", "LABEL_1": "Duplicate"}
    label = label_map.get(result['label'], "Unknown")
    
    return f"Label: {label}, Confidence: {result['score']:.4f}"

iface = gr.Interface(
    fn=check_duplicate,
    inputs=[
        gr.Textbox(label="Sentence 1"),
        gr.Textbox(label="Sentence 2")
    ],
    outputs=gr.Textbox(label="Prediction"),
    title="Duplicate Sentence Detection",
    description="Enter two sentences to check if they are duplicates."
)

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