Blame for app.py

try:
    import spaces          
    GPU = spaces.GPU
except ImportError:
    GPU = lambda f: f       

import gradio as gr
import torch
import numpy as np
import joblib
import json
import html
import io
import os
import re
import nltk
from nltk.corpus import stopwords
from nltk.stem import SnowballStemmer
from nltk.sentiment import SentimentIntensityAnalyzer
from nltk.tokenize import sent_tokenize
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel
import shap

#nltk bits used by vader, tf-idf cleaning and the sentence splitter 
for pkg in ["stopwords", "vader_lexicon", "punkt", "punkt_tab"]:
    nltk.download(pkg, quiet=True)

HERE = os.path.dirname(os.path.abspath(__file__))
CATEGORIES = ["WORKLOAD", "TEAM", "ADJUSTMENT", "OUTSIDE_COMMS"]
NAMES = {"WORKLOAD": "Sleep and workload", "TEAM": "Team dynamics",
         "ADJUSTMENT": "Adjustment and coping", "OUTSIDE_COMMS": "Outside communication"}


# the three models 

#1. vader baseline: keyword lists + sentiment gate 
KEYWORDS = {
    "WORKLOAD": ["exhausted", "exhaustion", "tired", "fatigue", "overworked", "sleep", "schedule", "timeline",
                 "overtime", "demanding", "heavy", "restless", "shift", "tasks", "pressure"],
    "ADJUSTMENT": ["depressed", "anxious", "lonely", "monotony", "boredom", "homesick", "frustrated", "frustration",
                   "mood", "isolated", "confinement", "confined", "unmotivated", "homesickness", "sad"],
    "TEAM": ["conflict", "friction", "argument", "tension", "crew", "commander", "disagreement", "micromanagement",
             "irritated", "colleague", "team", "relationship", "clash", "dispute", "hostility"],
    "OUTSIDE_COMMS": ["family", "spouse", "children", "delay", "communication", "dropout", "satellite", "earth",
                      "ground", "conference", "disconnected", "isolated", "mail", "link", "audio"],
}
VADER_COMPOUND, VADER_NEG = -0.30, 0.02
NO_GATE = {"OUTSIDE_COMMS"}   #outside comms fires on the keyword alone
sia = SentimentIntensityAnalyzer()

#2. tf-idf + logistic regression 
vec = joblib.load(os.path.join(HERE, "tfidf_vectorizer.pkl"))
tf_models = joblib.load(os.path.join(HERE, "tfidf_logistic_models.pkl"))
STOP = set(stopwords.words("english"))
stemmer = SnowballStemmer("english")

#3. bert-base-uncased + my lora adapter

BERT_THRESHOLDS = {"WORKLOAD": 0.65, "TEAM": 0.90, "ADJUSTMENT": 0.85, "OUTSIDE_COMMS": 0.90}
tok = AutoTokenizer.from_pretrained(HERE)
base = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased", num_labels=4, problem_type="multi_label_classification")
bert = PeftModel.from_pretrained(base, HERE, torch_device="cpu").eval()   #cpu --> zerogpu pretends cuda exists at startup, eval() turns dropout off


def vader_flags(seg):
    seg = re.sub(r"\s+", " ", seg).strip()   #vader only gets whitespace cleaning (it reads caps and punctuation)
    s = sia.polarity_scores(seg)
    flags = []
    for c in CATEGORIES:
        hit = any(re.search(r"\b" + re.escape(w) + r"\b", seg.lower()) for w in KEYWORDS[c])   #whole words only
        if c in NO_GATE:
            flags.append(int(hit))
        else:
            flags.append(int(hit and (s["compound"] <= VADER_COMPOUND or s["neg"] >= VADER_NEG)))
    return flags


def clean_tfidf(text):
    #lowercase, strip punctuation, drop stopwords, stem
    text = re.sub(r"[^a-z\s]", " ", str(text).lower())
    return " ".join(stemmer.stem(w) for w in text.split() if w not in STOP and len(w) > 1)


def tfidf_probs(segs):
    x = vec.transform([clean_tfidf(s) for s in segs])
    return np.column_stack([tf_models[c].predict_proba(x)[:, 1] for c in CATEGORIES])


def bert_probs(segs):
    out = []
    for i in range(0, len(segs), 16):   #16 at a time so long transcripts don't run out of memory
        batch = tok(list(segs[i:i + 16]), truncation=True, padding=True, max_length=512, return_tensors="pt")
        with torch.no_grad():   #only predicting, no gradients needed
            out.append(torch.sigmoid(bert(**batch).logits).numpy())   #sigmoid not softmax, each category is its own yes/no
    return np.vstack(out)


#the text pipeline 

#five cleaning rules in a fixed order
WRAPPED_RE = re.compile(r"\n([A-Z])\n([A-Z][A-Za-z\-]{1,30}:)")
HEADER_RE = re.compile(r"^.*Johnson Space Center Oral History Project.*$", re.MULTILINE | re.IGNORECASE)
FOOTER_RE = re.compile(r"^\s*\d{1,2}\s+[A-Za-z]+\s+\d{4}\s+\d+[-\u2013]?\d*\s*$", re.MULTILINE)
FOOTER_RE2 = re.compile(r"^\s*[A-Za-z]+\.?\s+\d{1,2},?\s+\d{4}\s+\d+[-\u2013]?\d*\s*$", re.MULTILINE)
PAGE_RE = re.compile(r"\bPage\s+\d+\b", re.IGNORECASE)
BRACKET_RE = re.compile(r"\[.*?\]")


def clean_text(text):
    text = text.replace("\r\n", "\n")
    counts = {"wrapped speaker tags": len(WRAPPED_RE.findall(text)), "headers": len(HEADER_RE.findall(text)),
              "footers": len(FOOTER_RE.findall(text)) + len(FOOTER_RE2.findall(text)),
              "page numbers": len(PAGE_RE.findall(text)), "[brackets]": len(BRACKET_RE.findall(text))}
    text = WRAPPED_RE.sub(r"\n\1\2", text)   #has to go first or broken tags are never seen as speakers
    for rule in [HEADER_RE, FOOTER_RE, FOOTER_RE2, PAGE_RE, BRACKET_RE]:
        text = rule.sub("", text)
    return re.sub(r"[ \t]{2,}", " ", text), counts


#interviewer removal: name read from the cover page (first 1500 characters)
SPEAKER_RE = re.compile(r"\n([A-Z][A-Z\s\-\.\'\"]{1,40}):\s")
INTERVIEWER_RE = re.compile(r"Interview(?:ed|er)?\s*(?:conducted\s+)?(?:by|:)\s*([^\n]+)", re.IGNORECASE)
ASSISTED_RE = re.compile(r"assisted\s+by\s+([^\n.]+)", re.IGNORECASE)


def surname(name):
    t = [w for w in re.sub(r"[^A-Za-z\-]", " ", name).split()
         if w.upper() not in {"DR", "MR", "MRS", "MS", "PROF", "JR", "SR", "II", "III", "IV"}]
    return t[-1].upper() if t else ""


def diarise(text):
    header = text[:1500]
    lines = ([INTERVIEWER_RE.search(header).group(1)] if INTERVIEWER_RE.search(header) else [])
    lines += [m.group(1) for m in ASSISTED_RE.finditer(header)]   #co-interviewers named in another sentence
    interviewers = {surname(n) for l in lines for n in re.split(r"\s+AND\s+|&|,", l, flags=re.I) if surname(n)}

    parts = SPEAKER_RE.split("\n" + text)
    turns = [(re.sub(r"[^A-Z\- ]", "", parts[i].upper()).strip(), parts[i + 1].strip())
             for i in range(1, len(parts) - 1, 2)]
    turns = [(s, c) for s, c in turns if s and c]
    speakers = sorted({s for s, _ in turns})

    if not turns:   #no "NAME:" tags --> keep everything
        return text.strip(), "no_tags", [], speakers, 0, 0
    if interviewers:
        kept = [c for s, c in turns if s not in interviewers]
        return "\n\n".join(kept), "cover_page", sorted(interviewers), speakers, len(kept), len(turns)
    chars = {}
    for s, c in turns:
        chars[s] = chars.get(s, 0) + len(c)
    astro = max(chars, key=chars.get)   #fallback --> keep the speaker who talks the most
    kept = [c for s, c in turns if s == astro]
    return "\n\n".join(kept), "fallback", [s for s in speakers if s != astro], speakers, len(kept), len(turns)


#segmentation: whole sentences packed up to 300 words, nothing under 20 words
MIN_WORDS, MAX_WORDS = 20, 300


def segment(text):
    blocks = [b.strip() for b in text.split("\n\n") if b.strip()]
    paras = blocks[:1]
    for b in blocks[1:]:   #a blank line is only a real break if the sentence ended and the next starts with a capital
        if paras[-1][-1] in ".!?\"\u201d" and b[0].isupper():
            paras.append(b)
        else:
            paras[-1] += " " + b
    chunks = []
    for p in paras:
        if len(p.split()) < MIN_WORDS:
            continue
        if len(p.split()) <= MAX_WORDS:
            chunks.append(p)
            continue
        sents = []
        for s in sent_tokenize(p):   #glue back "sentences" that start lowercase (punkt splits on initials)
            if sents and s[:1].islower():
                sents[-1] += " " + s
            else:
                sents.append(s)
        cur, n = [], 0
        for s in sents:
            w = len(s.split())
            if n + w > MAX_WORDS and cur:
                chunks.append(" ".join(cur))
                cur, n = [], 0
            cur.append(s)
            n += w
        if cur:
            chunks.append(" ".join(cur))
    return [c for c in chunks if len(c.split()) >= MIN_WORDS]


#running everything 

def read_input(file, text):
    if file is None:
        return text or ""
    path = file if isinstance(file, str) else file.name
    if path.lower().endswith(".pdf"):
        import pdfplumber
        with pdfplumber.open(path) as pdf:   #page by page, same as layer 2
            return "\n\n".join(p.extract_text() or "" for p in pdf.pages)
    return open(path, encoding="utf-8", errors="replace").read()


@GPU
def analyse(file, text):
    raw = read_input(file, text)
    empty = gr.update(choices=[], value=None)
    if not raw.strip():
        return "Nothing to analyse, paste some text or upload a .txt / .pdf.", "", "", "", "", None, empty

    #step 1 + 2
    s1 = f"{len(raw.split()):,} words, {len(raw):,} characters"
    cleaned, counts = clean_text(raw)
    s2 = " | ".join(f"{k}: {v}" for k, v in counts.items())

    #step 3
    narrative, method, removed, speakers, kept, total = diarise(cleaned)
    if method == "cover_page":
        s3 = f"PASSED: interviewer read from the cover page ({', '.join(removed)}). Astronaut turns kept: {kept} of {total}."
    elif method == "fallback":
        s3 = f"PASSED (fallback): no 'Interviewed by' line, kept the longest speaker. Removed: {', '.join(removed)}."
    else:
        s3 = "NO SPEAKER TAGS FOUND (like ARMSTRONG:), so the whole text is treated as the astronaut's words."
    if method != "no_tags" and not narrative.strip():
        #known failure from my report: a curly apostrophe or accent in the astronaut's tag (O’CONNOR) isn't matched
        s3 = (f"FAILED: interviewer found ({', '.join(removed)}) but nothing was left after removing them, "
              f"because the astronaut's name tag was not recognised.")
    s3 += f" Words left: {len(narrative.split()):,}."
    if method == "no_tags":
        #no cover page and no NAME: tags --> probably not an oral history, the models still run but read with care
        s3 += " Note: this does not look like a NASA oral history transcript, so the scores below are less meaningful."

    #step 4
    segs = segment(narrative)
    note = ""
    if not segs and len(narrative.split()) >= 5:
        segs = [re.sub(r"\s+", " ", narrative).strip()]   #short text: analysed as one segment for the demo
        note = f" (shorter than the {MIN_WORDS} word floor used in training, so it is analysed as one segment anyway)"
    if not segs:
        return s1, s2, s3, "<h4>4. Segmenting</h4>FAILED: no text left to segment.", "", None, empty
    segs = segs[:300]   #keeps a free cpu quick
    wc = [len(s.split()) for s in segs]
    long = sum(w > MAX_WORDS for w in wc)
    if long:
        #one "sentence" over 300 words (tables, references, no full stops) can't be split, bert only reads the first 512 tokens
        note += (f"<br><span class='small'>{long} segment{'s are' if long != 1 else ' is'} over {MAX_WORDS} words because the text "
                 f"has very long runs with no full stop, BERT only reads the first 512 tokens of those.</span>")
    s4 = (f"<h4>4. Segmenting</h4><b>{len(segs)} segment{"s" if len(segs) != 1 else ""}</b>, {min(wc)} to {max(wc)} words each{note}<br><br>" +
          "".join(f"<div class='seg'><b>#{i + 1}</b> ({wc[i]} words) {html.escape(s[:220])}{'...' if len(s) > 220 else ''}</div>"
                  for i, s in enumerate(segs[:3])) +
          (f"<i>...and {len(segs) - 3} more</i>" if len(segs) > 3 else ""))

    #step 5: all three models
    v = np.array([vader_flags(s) for s in segs])
    t = tfidf_probs(segs)
    b = bert_probs(segs)
    rows = ""
    for i, c in enumerate(CATEGORIES):
        rows += (f"<tr><td>{NAMES[c]}</td><td>{int(v[:, i].sum())}</td><td>{int((t[:, i] >= 0.5).sum())}</td>"
                 f"<td><b>{int((b[:, i] >= BERT_THRESHOLDS[c]).sum())}</b></td><td>{b[:, i].max():.2f} / {BERT_THRESHOLDS[c]:.2f}</td></tr>")
    s5 = (f"<h4>5. Classification: all three models</h4><table><tr><th>Category</th><th>VADER</th><th>TF-IDF</th><th>BERT + LoRA</th>"
          f"<th>BERT top prob / threshold</th></tr>{rows}</table>"
          f"<p class='small'>Numbers are segments flagged out of {len(segs)}. BERT is my final model. "
          f"A research prototype, it flags passages for a person to read and does not diagnose anyone.</p>")

    order = np.argsort(-b.max(axis=1))   #strongest segments first for shap
    choices = [f"#{k + 1}  (BERT max {b[k].max():.2f})  {segs[k][:60]}..." for k in order]
    return s1, s2, s3, s4, s5, {"segs": segs, "bert": b.tolist()}, gr.update(choices=choices, value=choices[0])


explainer = shap.Explainer(lambda x: bert_probs([str(s) for s in x]), shap.maskers.Text(tok), output_names=CATEGORIES)


@GPU
def explain(state, choice, category):
    if not state or not choice:
        return "Run the analysis first."
    k = int(choice.split()[0][1:]) - 1
    seg = state["segs"][k]
    sv = explainer([seg], max_evals=150, silent=True)   #shap hides groups of words and watches the probability move
    i = CATEGORIES.index(category)
    words, vals = sv.data[0], np.asarray(sv.values[0])[:, i]
    top = max(np.abs(vals).max(), 1e-9)
    spans = "".join(
        f"<span style='background:{'rgba(248,113,113,' if val > 0 else 'rgba(96,165,250,'}{min(abs(val) / top, 1) * 0.8:.2f});"
        f"border-radius:3px'>{html.escape(str(w))}</span>" for w, val in zip(words, vals))
    p = state["bert"][k][i]
    return (f"Segment #{k + 1}, {NAMES[category]}: BERT probability <b>{p:.2f}</b> (threshold {BERT_THRESHOLDS[category]:.2f}). "
            f"<span style='color:#f87171'>Red</span> words push it up, <span style='color:#60a5fa'>blue</span> words push it down."
            f"<div class='seg' style='line-height:1.9'>{spans}</div>")


#files 
TEST_SET = [("ridesk-12-6-02.txt", "Sally Ride, Shuttle"), ("crippenrl-10-16-12.txt", "Robert Crippen, Shuttle"),
            ("whitsonpa-9-6-12.txt", "Peggy Whitson, ISS"), ("pettitdr-8-17-15.pdf", "Donald Pettit, ISS (PDF)"),
            ("harttj-4-10-03.txt", "Terry Hart, Shuttle"), ("rossjl-12-4-03.txt", "Jerry Ross, Shuttle"),
            ("mullanerm-1-24-03.txt", "Mike Mullane, Shuttle"), ("armstrongna-9-19-01.txt", "Neil Armstrong, Apollo")]
PIPELINE_TESTS = [("1_WORKS_cover_page_transcript.txt", "Works: transcript with cover page"),
                  ("2_TEST_one_short_sentence.txt", "Test: one short sentence"),
                  ("3_TEST_speaker_names_not_in_capitals.txt", "Test: speaker names not in capitals"),
                  ("4_TEST_curly_apostrophe_in_name.txt", "Test: curly apostrophe in a name"),
                  ("5_TEST_technical_no_strain.txt", "Test: technical, no strain")]


def find_files(pairs):
    found = []
    for name, label in pairs:
        for folder in [os.path.join(HERE, "test_set"), HERE]:
            if os.path.exists(os.path.join(folder, name)):
                found.append((os.path.join(folder, name), label))
                break
    return found


#  the page 

CSS = """
.gradio-container {background: radial-gradient(ellipse at top, #2a0f1a 0%, #05070f 70%) !important;}
* {box-shadow: none !important;}
.seg {border: 1px solid #ffffff44; border-radius: 12px; padding: 8px 12px; margin: 6px 0; color: #e6ebff;}
table {width: 100%; border-collapse: collapse; color: #e6ebff;}
th, td {border-bottom: 1px solid #ffffff33; padding: 6px 8px; text-align: left;}
.small {color: #9aa4c4; font-size: 13px;}
"""

with gr.Blocks(title="SPACEMIND") as demo:
    gr.Markdown("# SPACEMIND")

    with gr.Row():
        #inputs
        with gr.Column():
            textInput = gr.Textbox(label="Paste a transcript or a passage", lines=10)
            fileInput = gr.File(label="...or upload a .txt / .pdf transcript", file_types=[".txt", ".pdf"])
            with gr.Row():
                clearButton = gr.ClearButton()
                runButton = gr.Button("Run the pipeline", variant="primary")

            #held out test astronauts --> the model never saw them in training
            heldOut = find_files(TEST_SET)
            if heldOut:
                gr.Examples([[p] for p, _ in heldOut], inputs=fileInput, example_labels=[l for _, l in heldOut],
                            label="Load a transcript from my held-out test set")
            tests = find_files(PIPELINE_TESTS)
            if tests:
                gr.Examples([[p] for p, _ in tests], inputs=fileInput, example_labels=[l for _, l in tests],
                            label="Pipeline tests (where it works and where it breaks)")

        #the stages
        with gr.Column():
            out1 = gr.Textbox(label="1. Input (word count)")
            out2 = gr.Textbox(label="2. Cleaning")
            out3 = gr.Textbox(label="3. Removing the interviewer", lines=3)
            out4 = gr.HTML(label="4. Segmenting")
            out5 = gr.HTML(label="5. Classification")

    state = gr.State()
    gr.Markdown("### 6. SHAP: why did BERT score this segment?")
    with gr.Row():
        segPick = gr.Dropdown(label="Segment (highest BERT score first)", choices=[])
        catPick = gr.Radio(CATEGORIES, value="WORKLOAD", label="Category")
    shapButton = gr.Button("Explain (takes about 30 seconds on a free CPU)")
    shapOut = gr.HTML()

    clearButton.add([textInput, fileInput, out1, out2, out3, out4, out5, shapOut])
    runButton.click(analyse, [fileInput, textInput], [out1, out2, out3, out4, out5, state, segPick])
    shapButton.click(explain, [state, segPick, catPick], shapOut)


BURGUNDY = gr.themes.Color(c50="#fdf2f4", c100="#fbe5e9", c200="#f5c6d0", c300="#ec9aab", c400="#d9637c",
                           c500="#9b2743", c600="#801f38", c700="#661a2e", c800="#4f1524", c900="#3d1019",
                           c950="#2e0a14", name="burgundy")

demo.launch(theme=gr.themes.Base(primary_hue=BURGUNDY, neutral_hue="slate"), css=CSS,
            js="() => document.body.classList.add('dark')")