Update app.py

app.py

@@ -2,9 +2,9 @@ import gradio as gr
 import json, re, math, os
 from collections import Counter, defaultdict
 
-# =====================================================================
-# Utility Functions
-# =====================================================================
+# ===============================================================
+# UTILITIES
+# ===============================================================
 
 def tokenize(text):
     return re.findall(r"[A-Za-z0-9']+", text.lower())
@@ -15,31 +15,69 @@ def text_vector(text):
 def centroid(docs):
     C = Counter()
     for d in docs:
-        C.update(text_vector(d.get("text", "")))
+        C.update(text_vector(d["text"]))
     return C
 
 def cosine(a, b):
-    num = 0; da = 0; db = 0
+    num = 0
+    da = 0
+    db = 0
     for k in set(a.keys()) | set(b.keys()):
-        va = a.get(k, 0)
-        vb = b.get(k, 0)
-        num += va * vb
-        da += va * va
-        db += vb * vb
+        va = a.get(k,0)
+        vb = b.get(k,0)
+        num += va*vb
+        da += va*va
+        db += vb*vb
     if da == 0 or db == 0:
         return 0
-    return num / math.sqrt(da * db)
+    return num / math.sqrt(da*db)
+
+# ===============================================================
+# LOAD JSONL FROM FILE
+# ===============================================================
+
+def load_records_from_path(path):
+    if not os.path.exists(path):
+        return None, None, "⚠ JSONL file not found."
+
+    records = []
+    with open(path, "r", encoding="utf8") as f:
+        for line in f:
+            try:
+                records.append(json.loads(line))
+            except:
+                pass
+
+    return initialize_state(records)
+
+
+def load_jsonl(user_file):
+    if user_file is None:
+        return None, "⚠ No file uploaded."
+
+    records = []
+    with open(user_file.name, "r", encoding="utf8") as f:
+        for line in f:
+            try:
+                records.append(json.loads(line))
+            except:
+                pass
+
+    state, msg = initialize_state(records)
+    return state, msg
 
-# =====================================================================
-# Load Data
-# =====================================================================
 
 def initialize_state(records):
+    # Ensure IDs exist
+    for i, r in enumerate(records):
+        if "id" not in r:
+            r["id"] = i
+
     cluster_map = defaultdict(list)
     for r in records:
         cluster_map[r.get("cluster", -1)].append(r)
 
-    docs_text = [r.get("text", "") for r in records]
+    docs_text = [r["text"] for r in records]
     tokenized_docs = [tokenize(t) for t in docs_text]
 
     doc_freq = Counter()
@@ -52,49 +90,25 @@ def initialize_state(records):
 
     centroids = {cid: centroid(docs) for cid, docs in cluster_map.items()}
 
-    return {
+    state = {
         "records": records,
         "cluster_map": cluster_map,
         "tokenized_docs": tokenized_docs,
         "doc_freq": doc_freq,
         "Ndocs": Ndocs,
         "avg_len": avg_len,
-        "centroids": centroids
+        "centroids": centroids,
     }
 
-def load_jsonl(user_file):
-    if user_file is None:
-        return None, "⚠ No file uploaded."
-
-    records = []
-    with open(user_file.name, "r", encoding="utf8") as f:
-        for line in f:
-            try:
-                records.append(json.loads(line))
-            except:
-                pass
-    return initialize_state(records), f"Loaded {len(records)} records."
+    return state, f"Loaded {len(records)} records."
 
-def load_default():
-    path = "epstein_semantic.jsonl"
-    if not os.path.exists(path):
-        return None, "⚠ No default dataset found."
-
-    records = []
-    with open(path, "r", encoding="utf8") as f:
-        for line in f:
-            try:
-                records.append(json.loads(line))
-            except:
-                pass
-    return initialize_state(records), f"Loaded {len(records)} records."
 
-# =====================================================================
-# BM25
-# =====================================================================
+# ===============================================================
+# BM25 SEARCH
+# ===============================================================
 
 def bm25_score(query, doc_toks, doc_freq, Ndocs, avg_len):
-    k = 1.5; b = 0.75
+    k=1.5; b=0.75
     score = 0
     q_toks = tokenize(query)
 
@@ -102,7 +116,6 @@ def bm25_score(query, doc_toks, doc_freq, Ndocs, avg_len):
         df = doc_freq.get(q, 0)
         if df == 0:
             continue
-
         idf = math.log((Ndocs - df + 0.5) / (df + 0.5) + 1)
         tf = doc_toks.count(q)
         denom = tf + k * (1 - b + b * (len(doc_toks) / avg_len))
@@ -110,9 +123,10 @@ def bm25_score(query, doc_toks, doc_freq, Ndocs, avg_len):
 
     return score
 
-# =====================================================================
-# UI Functions
-# =====================================================================
+
+# ===============================================================
+# FEATURE FUNCTIONS
+# ===============================================================
 
 def do_view_cluster(state, cid):
     if state is None:
@@ -126,13 +140,11 @@ def do_view_cluster(state, cid):
     cluster_map = state["cluster_map"]
 
     if cid not in cluster_map:
-        return f"❌ Cluster {cid} not found."
+        return "❌ Cluster not found."
 
     out = [f"=== Cluster {cid} ({len(cluster_map[cid])} docs) ===\n"]
-
     for d in cluster_map[cid]:
-        rid = d.get("id", "unknown")
-        out.append(f"\n--- id={rid} ---\n{d.get('text','')}\n")
+        out.append(f"\n--- id={d['id']} ---\n{d['text']}\n")
 
     return "\n".join(out)
 
@@ -142,7 +154,6 @@ def do_search(state, query):
         return "⚠ No dataset loaded."
 
     results = []
-
     for r, toks in zip(state["records"], state["tokenized_docs"]):
         score = bm25_score(query, toks, state["doc_freq"], state["Ndocs"], state["avg_len"])
         if score > 0:
@@ -151,10 +162,27 @@ def do_search(state, query):
     results.sort(key=lambda x: x[0], reverse=True)
 
     out = [f"=== Results for '{query}' ==="]
+    for score, r in results[:30]:
+        out.append(f"\nScore {score:.2f} — Cluster {r['cluster']} — id={r['id']}\n{r['text']}\n")
 
-    for score, r in results[:40]:
-        rid = r.get("id", "unknown")
-        out.append(f"\nScore {score:.2f} — Cluster {r.get('cluster')} — id={rid}\n{r.get('text','')}\n")
+    return "\n".join(out)
+
+
+def do_entity_search(state, name):
+    if state is None:
+        return "⚠ No dataset loaded."
+
+    hits = []
+    for cid, docs in state["cluster_map"].items():
+        count = sum(name.lower() in d["text"].lower() for d in docs)
+        if count:
+            hits.append((count, cid))
+
+    hits.sort(reverse=True)
+
+    out = [f"=== Clusters mentioning '{name}' ==="]
+    for count, cid in hits[:30]:
+        out.append(f"Cluster {cid}: {count} hits")
 
     return "\n".join(out)
 
@@ -163,92 +191,60 @@ def do_show_topics(state):
     if state is None:
         return "⚠ No dataset loaded."
 
-    STOPWORDS = set("""
+    STOP = set("""
 the and to of a in is this that for on with as be or by from at
 an it are was you your if but have we they his her she their our
-subject re fw message thereof all may any doc email said
+subject re fw message thereof all may any doc email
 """.split())
 
     out = ["=== Cluster Topics ==="]
 
     for cid, cent in state["centroids"].items():
-        filtered = {
-            w: c for w, c in cent.items()
-            if w not in STOPWORDS and len(w) > 2 and c > 1
-        }
-
-        top = [w for w, _ in Counter(filtered).most_common(12)]
+        filtered = {w: c for w, c in cent.items()
+                    if w not in STOP and len(w) > 2 and c > 1}
+        top = [w for w, _ in Counter(filtered).most_common(10)]
         out.append(f"Cluster {cid:<4} | {' '.join(top)}")
 
     return "\n".join(out)
 
 
-def do_entity_search(state, name):
-    if state is None:
-        return "⚠ No dataset loaded."
-
-    hits = []
-
-    for cid, docs in state["cluster_map"].items():
-        count = sum(name.lower() in d.get("text", "").lower() for d in docs)
-        if count > 0:
-            hits.append((count, cid))
-
-    hits.sort(reverse=True)
+# ===============================================================
+# AUTO LOAD IF FILE EXISTS
+# ===============================================================
 
-    out = [f"=== Clusters mentioning '{name}' ==="]
-    for count, cid in hits[:40]:
-        out.append(f"Cluster {cid}: {count} hits")
+DEFAULT_PATH = "epstein_semantic.jsonl"
 
-    return "\n".join(out)
+startup_state = None
+startup_msg = "⚠ No default dataset found."
 
-# =====================================================================
-# Startup
-# =====================================================================
+if os.path.exists(DEFAULT_PATH):
+    startup_state, startup_msg = load_records_from_path(DEFAULT_PATH)
 
-startup_state_raw, startup_msg = load_default()
-startup_state = gr.State(startup_state_raw)
 
-# =====================================================================
-# UI
-# =====================================================================
+# ===============================================================
+# GRADIO UI (SINGLE PAGE)
+# ===============================================================
 
 with gr.Blocks(title="Epstein Semantic Explorer") as demo:
 
-    gr.Markdown("# **Epstein Semantic Explorer**")
+    gr.Markdown("# Epstein Semantic Explorer")
     gr.Markdown(startup_msg)
 
-    # Tabs
-    with gr.Tab("View Cluster"):
-        cluster_num = gr.Number(label="Cluster #", value=96)
-        out_cluster = gr.Textbox(label="Cluster Output", lines=40)
-        cluster_num.change(do_view_cluster, [startup_state, cluster_num], out_cluster)
-
-    with gr.Tab("Keyword Search"):
-        query_box = gr.Textbox(label="Keyword")
-        out_search = gr.Textbox(label="Search Output", lines=40)
-        query_box.submit(do_search, [startup_state, query_box], out_search)
-
-    with gr.Tab("Entity Search"):
-        entity_box = gr.Textbox(label="Person / Name")
-        out_entity = gr.Textbox(label="Entity Output", lines=40)
-        entity_box.submit(do_entity_search, [startup_state, entity_box], out_entity)
+    state_box = gr.State(startup_state)
 
-    with gr.Tab("Topics"):
-        out_topics = gr.Textbox(label="Topics", lines=40)
-        gr.Button("Show Topics").click(do_show_topics, [startup_state], out_topics)
+    cluster_input = gr.Number(label="Cluster #", value=0)
+    keyword_input = gr.Textbox(label="Keyword Search")
+    entity_input = gr.Textbox(label="Entity Search (name)")
+    jsonl_file = gr.File(label="Upload different JSONL dataset")
 
-    # Upload override
-    with gr.Tab("Upload Different Dataset"):
-        file_up = gr.File(label="Upload JSONL")
-        load_btn = gr.Button("Load")
-        load_msg = gr.Textbox(label="Status", lines=2)
+    out_box = gr.Textbox(label="Output", lines=40)
 
-        def apply_upload(file):
-            new_state, msg = load_jsonl(file)
-            startup_state.value = new_state
-            return msg
+    # Bindings
+    cluster_input.change(do_view_cluster, [state_box, cluster_input], out_box)
+    keyword_input.submit(do_search, [state_box, keyword_input], out_box)
+    entity_input.submit(do_entity_search, [state_box, entity_input], out_box)
 
-        load_btn.click(apply_upload, [file_up], load_msg)
+    gr.Button("Show Topics").click(do_show_topics, state_box, out_box)
+    gr.Button("Load Dataset").click(load_jsonl, jsonl_file, [state_box, out_box])
 
-demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+demo.launch()