Update app.py

app.py

@@ -1,74 +1,85 @@
 import gradio as gr
-from sentence_transformers import CrossEncoder
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, BertConfig
 
-# Load fine-tuned model from Hub
+# ── Model setup ──────────────────────────────────────────────────────────────
 MODEL_NAME = "pujithapsx/address-crossencoder-bge-reranker-v2-m3-finetuned"
+THRESHOLD  = 0.75
 
-print("Loading model...")
-#model = CrossEncoder(MODEL_NAME)
-# model = CrossEncoder(
-#     MODEL_NAME,
-#     trust_remote_code=True,                          # allow custom model code
-#     automodel_args={"ignore_mismatched_sizes": True}, # skip size mismatch errors
-#     config_args={"model_type": "bert"},              # inject missing model_type
-# )
-config   = BertConfig.from_pretrained(MODEL_NAME)   # BertConfig doesn't need model_type
-hf_model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, config=config)
-print("Model loaded successfully!")
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# STATIC THRESHOLD
-THRESHOLD = 0.75
+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(
+    MODEL_NAME,
+    config=config,
+    ignore_mismatched_sizes=True,
+)
+hf_model.eval()
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+hf_model.to(device)
+print(f"Model loaded on {device}!")
+
+
+# ── 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()
 
 
 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 = model.predict([[input1, input2]])[0]
+    score         = _score_pair(input1.strip(), input2.strip())
     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"
 
-    similarity_str = f"{similarity_pct:.2f}%"
-    result_str = f"{result}  •  Confidence: {confidence_label}"
-
-    return similarity_str, result_str, float(similarity_pct)
+    return f"{similarity_pct:.2f}%", f"{result}  •  Confidence: {confidence_label}", 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;
-    --green:     #34d399;
-    --red:       #f87171;
-    --yellow:    #fbbf24;
-    --text:      #e2e8f0;
-    --muted:     #64748b;
-    --radius:    12px;
+    --bg:      #0b0f1a;
+    --surface: #111827;
+    --border:  #1f2d45;
+    --accent:  #38bdf8;
+    --accent2: #818cf8;
+    --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);
@@ -92,7 +103,6 @@ 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;
@@ -112,10 +122,7 @@ body, .gradio-container {
     margin-right: 8px;
     margin-top: 12px;
 }
-
-/* ── input cards ── */
-.input-card textarea,
-.input-card input {
+.input-card textarea, .input-card input {
     background: var(--surface) !important;
     border: 1px solid var(--border) !important;
     border-radius: var(--radius) !important;
@@ -125,8 +132,7 @@ 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;
 }
@@ -137,8 +143,6 @@ 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;
@@ -147,19 +151,14 @@ 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 cards ── */
-.output-card textarea,
-.output-card input {
+.output-card textarea, .output-card input {
     background: #0d1424 !important;
     border: 1px solid var(--border) !important;
     border-radius: var(--radius) !important;
@@ -169,13 +168,7 @@ label span {
     font-weight: 700 !important;
     text-align: center;
 }
-
-/* ── slider (score bar) ── */
-.score-slider input[type=range] {
-    accent-color: var(--accent);
-}
-
-/* ── examples table ── */
+.score-slider input[type=range] { accent-color: var(--accent); }
 .gr-samples-table {
     background: var(--surface) !important;
     border: 1px solid var(--border) !important;
@@ -186,18 +179,10 @@ 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;
-}
-
-/* ── info footer ── */
+.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; }
 #footer-info {
     background: var(--surface);
     border: 1px solid var(--border);
@@ -215,10 +200,9 @@ 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>
@@ -226,13 +210,11 @@ 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(
@@ -249,55 +231,39 @@ 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>pujithapsx/address-crossencoder-bge-reranker-v2-m3-finetuned</span></p>
+        <p>🤖 <strong>Model:</strong> <span>{MODEL_NAME}</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],
@@ -305,4 +271,4 @@ with gr.Blocks(css=custom_css, title="Address Entity Matcher") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()