Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,436 +1,216 @@
 # app.py
-import html
-import json
-import re
-from dataclasses import asdict, is_dataclass
-from typing import Any, Dict, List, Optional, Tuple
-
-import numpy as np
-import streamlit as st
 import torch
-from transformers import AutoModel, AutoTokenizer
-
-MODEL_ID = "zilliz/semantic-highlight-bilingual-v1"
+import streamlit as st
+from transformers import AutoModel
 
-st.set_page_config(page_title="Semantic Highlight (Token-level)", layout="wide")
+st.set_page_config(page_title="Semantic Highlight Bilingual Demo", layout="wide")
 
 
 @st.cache_resource(show_spinner=True)
-def load_model_and_tokenizer():
-    # Tokenizer is usually fine as-is
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-    # HF Spaces sometimes lacks safetensors or has a transformers/safetensors mismatch.
-    # Force PyTorch loading to avoid metadata=None crashes inside transformers.
+def load_model():
     model = AutoModel.from_pretrained(
-        MODEL_ID,
+        "zilliz/semantic-highlight-bilingual-v1",
         trust_remote_code=True,
-        use_safetensors=False,
     )
-    model.eval()
-    return model, tokenizer
-
-
-def to_jsonable(x: Any) -> Any:
-    if x is None or isinstance(x, (str, int, float, bool)):
-        return x
-    if is_dataclass(x):
-        return asdict(x)
-    if isinstance(x, dict):
-        return {str(k): to_jsonable(v) for k, v in x.items()}
-    if isinstance(x, (list, tuple)):
-        return [to_jsonable(v) for v in x]
-    if isinstance(x, np.ndarray):
-        return x.tolist()
-    if torch.is_tensor(x):
-        return x.detach().cpu().tolist()
-    return str(x)
-
-
-def detect_lang(question: str, context: str) -> str:
-    s = (question or "") + " " + (context or "")
-    for ch in s:
-        if "\u4e00" <= ch <= "\u9fff":
-            return "zh"
-    return "en"
-
-
-def split_sentences_with_spans(text: str, lang: str) -> List[Tuple[int, int]]:
-    text = text or ""
-    if not text.strip():
-        return []
+    return model
 
-    spans: List[Tuple[int, int]] = []
-    if lang == "zh":
-        pattern = re.compile(r"[^。！？]*[。！？]?")
-    else:
-        pattern = re.compile(r"[^.!?]*[.!?]?")
 
-    for m in pattern.finditer(text):
-        s, e = m.span()
-        if s == e:
-            continue
-        chunk = text[s:e]
-        if chunk.strip() == "":
-            continue
-        spans.append((s, e))
-
-    if not spans:
-        spans = [(0, len(text))]
-    return spans
-
-
-def build_context_mask(tokenizer, enc) -> List[bool]:
-    ids = enc["input_ids"][0].tolist()
-    L = len(ids)
-
-    seq_ids: Optional[List[Optional[int]]] = None
-    try:
-        if hasattr(enc, "encodings") and enc.encodings and hasattr(enc.encodings[0], "sequence_ids"):
-            seq_ids = enc.encodings[0].sequence_ids
-    except Exception:
-        seq_ids = None
-
-    if seq_ids and len(seq_ids) == L:
-        return [(sid == 1) for sid in seq_ids]
-
-    sep_id = tokenizer.sep_token_id
-    context_mask = [False] * L
-    try:
-        first_sep = ids.index(sep_id)
-        for i in range(first_sep + 1, L):
-            if ids[i] == sep_id:
-                break
-            context_mask[i] = True
-        return context_mask
-    except Exception:
-        special = set(tokenizer.all_special_ids)
-        return [tok_id not in special for tok_id in ids]
-
-
-def infer_token_scores(
-    model,
-    tokenizer,
-    question: str,
-    context: str,
-    device: torch.device,
-) -> Dict[str, Any]:
-    enc = tokenizer(
-        question,
-        context,
-        return_tensors="pt",
-        return_offsets_mapping=True,
-        truncation=True,
-    )
+def split_sentences(text: str):
+    text = text.strip()
+    if not text:
+        return []
 
-    if "offset_mapping" not in enc:
-        raise RuntimeError("Tokenizer did not return offset_mapping; cannot build token-level spans.")
-
-    input_ids = enc["input_ids"].to(device)
-    attention_mask = enc.get("attention_mask", None)
-    if attention_mask is not None:
-        attention_mask = attention_mask.to(device)
-
-    offset_mapping = enc["offset_mapping"][0].tolist()
-    ids_list = enc["input_ids"][0].tolist()
-    token_texts_all = tokenizer.convert_ids_to_tokens(ids_list)
-
-    context_mask = build_context_mask(tokenizer, enc)
-
-    with torch.no_grad():
-        out = model(input_ids=input_ids, attention_mask=attention_mask)
-
-    debug_keys = []
-    if hasattr(out, "keys"):
-        try:
-            debug_keys = list(out.keys())
-        except Exception:
-            debug_keys = []
-
-    logits = None
-    if hasattr(out, "logits"):
-        logits = out.logits
-    elif isinstance(out, (tuple, list)) and len(out) > 0 and torch.is_tensor(out[0]):
-        logits = out[0]
-    elif isinstance(out, dict) and "logits" in out and torch.is_tensor(out["logits"]):
-        logits = out["logits"]
-
-    if logits is None or not torch.is_tensor(logits):
-        raise RuntimeError(f"Could not find logits in model output. Output keys: {debug_keys or type(out)}")
-
-    if logits.dim() == 3:
-        if logits.size(-1) == 1:
-            token_logits = logits[:, :, 0]
-        elif logits.size(-1) == 2:
-            token_logits = logits[:, :, 1]
-        else:
-            token_logits = logits.max(dim=-1).values
-    elif logits.dim() == 2:
-        token_logits = logits
+    # Very simple heuristic: use Chinese period if present, else English period.
+    if "。" in text:
+        parts = [s.strip() for s in text.split("。") if s.strip()]
+        # Add back "。" to each sentence for nicer display.
+        sentences = [s + "。" for s in parts]
     else:
-        raise RuntimeError(f"Unexpected logits dim: {logits.dim()}")
+        parts = [s.strip() for s in text.split(".") if s.strip()]
+        sentences = [s + "." for s in parts]
 
-    token_probs = torch.sigmoid(token_logits)[0].detach().cpu().tolist()
+    return sentences
 
-    L_ids = len(ids_list)
-    L_offsets = len(offset_mapping)
-    L_mask = len(context_mask)
-    L_probs = len(token_probs)
-    L_texts = len(token_texts_all)
-    L = min(L_ids, L_offsets, L_mask, L_probs, L_texts)
 
-    context_token_scores: List[float] = []
-    context_token_offsets: List[Tuple[int, int]] = []
-    context_token_texts: List[str] = []
+def highlight_context(context: str, highlighted_sentences):
+    if not context or not highlighted_sentences:
+        return context
 
-    for i in range(L):
-        if not context_mask[i]:
-            continue
-        start, end = offset_mapping[i]
-        if start == 0 and end == 0:
+    # Simple HTML highlighting by sentence replacement
+    highlighted_html = context
+    for sent in highlighted_sentences:
+        sent_clean = sent.strip()
+        if not sent_clean:
             continue
-        if end <= start:
-            continue
-        context_token_scores.append(float(token_probs[i]))
-        context_token_offsets.append((int(start), int(end)))
-        context_token_texts.append(str(token_texts_all[i]))
-
-    return {
-        "token_scores": context_token_scores,
-        "token_offsets": context_token_offsets,
-        "token_texts": context_token_texts,
-        "debug": {
-            "output_keys": debug_keys,
-            "logits_shape": list(logits.shape),
-            "lens": {
-                "input_ids": L_ids,
-                "offset_mapping": L_offsets,
-                "context_mask": L_mask,
-                "token_probs": L_probs,
-                "token_texts": L_texts,
-                "aligned_L_used": L,
-            },
-            "num_ctx_tokens_kept": len(context_token_scores),
-        },
-    }
-
-
-def spans_from_token_scores(
-    token_offsets: List[Tuple[int, int]],
-    token_scores: List[float],
-    threshold: float,
-    merge_gap_chars: int = 1,
-) -> List[Tuple[int, int, float]]:
-    assert len(token_offsets) == len(token_scores)
-
-    raw: List[Tuple[int, int, float]] = []
-    cur_start = None
-    cur_end = None
-    cur_scores: List[float] = []
-
-    for (s, e), score in zip(token_offsets, token_scores):
-        if score >= threshold and e > s:
-            if cur_start is None:
-                cur_start, cur_end = s, e
-                cur_scores = [score]
-            else:
-                if s <= (cur_end + merge_gap_chars):
-                    cur_end = max(cur_end, e)
-                    cur_scores.append(score)
-                else:
-                    raw.append((cur_start, cur_end, float(sum(cur_scores) / max(1, len(cur_scores)))))
-                    cur_start, cur_end = s, e
-                    cur_scores = [score]
-        else:
-            if cur_start is not None:
-                raw.append((cur_start, cur_end, float(sum(cur_scores) / max(1, len(cur_scores)))))
-                cur_start = None
-                cur_end = None
-                cur_scores = []
-
-    if cur_start is not None:
-        raw.append((cur_start, cur_end, float(sum(cur_scores) / max(1, len(cur_scores)))))
-
-    if not raw:
-        return []
-
-    raw.sort(key=lambda x: (x[0], x[1]))
-    merged: List[Tuple[int, int, List[float]]] = [(raw[0][0], raw[0][1], [raw[0][2]])]
-
-    for s, e, avg in raw[1:]:
-        ps, pe, pavgs = merged[-1]
-        if s <= pe + merge_gap_chars:
-            merged[-1] = (ps, max(pe, e), pavgs + [avg])
-        else:
-            merged.append((s, e, [avg]))
-
-    return [(s, e, float(sum(avgs) / len(avgs))) for s, e, avgs in merged]
-
-
-def render_highlighted_html(text: str, spans: List[Tuple[int, int, float]]) -> str:
-    text = text or ""
-    if not spans:
-        return f"<div class='context-box'>{html.escape(text)}</div>"
-
-    spans = [(max(0, s), min(len(text), e), sc) for s, e, sc in spans if e > s]
-    spans.sort(key=lambda x: x[0])
-
-    pieces: List[str] = []
-    cur = 0
-    for s, e, sc in spans:
-        if s > cur:
-            pieces.append(html.escape(text[cur:s]))
-        frag = html.escape(text[s:e])
-        pieces.append(f"<mark class='hl' title='score={sc:.3f}'>{frag}</mark>")
-        cur = e
-    if cur < len(text):
-        pieces.append(html.escape(text[cur:]))
+        # Avoid double-wrapping: only replace plain text, not already highlighted
+        replacement = f'<span class="hl-sentence">{sent_clean}</span>'
+        highlighted_html = highlighted_html.replace(sent_clean, replacement)
 
+    # Basic styling
     style = """
     <style>
-      .context-box {
+    .hl-sentence {
+        background-color: rgba(255, 215, 0, 0.35);
+        padding: 2px 3px;
+        border-radius: 3px;
+    }
+    .context-box {
         white-space: pre-wrap;
         font-family: ui-monospace, Menlo, Monaco, "Courier New", monospace;
-        font-size: 0.95rem;
-        line-height: 1.55;
-      }
-      mark.hl {
-        background-color: rgba(255, 215, 0, 0.35);
-        padding: 0.05em 0.15em;
-        border-radius: 0.2em;
-      }
+        font-size: 0.9rem;
+        line-height: 1.5;
+    }
     </style>
     """
-    return style + f"<div class='context-box'>{''.join(pieces)}</div>"
-
-
-def sentence_scores_from_tokens(
-    sentence_spans: List[Tuple[int, int]],
-    token_offsets: List[Tuple[int, int]],
-    token_scores: List[float],
-) -> List[float]:
-    if not sentence_spans:
-        return []
-    if not token_offsets:
-        return [0.0 for _ in sentence_spans]
-
-    mids = [((s + e) / 2.0) for s, e in token_offsets]
-    sent_scores: List[List[float]] = [[] for _ in sentence_spans]
-
-    for mid, score in zip(mids, token_scores):
-        for i, (ss, se) in enumerate(sentence_spans):
-            if ss <= mid < se:
-                sent_scores[i].append(score)
-                break
-
-    out: List[float] = []
-    for scores in sent_scores:
-        out.append(float(sum(scores) / len(scores)) if scores else 0.0)
-    return out
+    return style + f'<div class="context-box">{highlighted_html}</div>'
 
 
 def main():
-    st.title("Semantic Highlight (Token-level)")
-    st.caption(f"Model: {MODEL_ID} (bypassing `process()` to get token scores)")
+    st.title("Semantic Highlight Bilingual Demo")
+    st.caption("Model: zilliz/semantic-highlight-bilingual-v1")
 
     with st.sidebar:
         st.header("Settings")
-        lang_choice = st.selectbox("Language", ["auto", "en", "zh"], index=0)
-        token_threshold = st.slider("Token relevance threshold", 0.0, 1.0, 0.50, 0.01)
-        merge_gap = st.number_input("Merge gap (chars)", min_value=0, max_value=10, value=1, step=1)
-        show_sentence_table = st.checkbox("Show sentence scores table", value=True)
-        debug = st.checkbox("Debug (show shapes/keys/lens)", value=True)
-
-    default_q = "What are the key claims about Dan Petrovic?"
-    default_ctx = (
-        "Dan Petrovic is the managing director of DEJAN a well-known Australian AI SEO with worldwide recognition "
-        "for industry-defining innovation and thought leadership. He works with major global brands on strategic "
-        "product, service and brand visibility in AI search. Dan holds two Google awards for his work in discovery "
-        "of Google’s APIs exposing internal systems and algorithms."
-    )
-
-    col1, col2 = st.columns(2)
-    with col1:
-        question = st.text_input("Query / Question", value=default_q)
-        context = st.text_area("Context / Document", value=default_ctx, height=280)
-
-    with col2:
-        st.subheader("Run")
-        run = st.button("Highlight", type="primary")
-
-    if not run:
-        return
-
-    if not question.strip():
-        st.error("Query is empty.")
-        return
-    if not context.strip():
-        st.error("Context is empty.")
-        return
-
-    lang = detect_lang(question, context) if lang_choice == "auto" else lang_choice
-
-    try:
-        model, tokenizer = load_model_and_tokenizer()
-    except Exception as e:
-        st.error("Model failed to load in this environment.")
-        st.code(str(e), language="text")
-        st.info("On HF Spaces: add safetensors to requirements.txt and pin transformers>=4.41.0. This app also forces use_safetensors=False.")
-        return
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-
-    with st.spinner("Scoring tokens..."):
-        token_pack = infer_token_scores(
-            model=model,
-            tokenizer=tokenizer,
-            question=question,
-            context=context,
-            device=device,
+        threshold = st.slider(
+            "Relevance threshold",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.5,
+            step=0.01,
+            help="Lower values highlight more sentences; higher values highlight fewer.",
+        )
+        language = st.selectbox(
+            "Language",
+            options=["auto", "en", "zh"],
+            index=0,
+            help="Let the model auto-detect, or force English (en) / Chinese (zh).",
+        )
+        return_sentence_metrics = st.checkbox(
+            "Return per-sentence probabilities",
+            value=True,
         )
 
-    token_scores = token_pack["token_scores"]
-    token_offsets = token_pack["token_offsets"]
+        st.markdown("---")
+        st.info(
+            "1. Enter a query.\n"
+            "2. Paste a document as context.\n"
+            "3. Click **Run Semantic Highlight**."
+        )
 
-    spans = spans_from_token_scores(
-        token_offsets=token_offsets,
-        token_scores=token_scores,
-        threshold=float(token_threshold),
-        merge_gap_chars=int(merge_gap),
+    default_question = "What are the symptoms of dehydration?"
+    default_context = (
+        "Dehydration occurs when your body loses more fluid than you take in.\n"
+        "Common signs include feeling thirsty and having a dry mouth.\n"
+        "The human body is composed of about 60% water.\n"
+        "Dark yellow urine and infrequent urination are warning signs.\n"
+        "Water is essential for many bodily functions.\n"
+        "Dizziness, fatigue, and headaches can indicate severe dehydration.\n"
+        "Drinking enough water daily is often recommended."
     )
 
-    st.subheader("Context with token-level highlights")
-    st.markdown(render_highlighted_html(context, spans), unsafe_allow_html=True)
-
-    st.subheader("Metrics")
-    m1, m2, m3 = st.columns(3)
-    with m1:
-        st.metric("Highlighted spans", value=len(spans))
-    with m2:
-        st.metric("Context tokens kept", value=len(token_scores))
-    with m3:
-        st.metric("Language", value=lang)
-
-    if show_sentence_table:
-        st.subheader("Sentence scores (derived from token scores)")
-        sentence_spans = split_sentences_with_spans(context, lang=lang)
-        sent_scores = sentence_scores_from_tokens(sentence_spans, token_offsets, token_scores)
-
-        rows = []
-        for i, ((s, e), sc) in enumerate(zip(sentence_spans, sent_scores), start=1):
-            rows.append(
-                {
-                    "Sentence #": i,
-                    "Score": float(sc),
-                    "Sentence": (context[s:e] or "").strip(),
-                }
-            )
-        st.dataframe(rows, use_container_width=True)
+    col_left, col_right = st.columns(2)
 
-    if debug:
-        st.subheader("Debug")
-        st.code(json.dumps(to_jsonable(token_pack["debug"]), indent=2, ensure_ascii=False), language="json")
+    with col_left:
+        question = st.text_input(
+            "Query / Question",
+            value=default_question,
+        )
+        context = st.text_area(
+            "Context / Document",
+            value=default_context,
+            height=260,
+        )
+
+    with col_right:
+        st.subheader("Controls")
+        run = st.button("Run Semantic Highlight", type="primary")
+
+        if run:
+            if not question.strip():
+                st.error("Please enter a query/question.")
+                return
+            if not context.strip():
+                st.error("Please enter some context text.")
+                return
+
+            with st.spinner("Loading model and running inference..."):
+                model = load_model()
+                kwargs = {
+                    "question": question,
+                    "context": context,
+                    "threshold": threshold,
+                    "return_sentence_metrics": return_sentence_metrics,
+                }
+                if language != "auto":
+                    kwargs["language"] = language
+
+                with torch.no_grad():
+                    result = model.process(**kwargs)
+
+            highlighted_sentences = result.get("highlighted_sentences", [])
+            compression_rate = result.get("compression_rate", None)
+            sentence_probs = result.get("sentence_probabilities", None)
+
+            st.subheader("Results")
+
+            # Metrics row
+            metric_cols = st.columns(3)
+            with metric_cols[0]:
+                st.metric(
+                    "Highlighted sentences",
+                    value=len(highlighted_sentences),
+                )
+            with metric_cols[1]:
+                if compression_rate is not None:
+                    st.metric(
+                        "Compression rate",
+                        value=f"{compression_rate * 100:.1f}%",
+                        help="Approximate percentage of text removed.",
+                    )
+            with metric_cols[2]:
+                st.metric(
+                    "Threshold used",
+                    value=f"{threshold:.2f}",
+                )
+
+            # Highlighted sentence list
+            st.markdown("### Highlighted Sentences")
+            if highlighted_sentences:
+                for i, sent in enumerate(highlighted_sentences, start=1):
+                    st.markdown(f"**{i}.** {sent}")
+            else:
+                st.write("No sentences passed the current threshold.")
+
+            # Full context with inline highlights
+            st.markdown("### Context with Highlights")
+            highlighted_html = highlight_context(context, highlighted_sentences)
+            st.markdown(highlighted_html, unsafe_allow_html=True)
+
+            # Sentence probabilities table (if available)
+            if return_sentence_metrics and sentence_probs is not None:
+                st.markdown("### Sentence Probabilities")
+
+                sentences = split_sentences(context)
+                # Align lengths if possible; otherwise just show probabilities
+                if len(sentences) == len(sentence_probs):
+                    import pandas as pd
+
+                    data = {
+                        "Sentence #": list(range(1, len(sentences) + 1)),
+                        "Sentence": sentences,
+                        "Probability": sentence_probs,
+                    }
+                    df = pd.DataFrame(data)
+                    st.dataframe(
+                        df,
+                        use_container_width=True,
+                    )
+                else:
+                    st.write(
+                        "Count of split sentences does not match model probabilities; "
+                        "showing raw probability list."
+                    )
+                    st.write(sentence_probs)
 
 
 if __name__ == "__main__":