import gradio as gr
from sentence_transformers import CrossEncoder

# Available models
MODEL_OPTIONS = [
    "cross-encoder/ms-marco-MiniLM-L-6-v2",
    "cross-encoder/ms-marco-MiniLM-L-12-v2",
    "cross-encoder/ms-marco-electra-base"
]

model_cache = {}

def load_model(model_name):
    if model_name not in model_cache:
        model_cache[model_name] = CrossEncoder(model_name)
    return model_cache[model_name]

def compute_similarity(model_name, claim, sentences_text):
    if not claim.strip():
        return "Please enter a claim."
    
    sentences = [s.strip() for s in sentences_text.split("\n") if s.strip()]
    
    if len(sentences) == 0:
        return "Please enter at least one sentence."
    
    model = load_model(model_name)
    
    pairs = [(claim, s) for s in sentences]
    scores = model.predict(pairs)
    
    results = list(zip(sentences, scores))
    results.sort(key=lambda x: x[1], reverse=True)
    
    output = "### 🔎 Ranked Results\n\n"
    for sent, score in results:
        output += f"**Score: {score:.4f}**  \n{sent}\n\n"
    
    return output


with gr.Blocks() as demo:
    gr.Markdown("# Semantic Similarity (Cross-Encoder MS MARCO)")
    
    model_dropdown = gr.Dropdown(
        choices=MODEL_OPTIONS,
        value=MODEL_OPTIONS[0],
        label="Choose Cross-Encoder Model"
    )
    
    claim_input = gr.Textbox(
        label="Claim",
        placeholder="Enter claim here..."
    )
    
    sentences_input = gr.Textbox(
        label="Candidate Sentences (one per line)",
        lines=6,
        placeholder="Enter one sentence per line..."
    )
    
    output = gr.Markdown()
    
    run_button = gr.Button("Compute Similarity")
    
    run_button.click(
        compute_similarity,
        inputs=[model_dropdown, claim_input, sentences_input],
        outputs=output
    )

demo.launch()