Upload inference.py

inference.py

@@ -0,0 +1,240 @@
+import streamlit as st
+import torch
+import re
+import os
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+LABEL2ID = {"O": 0, "B-SPAN": 1, "I-SPAN": 2}
+ID2LABEL = {v: k for k, v in LABEL2ID.items()}
+
+import glob
+
+MODEL_DIRS = {
+    "CE": "./span_model_ce",
+    "Focal": "./span_model_focal",
+}
+
+def discover_checkpoints(model_dir, prefix):
+    found = {}
+    for path in sorted(glob.glob(f"{model_dir}/checkpoint-*"), key=lambda p: int(p.split("-")[-1])):
+        name = f"{prefix} / {path.split('/')[-1]}"
+        found[name] = path
+    final_path = f"{model_dir}/final"
+    if os.path.exists(final_path):
+        found[f"{prefix} / final"] = final_path
+    return found
+
+CHECKPOINTS = {}
+for prefix, model_dir in MODEL_DIRS.items():
+    CHECKPOINTS.update(discover_checkpoints(model_dir, prefix))
+if not CHECKPOINTS:
+    st.error("No checkpoints found.")
+    st.stop()
+
+
+_current_model = {"path": None, "model": None, "tokenizer": None}
+
+def load_model(checkpoint_path):
+    if _current_model["path"] == checkpoint_path:
+        return _current_model["tokenizer"], _current_model["model"]
+    # Free old model
+    if _current_model["model"] is not None:
+        del _current_model["model"]
+        del _current_model["tokenizer"]
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
+    model = AutoModelForTokenClassification.from_pretrained(checkpoint_path)
+    model.eval()
+    if torch.cuda.is_available():
+        model = model.cuda()
+    _current_model["path"] = checkpoint_path
+    _current_model["model"] = model
+    _current_model["tokenizer"] = tokenizer
+    return tokenizer, model
+
+
+def strip_md(text):
+    text = re.sub(r'\[([^\]]*)\]\([^)]*\)', r'\1', text)
+    text = re.sub(r'\*\*([^*]*)\*\*', r'\1', text)
+    text = re.sub(r'\*([^*]*)\*', r'\1', text)
+    return text
+
+
+def build_clean_to_original_map(original, cleaned):
+    """Build character mapping from cleaned text positions back to original text positions."""
+    # Align cleaned to original using simple forward matching
+    mapping = []
+    j = 0
+    for i, ch in enumerate(cleaned):
+        while j < len(original) and original[j] != ch:
+            j += 1
+        mapping.append(j)
+        j += 1
+    return mapping
+
+
+def predict_spans(tokenizer, model, title, text, threshold=0.5):
+    """Run inference and return list of (text, is_span) tuples for rendering."""
+    device = next(model.parameters()).device
+
+    # Strip markdown for model input, keep original for display
+    clean_text = strip_md(text)
+
+    # Tokenize title and cleaned text
+    title_enc = tokenizer(title, add_special_tokens=False)
+    text_enc = tokenizer(clean_text, add_special_tokens=False, return_offsets_mapping=True)
+
+    title_ids = title_enc["input_ids"]
+    text_ids = text_enc["input_ids"]
+    text_offsets = text_enc["offset_mapping"]
+
+    # Build input: [CLS] title [SEP] text [SEP]
+    input_ids = [tokenizer.cls_token_id] + title_ids + [tokenizer.sep_token_id] + text_ids + [tokenizer.sep_token_id]
+    attention_mask = [1] * len(input_ids)
+
+    # Truncate to model max length
+    max_len = tokenizer.model_max_length
+    if max_len > 10000:
+        max_len = 512
+    input_ids = input_ids[:max_len]
+    attention_mask = attention_mask[:max_len]
+
+    text_start = len(title_ids) + 2  # CLS + title + SEP
+    text_end = len(input_ids) - 1     # before final SEP
+
+    inputs = {
+        "input_ids": torch.tensor([input_ids], device=device),
+        "attention_mask": torch.tensor([attention_mask], device=device),
+    }
+
+    with torch.no_grad():
+        logits = model(**inputs).logits[0]  # (seq_len, 3)
+        probs = torch.softmax(logits, dim=-1)
+
+    # Map token probs from clean text back to original text
+    clean_to_orig = build_clean_to_original_map(text, clean_text)
+
+    char_labels = [0] * len(text)
+    char_probs = [0.0] * len(text)
+    all_char_probs = [0.0] * len(text)
+    tokens_used = min(len(text_ids), text_end - text_start)
+
+    for i in range(tokens_used):
+        tok_idx = text_start + i
+        if tok_idx >= len(probs):
+            break
+        span_prob = (probs[tok_idx][LABEL2ID["B-SPAN"]] + probs[tok_idx][LABEL2ID["I-SPAN"]]).item()
+        if i < len(text_offsets):
+            clean_start, clean_end = text_offsets[i]
+            for cc in range(clean_start, min(clean_end, len(clean_text))):
+                if cc < len(clean_to_orig):
+                    oc = clean_to_orig[cc]
+                    if oc < len(text):
+                        all_char_probs[oc] = max(all_char_probs[oc], span_prob)
+            if span_prob >= threshold:
+                for cc in range(clean_start, min(clean_end, len(clean_text))):
+                    if cc < len(clean_to_orig):
+                        oc = clean_to_orig[cc]
+                        if oc < len(text):
+                            char_labels[oc] = 1
+                            char_probs[oc] = max(char_probs[oc], span_prob)
+
+    # Expand labeled chars to cover full words (fix subword splits)
+    # A "word" is a run of non-whitespace characters
+    i = 0
+    while i < len(text):
+        if text[i].isspace():
+            i += 1
+            continue
+        # Find word boundary
+        word_start = i
+        while i < len(text) and not text[i].isspace():
+            i += 1
+        word_end = i
+        # If any char in this word is labeled, label the whole word
+        if any(char_labels[c] for c in range(word_start, word_end)):
+            max_prob = max(char_probs[c] for c in range(word_start, word_end))
+            for c in range(word_start, word_end):
+                char_labels[c] = 1
+                char_probs[c] = max(char_probs[c], max_prob)
+
+    # Build segments with average confidence per span
+    segments = []
+    if not text:
+        return segments
+
+    current_label = char_labels[0]
+    current_start = 0
+
+    for i in range(1, len(text)):
+        if char_labels[i] != current_label:
+            conf = sum(char_probs[current_start:i]) / max(1, i - current_start) if current_label == 1 else 0.0
+            segments.append((text[current_start:i], current_label == 1, conf))
+            current_start = i
+            current_label = char_labels[i]
+    conf = sum(char_probs[current_start:]) / max(1, len(text) - current_start) if current_label == 1 else 0.0
+    segments.append((text[current_start:], current_label == 1, conf))
+
+    return segments, all_char_probs
+
+
+st.set_page_config(page_title="Span Extractor", layout="wide")
+st.title("Span Extractor Inference")
+
+checkpoint_names = list(CHECKPOINTS.keys())
+checkpoint = st.selectbox("Checkpoint", checkpoint_names, index=len(checkpoint_names) - 1)
+tokenizer, model = load_model(CHECKPOINTS[checkpoint])
+
+threshold = st.slider("Span confidence threshold", 0.0, 1.0, 0.5, 0.05)
+
+title = st.text_input("Title", placeholder="Enter article title...")
+text = st.text_area("Text", height=300, placeholder="Enter article text...")
+
+if st.button("Extract Spans") and title and text:
+    segments, all_char_probs = predict_spans(tokenizer, model, title, text, threshold)
+
+    if not any(is_span for _, is_span, _ in segments):
+        st.warning("No spans predicted.")
+    else:
+        span_count = sum(1 for seg, is_span, _ in segments if is_span)
+        st.caption(f"{span_count} span(s) detected")
+
+        # Render with green background for spans, tooltips on all words
+        html_parts = []
+        pos = 0
+        for seg, is_span, conf in segments:
+            # Split segment into words to add per-word tooltips
+            import re as _re
+            words = _re.split(r'(\s+)', seg)
+            for word in words:
+                if not word:
+                    continue
+                escaped = word.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace("\n", "<br>")
+                # Get avg prob for this word's characters
+                word_start = pos
+                word_end = pos + len(word)
+                word_probs = all_char_probs[word_start:word_end]
+                avg_prob = sum(word_probs) / max(1, len(word_probs))
+                tooltip = f"{avg_prob:.2f}"
+                if is_span:
+                    html_parts.append(f'<span title="{tooltip}" style="background-color: #22c55e; color: white; padding: 1px 3px; border-radius: 3px; cursor: help;">{escaped}</span>')
+                else:
+                    html_parts.append(f'<span title="{tooltip}" style="cursor: help;">{escaped}</span>')
+                pos += len(word)
+
+        html = f'<div style="font-size: 16px; line-height: 1.8; font-family: Georgia, serif;">{"".join(html_parts)}</div>'
+        st.markdown(html, unsafe_allow_html=True)
+
+        # Show extracted spans as dataframe
+        st.divider()
+        st.subheader("Extracted Spans")
+        import pandas as pd
+        span_data = [{"span": seg.strip(), "confidence": conf} for seg, is_span, conf in segments if is_span]
+        df = pd.DataFrame(span_data)
+        st.dataframe(
+            df,
+            use_container_width=True,
+            hide_index=True,
+            column_config={"confidence": st.column_config.ProgressColumn(min_value=0, max_value=1, format="%.2f")},
+        )