Create main.py

main.py

@@ -0,0 +1,270 @@
+import os, re, json, io, zipfile, shutil, math, gc, uuid
+from datetime import datetime
+from typing import List, Dict, Any, Tuple
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from bs4 import BeautifulSoup
+import ftfy
+from langdetect import detect, DetectorFactory
+DetectorFactory.seed = 0
+
+import gradio as gr
+from tqdm import tqdm
+
+from sentence_transformers import SentenceTransformer
+import faiss
+
+try:
+    import hdbscan
+    HDBSCAN_AVAILABLE = True
+except Exception:
+    HDBSCAN_AVAILABLE = False
+
+# ------------------- Helpers -------------------
+URL = re.compile(r"https?://\S+", re.I)
+
+
+def torch_cuda_available():
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except Exception:
+        return False
+
+
+def html_to_text(html: str) -> str:
+    if not html:
+        return ""
+    soup = BeautifulSoup(html, "html.parser")
+    for tag in soup(["script", "style"]):
+        tag.decompose()
+    return soup.get_text(separator="\n")
+
+
+def strip_quotes_and_sigs(text: str) -> str:
+    if not text:
+        return ""
+    text = re.sub(r"^>.*$", "", text, flags=re.M)  # drop quote lines
+    res = re.split(r"\n-- ?\n", text)[0]
+    res = re.split(r"\nSent from my ", res)[0]
+    return res.strip()
+
+
+def parse_name_email(s: str) -> Tuple[str, str]:
+    if not s:
+        return "", ""
+    m = re.match(r'(?:"?([^"]*)"?\s)?<?([^<>]+@[^<>]+)>?', s)
+    if m:
+        return (m.group(1) or "").strip(), (m.group(2) or "").strip()
+    return "", s.strip()
+
+# ------------------- Normalization -------------------
+
+def normalize_email_record(raw: Dict[str, Any]) -> Dict[str, Any]:
+    subject = raw.get("subject") or raw.get("Subject") or ""
+    body_html = raw.get("body_html") or raw.get("html") or ""
+    body_text = raw.get("body_text") or raw.get("text") or ""
+    headers = raw.get("headers") or {}
+
+    date_val = raw.get("date") or raw.get("Date") or headers.get("Date") or ""
+    msg_id = raw.get("message_id") or raw.get("Message-ID") or headers.get("Message-ID") or ""
+    sender = raw.get("from") or raw.get("From") or headers.get("From") or ""
+
+    if body_html and not body_text:
+        body_text = html_to_text(body_html)
+
+    subject = ftfy.fix_text(subject or "").strip()
+    body_text = ftfy.fix_text(body_text or "")
+    body_text = strip_quotes_and_sigs(body_text)
+    body_text = re.sub(r"\s+", " ", body_text).strip()
+    subject_norm = re.sub(r"\s+", " ", subject)
+
+    try:
+        lang = detect((subject_norm + " " + body_text[:5000]).strip()) if (subject_norm or body_text) else "unknown"
+    except Exception:
+        lang = "unknown"
+
+    from_name, from_email = parse_name_email(sender)
+    from_domain = from_email.split("@")[-1].lower() if "@" in from_email else ""
+
+    iso_date = ""
+    if isinstance(date_val, (int, float)):
+        try:
+            iso_date = pd.to_datetime(int(date_val), unit="s", utc=True).isoformat()
+        except Exception:
+            pass
+    elif isinstance(date_val, str) and date_val:
+        try:
+            iso_date = pd.to_datetime(date_val, utc=True, errors="coerce").isoformat()
+        except Exception:
+            pass
+
+    thread_key = subject_norm or msg_id
+    thread_id = str(pd.util.hash_pandas_object(pd.Series([thread_key], dtype="string")).astype("uint64").iloc[0])
+
+    text_hash = str(pd.util.hash_pandas_object(pd.Series([body_text], dtype="string")).astype("uint64").iloc[0]) if body_text else ""
+
+    if not msg_id:
+        msg_id = f"gen-{uuid.uuid4().hex}"
+
+    return {
+        "message_id": str(msg_id),
+        "thread_id": thread_id,
+        "date": iso_date,
+        "from_name": from_name,
+        "from_email": from_email,
+        "from_domain": from_domain,
+        "subject": subject_norm,
+        "body_text": body_text,
+        "lang": lang,
+        "text_hash": text_hash,
+    }
+
+# ------------------- Embeddings & Clustering -------------------
+
+def embed_texts(model: SentenceTransformer, texts: List[str], batch_size: int, use_gpu: bool) -> np.ndarray:
+    embs = []
+    for i in tqdm(range(0, len(texts), batch_size), desc="Embedding", leave=False):
+        chunk = texts[i:i + batch_size]
+        embs.append(model.encode(
+            chunk,
+            batch_size=min(256, len(chunk)),
+            show_progress_bar=False,
+            normalize_embeddings=True,
+            convert_to_numpy=True,
+            device="cuda" if use_gpu else "cpu",
+        ))
+    return np.vstack(embs).astype(np.float32)
+
+
+def cluster_embeddings(embs: np.ndarray, method: str, min_cluster_size: int, k_hint: int, use_gpu: bool) -> np.ndarray:
+    if method == "HDBSCAN" and HDBSCAN_AVAILABLE:
+        clust = hdbscan.HDBSCAN(min_cluster_size=min_cluster_size, min_samples=max(5, min_cluster_size // 5), metric='euclidean')
+        return clust.fit_predict(embs)
+    k = max(10, k_hint or int(max(20, math.sqrt(len(embs) / 50))))
+    kmeans = faiss.Kmeans(d=embs.shape[1], k=k, niter=25, verbose=False, gpu=use_gpu)
+    kmeans.train(embs)
+    _, labels = kmeans.index.search(embs, 1)
+    return labels.reshape(-1)
+
+
+def zero_shot_embed_sim(embs: np.ndarray, model: SentenceTransformer, label_texts: List[str], use_gpu: bool) -> Tuple[np.ndarray, np.ndarray]:
+    prompts = [f"This email is about: {t}" for t in label_texts]
+    label_embs = model.encode(prompts, normalize_embeddings=True, convert_to_numpy=True, device="cuda" if use_gpu else "cpu").astype(np.float32)
+    sims = embs @ label_embs.T
+    top_idx = sims.argmax(axis=1)
+    top_score = sims[np.arange(len(embs)), top_idx]
+    return top_idx, top_score
+
+# ------------------- Defaults -------------------
+DEFAULT_LABELS = [
+    "Newsletters/Subscriptions",
+    "Receipts & Billing",
+    "Personal/Family",
+    "Work/Colleagues",
+    "Meetings & Calendars",
+    "Travel/Itineraries",
+    "Legal/Contracts",
+    "System Notifications",
+    "Security/2FA",
+    "Hiring/Recruiting",
+    "Support Tickets",
+    "Politics/Government",
+    "Media/Press",
+    "Unknown"
+]
+
+# ------------------- Search -------------------
+class EmailSearch:
+    def __init__(self, df, embs, model):
+        self.df = df
+        self.embs = embs
+        self.model = model
+        self.index = faiss.IndexFlatIP(embs.shape[1])
+        self.index.add(embs)
+
+    def query(self, q: str, top_k=20):
+        q_emb = self.model.encode([q], normalize_embeddings=True, convert_to_numpy=True)
+        scores, idx = self.index.search(q_emb.astype(np.float32), top_k)
+        results = self.df.iloc[idx[0]].copy()
+        results["score"] = scores[0]
+        return results
+
+# ------------------- Gradio UI -------------------
+with gr.Blocks(title="Email Organizer & Browser") as demo:
+    gr.Markdown("""
+    # Email Organizer & Browser (No-Redaction)
+    Upload a **.jsonl** or **.json** of emails. The app normalizes, deduplicates, embeds, clusters, labels, and lets you **search** your inbox semantically.
+    """)
+
+    with gr.Row():
+        inbox_file = gr.File(label="Upload emails (.jsonl or .json)", file_types=[".jsonl", ".json"])
+
+    run_btn = gr.Button("Process", variant="primary")
+    status = gr.Textbox(label="Status", interactive=False)
+    label_counts_df = gr.Dataframe(label="Label counts", interactive=False)
+    html_samples = gr.HTML(label="Samples")
+
+    with gr.Row():
+        search_query = gr.Textbox(label="Search emails (keywords, names, etc.)")
+        search_btn = gr.Button("Search")
+    search_results = gr.Dataframe(label="Search results", interactive=False)
+
+    state_df = gr.State()
+    state_embs = gr.State()
+    state_model = gr.State()
+    state_search = gr.State()
+
+    def process_file(inbox_file):
+        if inbox_file is None:
+            return "Please upload a file", None, None, None, None, None, None
+        local_path = inbox_file.name
+        recs = []
+        if local_path.endswith(".jsonl"):
+            with open(local_path, "r", encoding="utf-8") as fh:
+                for line in fh:
+                    try:
+                        recs.append(json.loads(line))
+                    except:
+                        continue
+        else:
+            with open(local_path, "r", encoding="utf-8") as fh:
+                obj = json.load(fh)
+                if isinstance(obj, list):
+                    recs = obj
+                elif isinstance(obj, dict):
+                    recs = [obj]
+        normd = [normalize_email_record(r) for r in recs]
+        df = pd.DataFrame(normd)
+        df = df.drop_duplicates(subset=["message_id", "subject", "text_hash"])
+        texts = (df["subject"].fillna("") + "\n\n" + df["body_text"].fillna("").str.slice(0,2000)).tolist()
+        model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+        embs = embed_texts(model, texts, 512, torch_cuda_available())
+        searcher = EmailSearch(df, embs, model)
+        label_counts = df.groupby("from_domain").size().reset_index(name="count").sort_values("count", ascending=False)
+        return f"Processed {len(df)} emails", label_counts, df.head(20).to_html(), df, embs, model, searcher
+
+    run_btn.click(
+        process_file,
+        inputs=[inbox_file],
+        outputs=[status, label_counts_df, html_samples, state_df, state_embs, state_model, state_search]
+    )
+
+    def search_fn(q, df, embs, model, searcher):
+        if searcher is None:
+            return pd.DataFrame()
+        results = searcher.query(q, top_k=20)
+        return results[["date","from_email","subject","body_text","score"]]
+
+    search_btn.click(
+        search_fn,
+        inputs=[search_query, state_df, state_embs, state_model, state_search],
+        outputs=[search_results]
+    )
+
+if __name__ == "__main__":
+    demo.launch()