import gradio as gr
from sentence_transformers import CrossEncoder
import torch

# Load MS MARCO CrossEncoder (query-document relevance)
MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-12-v2"
model = CrossEncoder(MODEL_NAME)

def predict_similarity(s1, s2):
    # Get raw score (logit)
    score = model.predict([(s1, s2)])[0]
    # Apply sigmoid to map to 0–1
    similarity = torch.sigmoid(torch.tensor(score)).item()
    return round(similarity, 4)

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🔎 Query–Document Relevance (CrossEncoder)")
    gr.Markdown(
        f"Model: **{MODEL_NAME}**\n\n"
        "Scores are mapped with sigmoid to the range **0 (irrelevant) → 1 (highly relevant)**."
    )

    with gr.Row():
        s1 = gr.Textbox(label="Query", placeholder="Enter your search query...")
        s2 = gr.Textbox(label="Document Chunk", placeholder="Enter a document chunk...")

    btn = gr.Button("Compute Similarity 🚀")
    out = gr.Number(label="Relevance Score (0–1)")

    btn.click(fn=predict_similarity, inputs=[s1, s2], outputs=out)

    gr.Examples(
        examples=[
            ["What is the capital of France?", "Paris is the capital city of France."],
            ["I am happy today", "I am feeling joyful and excited right now."],
            ["Python programming", "Bananas are yellow fruits."],
            ["Machine learning applications", "ML is widely used in healthcare and finance."],
        ],
        inputs=[s1, s2],
    )

demo.launch(enable_queue=True)