Update app.py

app.py

@@ -2,9 +2,9 @@ import gradio as gr
 import json, re, math, os
 from collections import Counter, defaultdict
 
-# ===============================================================
-# UTILITIES
-# ===============================================================
+# =====================================================================
+# Utility Functions
+# =====================================================================
 
 def tokenize(text):
     return re.findall(r"[A-Za-z0-9']+", text.lower())
@@ -19,62 +19,27 @@ def centroid(docs):
     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)
-
-# ===============================================================
-# LOAD JSONL FROM FILE
-# ===============================================================
-
-def load_records_from_path(path):
-    """Loads a dataset from an existing file, used at startup."""
-    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):
-    """Loads a dataset from user upload."""
-    if user_file is None:
-        return gr.update(), 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)
+    return num / math.sqrt(da * db)
 
+# =====================================================================
+# Load Data
+# =====================================================================
 
 def initialize_state(records):
-    """Builds all indexes for search, clustering, etc."""
     cluster_map = defaultdict(list)
     for r in records:
         cluster_map[r.get("cluster", -1)].append(r)
 
-    docs_text = [r["text"] for r in records]
+    docs_text = [r.get("text", "") for r in records]
     tokenized_docs = [tokenize(t) for t in docs_text]
 
     doc_freq = Counter()
@@ -94,16 +59,44 @@ def initialize_state(records):
         "doc_freq": doc_freq,
         "Ndocs": Ndocs,
         "avg_len": avg_len,
-        "centroids": centroids,
+        "centroids": centroids
     }, sorted(cluster_map.keys()), f"Loaded {len(records)} records."
 
 
-# ===============================================================
-# BM25 SEARCH
-# ===============================================================
+def load_jsonl(user_file):
+    if user_file is None:
+        return None, 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)
+
+
+def load_default():
+    path = "epstein_semantic.jsonl"
+    if not os.path.exists(path):
+        return None, None, "⚠ Upload a dataset to begin."
+
+    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)
+
+# =====================================================================
+# BM25
+# =====================================================================
 
 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)
 
@@ -111,17 +104,18 @@ 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))
+
         score += idf * (tf * (k + 1)) / denom
 
     return score
 
-
-# ===============================================================
-# GRADIO FEATURE FUNCTIONS
-# ===============================================================
+# =====================================================================
+# UI Functions
+# =====================================================================
 
 def do_view_cluster(state, cid):
     if state is None:
@@ -135,12 +129,14 @@ def do_view_cluster(state, cid):
     cluster_map = state["cluster_map"]
 
     if cid not in cluster_map:
-        return "❌ Cluster not found."
+        return f"❌ Cluster {cid} not found."
 
-    # FULL TEXT (NO MORE TRUNCATION)
     out = [f"=== Cluster {cid} ({len(cluster_map[cid])} docs) ===\n"]
+
+    # show all docs, untruncated
     for d in cluster_map[cid]:
-        out.append(f"\n--- id={d.get('id')} ---\n{d['text']}\n")
+        rid = d.get("id", "unknown")
+        out.append(f"\n--- id={rid} ---\n{d.get('text','')}\n")
 
     return "\n".join(out)
 
@@ -150,16 +146,20 @@ 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:
             results.append((score, r))
 
-    results.sort(reverse=True)
+    # FIX: sort by score, not dict
+    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)
 
@@ -171,7 +171,7 @@ def do_show_topics(state):
     STOPWORDS = 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
+subject re fw message thereof all may any doc email said
 """.split())
 
     out = ["=== Cluster Topics ==="]
@@ -179,7 +179,9 @@ subject re fw message thereof all may any doc email
     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(10)]
+
+        top = [w for w, _ in Counter(filtered).most_common(12)]
+
         out.append(f"Cluster {cid:<4} | {' '.join(top)}")
 
     return "\n".join(out)
@@ -190,59 +192,55 @@ def do_entity_search(state, name):
         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:
+        count = sum(name.lower() in d.get("text", "").lower() for d in docs)
+        if count > 0:
             hits.append((count, cid))
 
     hits.sort(reverse=True)
 
     out = [f"=== Clusters mentioning '{name}' ==="]
-    for count, cid in hits[:30]:
+    for count, cid in hits[:40]:
         out.append(f"Cluster {cid}: {count} hits")
 
     return "\n".join(out)
 
+# =====================================================================
+# UI Layout
+# =====================================================================
 
-# ===============================================================
-# AUTO-LOAD DATASET IF PRESENT
-# ===============================================================
-
-DEFAULT_PATH = "epstein_semantic.jsonl"
-
-startup_state = None
-startup_clusters = None
-startup_msg = "⚠ No default dataset found."
-
-if os.path.exists(DEFAULT_PATH):
-    startup_state, startup_clusters, startup_msg = load_records_from_path(DEFAULT_PATH)
+startup_state, startup_clusters, startup_msg = load_default()
 
+with gr.Blocks(title="Epstein Semantic Explorer", css="#output {white-space: pre-wrap;}") as demo:
 
-# ===============================================================
-# GRADIO UI
-# ===============================================================
-
-with gr.Blocks(title="Epstein Semantic Explorer") as demo:
-
-    gr.Markdown("# Epstein Semantic Explorer")
+    gr.Markdown("# **Epstein Semantic Explorer**")
     gr.Markdown(startup_msg)
 
-    with gr.Row():
-        jsonl_file = gr.File(label="Upload different JSONL dataset")
-        load_btn = gr.Button("Load Dataset")
+    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)
 
-    state = gr.State(startup_state)
-    clusters_box = gr.Number(label="Cluster #", value=96)
-    query_box = gr.Textbox(label="Keyword Search")
-    entity_box = gr.Textbox(label="Entity Search (name)")
+    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)
 
-    output = gr.Textbox(label="Output", lines=40)
+    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)
 
-    load_btn.click(load_jsonl, inputs=[jsonl_file], outputs=[state, clusters_box, output])
-    clusters_box.change(do_view_cluster, inputs=[state, clusters_box], outputs=output)
-    query_box.submit(do_search, inputs=[state, query_box], outputs=output)
-    entity_box.submit(do_entity_search, inputs=[state, entity_box], outputs=output)
+    with gr.Tab("Topics"):
+        out_topics = gr.Textbox(label="Topics", lines=40)
+        gr.Button("Show Topics").click(do_show_topics, [startup_state], out_topics)
 
-    gr.Button("Show Topics").click(do_show_topics, inputs=[state], outputs=output)
+    # File Upload (override default)
+    with gr.Tab("Upload Different Dataset"):
+        jsonl_file = gr.File(label="Upload JSONL")
+        load_btn = gr.Button("Load Dataset")
+        load_out = gr.Textbox(label="Status", lines=2)
+        load_btn.click(load_jsonl, [jsonl_file], [startup_state, cluster_num, load_out])
 
-demo.launch()
+demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)