Update app.py

app.py

@@ -1,85 +1,72 @@
 import gradio as gr
+from sentence_transformers import CrossEncoder
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, BertConfig
 
-# ── Model setup ──────────────────────────────────────────────────────────────
+# Load fine-tuned model from Hub
 MODEL_NAME = "pujithapsx/address-crossencoder-bge-reranker-v2-m3-finetuned"
-THRESHOLD  = 0.75
 
-print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-
-print("Loading model weights...")
-# The model's config.json is missing "model_type", which breaks AutoConfig/CrossEncoder.
-# BertConfig.from_pretrained() reads the repo's config.json without needing model_type.
-config   = BertConfig.from_pretrained(MODEL_NAME)
-hf_model = AutoModelForSequenceClassification.from_pretrained(
+print("Loading model...")
+# Use CrossEncoder directly - it handles the model format correctly
+model = CrossEncoder(
     MODEL_NAME,
-    config=config,
-    ignore_mismatched_sizes=True,
+    trust_remote_code=True,
+    device="cpu"  # Use "cuda" if GPU is available in your HF Space
 )
-hf_model.eval()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-hf_model.to(device)
-print(f"Model loaded on {device}!")
-
+print("Model loaded successfully!")
 
-# ── Inference ─────────────────────────────────────────────────────────────────
-def _score_pair(text_a: str, text_b: str) -> float:
-    """Tokenise a pair and return a 0-1 similarity score."""
-    features = tokenizer(
-        text_a, text_b,
-        padding=True,
-        truncation=True,
-        max_length=512,
-        return_tensors="pt",
-    )
-    features = {k: v.to(device) for k, v in features.items()}
-    with torch.no_grad():
-        logits = hf_model(**features).logits
-    # single logit → sigmoid;  two logits → softmax of positive class
-    if logits.shape[-1] == 1:
-        return torch.sigmoid(logits[0, 0]).item()
-    else:
-        return torch.softmax(logits[0], dim=-1)[1].item()
+# STATIC THRESHOLD
+THRESHOLD = 0.75
 
 
 def predict_similarity(input1, input2):
+    """
+    Predict similarity between two addresses using static threshold.
+    Returns: Similarity %, Match/No Match result, and confidence bar value.
+    """
     if not input1.strip() or not input2.strip():
         return "—", "⚠️ Please provide both addresses", 0
 
-    score         = _score_pair(input1.strip(), input2.strip())
+    score = model.predict([[input1, input2]])[0]
     similarity_pct = score * 100
 
     if score >= THRESHOLD:
-        result           = "✅ MATCH"
+        result = "✅ MATCH"
         confidence_label = "High" if score > 0.85 else "Medium"
     else:
-        result           = "❌ NO MATCH"
+        result = "❌ NO MATCH"
         confidence_label = "High" if score < 0.40 else "Medium"
 
-    return f"{similarity_pct:.2f}%", f"{result}  •  Confidence: {confidence_label}", float(similarity_pct)
+    similarity_str = f"{similarity_pct:.2f}%"
+    result_str = f"{result}  •  Confidence: {confidence_label}"
+
+    return similarity_str, result_str, float(similarity_pct)
 
 
-# ── Custom CSS ────────────────────────────────────────────────────────────────
+# ── Custom CSS ──────────────────────────────────────────────────────────────
 custom_css = """
 @import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=DM+Sans:wght@300;400;500;600&display=swap');
 
 :root {
-    --bg:      #0b0f1a;
-    --surface: #111827;
-    --border:  #1f2d45;
-    --accent:  #38bdf8;
-    --accent2: #818cf8;
-    --text:    #e2e8f0;
-    --muted:   #64748b;
-    --radius:  12px;
+    --bg:        #0b0f1a;
+    --surface:   #111827;
+    --border:    #1f2d45;
+    --accent:    #38bdf8;
+    --accent2:   #818cf8;
+    --green:     #34d399;
+    --red:       #f87171;
+    --yellow:    #fbbf24;
+    --text:      #e2e8f0;
+    --muted:     #64748b;
+    --radius:    12px;
 }
+
 body, .gradio-container {
     background: var(--bg) !important;
     font-family: 'DM Sans', sans-serif !important;
     color: var(--text) !important;
 }
+
+/* ── header ── */
 #header-box {
     background: linear-gradient(135deg, #0f172a 0%, #1e293b 100%);
     border: 1px solid var(--border);
@@ -103,6 +90,7 @@ body, .gradio-container {
     font-weight: 700 !important;
     color: var(--accent) !important;
     margin: 0 0 8px !important;
+    letter-spacing: -0.5px;
 }
 #header-box p {
     color: var(--muted) !important;
@@ -122,7 +110,10 @@ body, .gradio-container {
     margin-right: 8px;
     margin-top: 12px;
 }
-.input-card textarea, .input-card input {
+
+/* ── input cards ── */
+.input-card textarea,
+.input-card input {
     background: var(--surface) !important;
     border: 1px solid var(--border) !important;
     border-radius: var(--radius) !important;
@@ -132,7 +123,8 @@ body, .gradio-container {
     padding: 14px 16px !important;
     transition: border-color 0.2s;
 }
-.input-card textarea:focus, .input-card input:focus {
+.input-card textarea:focus,
+.input-card input:focus {
     border-color: var(--accent) !important;
     box-shadow: 0 0 0 3px rgba(56,189,248,0.1) !important;
 }
@@ -143,6 +135,8 @@ label span {
     letter-spacing: 0.5px;
     text-transform: uppercase;
 }
+
+/* ── button ── */
 #run-btn {
     background: linear-gradient(135deg, var(--accent) 0%, var(--accent2) 100%) !important;
     border: none !important;
@@ -151,14 +145,19 @@ label span {
     font-family: 'Space Mono', monospace !important;
     font-size: 0.9rem !important;
     font-weight: 700 !important;
+    letter-spacing: 0.5px;
     padding: 14px 32px !important;
+    cursor: pointer;
+    transition: opacity 0.2s, transform 0.15s;
     width: 100%;
     margin-top: 8px;
-    transition: opacity 0.2s, transform 0.15s;
 }
-#run-btn:hover  { opacity: 0.9; transform: translateY(-1px); }
+#run-btn:hover { opacity: 0.9; transform: translateY(-1px); }
 #run-btn:active { transform: translateY(0); }
-.output-card textarea, .output-card input {
+
+/* ── output cards ── */
+.output-card textarea,
+.output-card input {
     background: #0d1424 !important;
     border: 1px solid var(--border) !important;
     border-radius: var(--radius) !important;
@@ -168,7 +167,13 @@ label span {
     font-weight: 700 !important;
     text-align: center;
 }
-.score-slider input[type=range] { accent-color: var(--accent); }
+
+/* ── slider (score bar) ── */
+.score-slider input[type=range] {
+    accent-color: var(--accent);
+}
+
+/* ── examples table ── */
 .gr-samples-table {
     background: var(--surface) !important;
     border: 1px solid var(--border) !important;
@@ -179,10 +184,18 @@ label span {
     font-size: 0.72rem !important;
     color: var(--muted) !important;
     text-transform: uppercase;
+    letter-spacing: 0.5px;
     background: #0d1424 !important;
 }
-.gr-samples-table td  { color: var(--text) !important; font-size: 0.88rem !important; }
-.gr-samples-table tr:hover td { background: rgba(56,189,248,0.04) !important; }
+.gr-samples-table td {
+    color: var(--text) !important;
+    font-size: 0.88rem !important;
+}
+.gr-samples-table tr:hover td {
+    background: rgba(56,189,248,0.04) !important;
+}
+
+/* ── info footer ── */
 #footer-info {
     background: var(--surface);
     border: 1px solid var(--border);
@@ -200,9 +213,10 @@ label span {
 #footer-info span { color: var(--accent) !important; }
 """
 
-# ── Gradio UI ─────────────────────────────────────────────────────────────────
+# ── Gradio UI ────────────────────────────────────────────────────────────────
 with gr.Blocks(css=custom_css, title="Address Entity Matcher") as demo:
 
+    # Header
     gr.HTML("""
     <div id="header-box">
         <h1>📍 Address Entity Matcher</h1>
@@ -210,11 +224,13 @@ with gr.Blocks(css=custom_css, title="Address Entity Matcher") as demo:
             Enter two addresses to determine whether they refer to the same location.<br>
             Powered by a fine-tuned <strong>BGE-Reranker-v2-m3</strong> cross-encoder model.
         </p>
+        <span class="badge">CrossEncoder</span>
         <span class="badge">BGE-Reranker-v2-m3</span>
         <span class="badge">Threshold: 0.75</span>
     </div>
     """)
 
+    # Inputs
     with gr.Row(equal_height=True):
         with gr.Column(elem_classes="input-card"):
             input1 = gr.Textbox(
@@ -231,39 +247,55 @@ with gr.Blocks(css=custom_css, title="Address Entity Matcher") as demo:
 
     btn = gr.Button("🔎 Check Match", elem_id="run-btn", variant="primary")
 
+    # Outputs
     with gr.Row(equal_height=True):
         with gr.Column(elem_classes="output-card"):
-            similarity_output = gr.Textbox(label="Similarity Score", interactive=False, placeholder="—")
+            similarity_output = gr.Textbox(
+                label="Similarity Score",
+                interactive=False,
+                placeholder="—",
+            )
         with gr.Column(elem_classes="output-card"):
-            result_output = gr.Textbox(label="Match Result", interactive=False, placeholder="—")
+            result_output = gr.Textbox(
+                label="Match Result",
+                interactive=False,
+                placeholder="—",
+            )
 
     score_bar = gr.Slider(
-        minimum=0, maximum=100, value=0, step=0.01,
+        minimum=0,
+        maximum=100,
+        value=0,
+        step=0.01,
         label="Score Visualisation  (threshold line: 75%)",
-        interactive=False, elem_classes="score-slider",
+        interactive=False,
+        elem_classes="score-slider",
     )
 
+    # Examples
     gr.Examples(
         examples=[
-            ["Flat 12-B Sector 5 Noida",                        "Flat 23-B Sector 5 Noida"],
-            ["Phase 4 Whitefield Bangalore",                     "Whitefield Phase V Bangalore"],
-            ["Thirteen I 7th Avenue Adyar Chennai",              "13 Seventh Avenue Adyar Chennai"],
-            ["Twenty Nine A Second Cross Koramangala Bengaluru", "47 Forty Seven B Third Street Indiranagar Bengaluru"],
-            ["Plot 8 Banjara Hills Hyderabad",                   "Plot 8 Banjara Hills Hyderabad"],
-            ["House No 4 Lane 2 DLF Phase 1 Gurugram",          "H.No 4/2 DLF Phase One Gurgaon"],
+            ["Flat 12-B Sector 5 Noida",                          "Flat 23-B Sector 5 Noida"],
+            ["Phase 4 Whitefield Bangalore",                       "Whitefield Phase V Bangalore"],
+            ["Thirteen I 7th Avenue Adyar Chennai",                "13 Seventh Avenue Adyar Chennai"],
+            ["Twenty Nine A Second Cross Koramangala Bengaluru",   "47 Forty Seven B Third Street Indiranagar Bengaluru"],
+            ["Plot 8 Banjara Hills Hyderabad",                     "Plot 8 Banjara Hills Hyderabad"],
+            ["House No 4 Lane 2 DLF Phase 1 Gurugram",            "H.No 4/2 DLF Phase One Gurgaon"],
         ],
         inputs=[input1, input2],
         label="📋 Try these examples",
     )
 
+    # Footer info
     gr.HTML(f"""
     <div id="footer-info">
-        <p>🤖 <strong>Model:</strong> <span>{MODEL_NAME}</span></p>
+        <p>🤖 <strong>Model:</strong> <span>pujithapsx/address-crossencoder-bge-reranker-v2-m3-finetuned</span></p>
         <p>📏 <strong>Threshold:</strong> <span>{THRESHOLD}</span> — Score ≥ {THRESHOLD} → MATCH &nbsp;|&nbsp; Score &lt; {THRESHOLD} → NO MATCH</p>
         <p>🏷️  <strong>Confidence:</strong> High (score &gt; 0.85 or &lt; 0.40) &nbsp;|&nbsp; Medium (otherwise)</p>
     </div>
     """)
 
+    # Wiring
     btn.click(
         fn=predict_similarity,
         inputs=[input1, input2],