import gradio as gr
from sentence_transformers import CrossEncoder
import torch
import requests

# -------------------------------
# CONFIG
# -------------------------------
HF_MODEL = "cross-encoder/ms-marco-MiniLM-L-12-v2"
JINA_MODEL = "jina-reranker-m0"
JINA_API_KEY = "jina_4075150fa702471c85ddea0a9ad4b306ouE7ymhrCpvxTxX3mScUv5LLDPKQ"
JINA_ENDPOINT = "https://api.jina.ai/v1/rerank"

# -------------------------------
# Load Hugging Face CrossEncoder
# -------------------------------
hf_model = CrossEncoder(HF_MODEL)

def compare_models(query, doc):
    # Hugging Face score
    raw_score = hf_model.predict([(query, doc)])[0]
    hf_score = torch.sigmoid(torch.tensor(raw_score)).item()

    # Jina reranker score
    headers = {
        "Authorization": f"Bearer {JINA_API_KEY}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": JINA_MODEL,
        "query": query,
        "documents": [doc],
    }
    try:
        r = requests.post(JINA_ENDPOINT, headers=headers, json=payload, timeout=20)
        r.raise_for_status()
        jina_score = r.json()["results"][0]["relevance_score"]
    except Exception as e:
        jina_score = f"Error: {str(e)}"

    return f"Hugging Face ({HF_MODEL}): {round(hf_score,4)}\nJina ({JINA_MODEL}): {jina_score}"

# -------------------------------
# Simple Lite UI
# -------------------------------
with gr.Blocks() as demo:
    gr.Markdown("### 🔎 Query vs Document Similarity (HF vs Jina)")
    query = gr.Textbox(label="Query", lines=3, placeholder="Paste your query here...")
    doc = gr.Textbox(label="Document Chunk", lines=6, placeholder="Paste your document chunk here...")
    out = gr.Textbox(label="Scores", lines=3)

    btn = gr.Button("Compute Similarity 🚀")
    btn.click(compare_models, inputs=[query, doc], outputs=out)

demo.launch()