wuenlp_tools.models.sentiment.harrymotions_prompting

  1from __future__ import annotations
  2
  3import json
  4import os
  5import re
  6import time
  7from collections import Counter
  8from concurrent.futures import ThreadPoolExecutor, as_completed
  9from dataclasses import dataclass, field
 10from enum import Enum
 11from itertools import combinations
 12from typing import Iterable
 13
 14from loguru import logger
 15from pydantic import BaseModel
 16from tqdm import tqdm
 17from wuenlp.impl.UIMANLPStructs import UIMADocument
 18from wuenlp.impl.uima import UIMACharacter, UIMACharacterReference, UIMAInteraction, UIMASentence, UIMASpan
 19
 20from wuenlp_tools.pipeline import PipelineProcessor, PipelineStep
 21from wuenlp_tools.utils.prompting import LLM, LLMArchitecture, LLMCompletion, gemini3flash
 22
 23
 24class HarryMotionsMarkerMode(str, Enum):
 25    NO_IND = "no_ind"
 26    ROLE = "role"
 27    ENTITY = "entity"
 28    MROLE = "mrole"
 29    MENTITY = "mentity"
 30
 31
 32class HarryMotionsPairingMode(str, Enum):
 33    ALL_PAIRS = "all_pairs"
 34    STAR = "star"  # top-1 character (by mention count) vs each other selected character in window
 35
 36
 37@dataclass(frozen=True)
 38class HarryMotionsPromptConfig:
 39    """HM prompting settings.
 40
 41    ``n_classes`` selects the label space: ``2``, ``5``, or ``8`` for discrete
 42    emotion labels, or ``0`` for a continuous polarity score in ``[-1, 1]``
 43    (negative → neutral → positive affect between the two characters).
 44    """
 45
 46    directed: bool = False
 47    n_classes: int = 2
 48    marker_mode: HarryMotionsMarkerMode = HarryMotionsMarkerMode.ENTITY
 49    model: LLMArchitecture = field(default_factory=lambda: gemini3flash)
 50
 51
 52@dataclass(frozen=True)
 53class HarryMotionsSamplingConfig:
 54    # Context extraction around selected mentions (kept unchanged).
 55    window_size: int = 10
 56    # Interaction extraction windows over sentences.
 57    window_sentences: int = 5
 58    overlapping_windows: bool = True
 59    top_k_characters: int = 10
 60    max_pairs: int | None = None
 61    workers: int = 8
 62    pairing_mode: HarryMotionsPairingMode = HarryMotionsPairingMode.ALL_PAIRS
 63
 64
 65class LabelResponse(BaseModel):
 66    label: str
 67
 68
 69class SentimentScalarResponse(BaseModel):
 70    """Structured LLM output for ``n_classes == 0`` (continuous polarity)."""
 71
 72    sentiment: float
 73
 74
 75MARKER_TOKEN_RE = re.compile(r"(?<!\S)(TARG|EXP)-(\S+?)-\1(?!\S)")
 76
 77_TRANSIENT_LLM_ERRORS = frozenset({
 78    "APITimeoutError",
 79    "RateLimitError",
 80    "APIConnectionError",
 81    "InternalServerError",
 82})
 83
 84
 85def _is_transient_llm_error(exc: Exception) -> bool:
 86    return type(exc).__name__ in _TRANSIENT_LLM_ERRORS
 87
 88
 89def _harrymotions_max_llm_retries() -> int:
 90    raw = os.getenv("WUENLP_HARRYMOTIONS_MAX_LLM_RETRIES", "8").strip()
 91    try:
 92        return max(1, int(raw))
 93    except ValueError:
 94        logger.warning(
 95            "Invalid WUENLP_HARRYMOTIONS_MAX_LLM_RETRIES={!r}; using 8",
 96            raw,
 97        )
 98        return 8
 99
100
101def _invoke_llm_resilient(llm: LLM, prompt: str) -> LLMCompletion:
102    max_retries = _harrymotions_max_llm_retries()
103    delay_s = 5.0
104    attempt = 0
105    last_exc: Exception | None = None
106    while attempt < max_retries:
107        attempt += 1
108        try:
109            return llm.with_reasoning(prompt)
110        except Exception as exc:
111            if not _is_transient_llm_error(exc):
112                raise
113            last_exc = exc
114            if attempt >= max_retries:
115                break
116            logger.warning(
117                "HarryMotions LLM transient error (attempt {}/{}), retry in {:.0f}s: {}: {}",
118                attempt,
119                max_retries,
120                delay_s,
121                type(exc).__name__,
122                exc,
123            )
124            time.sleep(delay_s)
125            delay_s = min(delay_s * 1.5, 120.0)
126    raise RuntimeError(
127        f"HarryMotions LLM failed after {max_retries} transient-error retries"
128    ) from last_exc
129
130
131@dataclass(frozen=True)
132class HarryMotionsPrediction:
133    label: str
134    reason: str | None = None
135
136
137def _supported_n_classes() -> tuple[int, ...]:
138    return (0, 2, 5, 8)
139
140
141def _labels_for_setting(directed: bool, n_classes: int) -> list[str]:
142    if n_classes == 0:
143        raise ValueError("Continuous mode (n_classes=0) does not use discrete labels.")
144    if n_classes not in (2, 5, 8):
145        raise ValueError(f"Unsupported class count: {n_classes}")
146    if n_classes == 8:
147        base = ["anger", "anticipation", "disgust", "fear", "joy", "sadness", "surprise", "trust"]
148    elif n_classes == 5:
149        base = ["anger", "fear", "joy", "sadness", "surprise"]
150    else:
151        base = ["neg", "pos"]
152    if directed:
153        return sorted([f"{b}_L" for b in base] + [f"{b}_R" for b in base])
154    return sorted(base)
155
156
157def _validate_sample_mode(sample: str, marker_mode: HarryMotionsMarkerMode) -> None:
158    has_role = bool(re.search(r"(?:TARG|EXP)-[\s\S]+?-(?:TARG|EXP)", sample))
159    has_ent = bool(re.search(r"ENT-[\s\S]+?-ENT", sample))
160    has_mrole = "experiencer_object" in sample or "target_object" in sample
161    has_mentity = "entity_object" in sample
162
163    if marker_mode == HarryMotionsMarkerMode.ROLE:
164        if not re.search(r"TARG-[\s\S]+?-TARG", sample) or not re.search(r"EXP-[\s\S]+?-EXP", sample):
165            raise ValueError("ROLE mode requires both TARG-...-TARG and EXP-...-EXP markers.")
166        if has_ent or has_mrole or has_mentity:
167            raise ValueError("ROLE mode sample contains incompatible marker style.")
168    elif marker_mode == HarryMotionsMarkerMode.ENTITY:
169        if not has_ent:
170            raise ValueError("ENTITY mode requires ENT-...-ENT markers.")
171        if has_role or has_mrole or has_mentity:
172            raise ValueError("ENTITY mode sample contains incompatible marker style.")
173    elif marker_mode == HarryMotionsMarkerMode.NO_IND:
174        if has_role or has_ent or has_mrole or has_mentity:
175            raise ValueError("NO_IND mode must not contain marker tokens/placeholders.")
176    elif marker_mode == HarryMotionsMarkerMode.MROLE:
177        if "experiencer_object" not in sample or "target_object" not in sample:
178            raise ValueError("MROLE mode requires experiencer_object and target_object placeholders.")
179        if has_role or has_ent or has_mentity:
180            raise ValueError("MROLE mode sample contains incompatible marker style.")
181    elif marker_mode == HarryMotionsMarkerMode.MENTITY:
182        if "entity_object" not in sample:
183            raise ValueError("MENTITY mode requires entity_object placeholder.")
184        if has_role or has_ent or has_mrole:
185            raise ValueError("MENTITY mode sample contains incompatible marker style.")
186
187
188def _apply_indicator(text: str, indicator: HarryMotionsMarkerMode) -> str:
189    def replace_tagged(in_text: str, left_tag: str, right_tag: str, repl: str) -> str:
190        return re.sub(fr"{left_tag}-([\s\S]+?)-{right_tag}", repl, in_text)
191
192    if indicator == HarryMotionsMarkerMode.NO_IND:
193        out = replace_tagged(text, "TARG", "TARG", r"\1")
194        out = replace_tagged(out, "EXP", "EXP", r"\1")
195        return out
196    if indicator == HarryMotionsMarkerMode.ROLE:
197        return text
198    if indicator == HarryMotionsMarkerMode.ENTITY:
199        out = replace_tagged(text, "TARG", "TARG", r"ENT-\1-ENT")
200        out = replace_tagged(out, "EXP", "EXP", r"ENT-\1-ENT")
201        return out
202    if indicator == HarryMotionsMarkerMode.MROLE:
203        out = re.sub(r"TARG-[\S\s]+?-TARG", "target_object", text)
204        out = re.sub(r"EXP-[\S\s]+?-EXP", "experiencer_object", out)
205        return out
206    if indicator == HarryMotionsMarkerMode.MENTITY:
207        out = re.sub(r"TARG-[\S\s]+?-TARG", "entity_object", text)
208        out = re.sub(r"EXP-[\S\s]+?-EXP", "entity_object", out)
209        return out
210    raise ValueError(f"Unknown indicator mode: {indicator}")
211
212
213def _marker_hint(indicator: HarryMotionsMarkerMode) -> str:
214    if indicator == HarryMotionsMarkerMode.ROLE:
215        return (
216            "Characters are marked in the text with their roles: TARG-Name-TARG (target) "
217            "and EXP-Name-EXP (experiencer). The character names inside the markers are preserved."
218        )
219    if indicator == HarryMotionsMarkerMode.ENTITY:
220        return (
221            "Characters are marked in the text as ENT-Name-ENT. The character names are preserved, "
222            "but role markers are replaced with a generic ENT label."
223        )
224    if indicator == HarryMotionsMarkerMode.MROLE:
225        return (
226            "Character names have been removed. Two generic placeholders appear: "
227            "'experiencer_object' and 'target_object'."
228        )
229    if indicator == HarryMotionsMarkerMode.MENTITY:
230        return (
231            "Both character names and role information have been removed. The placeholder "
232            "'entity_object' is used for every character position."
233        )
234    return ""
235
236
237def _direction_block(cfg: HarryMotionsPromptConfig) -> str:
238    if not cfg.directed:
239        hint = _marker_hint(cfg.marker_mode)
240        return (
241            "All labels are undirected: they describe the overall emotion between the two characters "
242            "without distinguishing who feels it toward whom.\n\n" + hint if hint else
243            "All labels are undirected: they describe the overall emotion between the two characters "
244            "without distinguishing who feels it toward whom."
245        )
246
247    common_lr = (
248        "- Labels ending in _L: the first-mentioned character is the experiencer. They feel the emotion toward "
249        "the second-mentioned character.\n"
250        "- Labels ending in _R: the second-mentioned character is the experiencer. They feel the emotion toward "
251        "the first-mentioned character."
252    )
253
254    if cfg.marker_mode == HarryMotionsMarkerMode.ROLE:
255        return (
256            "Labels are directional. Characters are marked in the text with role markers:\n"
257            "- EXP-Name-EXP = experiencer\n"
258            "- TARG-Name-TARG = target\n\n"
259            "- Labels ending in _L: the first-mentioned character is the experiencer (EXP appears first).\n"
260            "- Labels ending in _R: the second-mentioned character is the experiencer (EXP appears second)."
261        )
262    if cfg.marker_mode == HarryMotionsMarkerMode.ENTITY:
263        return (
264            "Labels are directional. Characters are marked in the text as ENT-Name-ENT. "
265            "The character names are preserved, but role markers are replaced with a generic ENT label.\n"
266            + common_lr
267        )
268    if cfg.marker_mode == HarryMotionsMarkerMode.MROLE:
269        return (
270            "Labels are directional. Character names have been removed and replaced by placeholders "
271            "'experiencer_object' and 'target_object'.\n"
272            + common_lr
273        )
274    if cfg.marker_mode == HarryMotionsMarkerMode.MENTITY:
275        return (
276            "Labels are directional. Both names and role markers are removed. "
277            "The placeholder 'entity_object' marks each character position.\n"
278            + common_lr
279        )
280    return "Labels are directional. Two characters appear in the text.\n" + common_lr
281
282
283def _continuous_undirected_intro(cfg: HarryMotionsPromptConfig) -> str:
284    marker = _marker_hint(cfg.marker_mode)
285    base = (
286        "Score the emotional polarity **between** the two characters on a continuous scale from -1 to 1:\n"
287        "- **-1** strongly negative affect, tension, or hostility between them\n"
288        "- **0** neutral or no clear valenced tie\n"
289        "- **+1** strongly positive warmth, trust, or goodwill between them\n"
290        "The score is **undirected**: it does not say who feels what toward whom; it summarizes the "
291        "overall relational tone in the passage."
292    )
293    return base + ("\n\n" + marker if marker else "")
294
295
296def _continuous_directed_intro(cfg: HarryMotionsPromptConfig) -> str:
297    marker = _marker_hint(cfg.marker_mode)
298    base = (
299        "Score **directed** emotional polarity: how the **first-mentioned** character appears to feel "
300        "toward or about the **second-mentioned** character, on a scale from -1 to 1:\n"
301        "- **-1** strongly negative (cold, hostile, resentful, fearful toward them)\n"
302        "- **0** neutral or unclear valence\n"
303        "- **+1** strongly positive (warm, trusting, fond, grateful toward them)\n"
304        "Use the order in which the characters appear in the marked text (first vs second)."
305    )
306    return base + ("\n\n" + marker if marker else "")
307
308
309def _build_llm_continuous(cfg: HarryMotionsPromptConfig) -> LLM:
310    dir_intro = _continuous_directed_intro(cfg) if cfg.directed else _continuous_undirected_intro(cfg)
311    system_prompt = (
312        "You are an emotion annotator for short literary excerpts. "
313        "Each input describes an interaction between two characters.\n\n"
314        f"{dir_intro}\n\n"
315        "Output format:\n"
316        "- Return a JSON object with exactly one numeric field `sentiment` (float) in [-1.0, 1.0].\n"
317        "- Do not add other fields or commentary."
318    )
319    return LLM(
320        model=cfg.model,
321        system_prompt=system_prompt,
322        cache_maxsize=0,
323        output_format=SentimentScalarResponse,
324    )
325
326
327def _build_llm(cfg: HarryMotionsPromptConfig, label_names: list[str]) -> LLM:
328    system_prompt = (
329        "You are an emotion classifier for short pieces of literary text. "
330        "Each input text contains an interaction between two characters.\n\n"
331        f"{_direction_block(cfg)}\n\n"
332        "Your task:\n"
333        "- Read the input text.\n"
334        "- Choose exactly one label from the following list that best matches the emotion expressed "
335        f"between the two characters:\n  [{', '.join(label_names)}]\n\n"
336        "Output format:\n"
337        "- Return a JSON object with a single field 'label', whose value is exactly one of the allowed labels.\n"
338        "- Do not include any additional fields or explanations."
339    )
340    return LLM(model=cfg.model, system_prompt=system_prompt, cache_maxsize=0, output_format=LabelResponse)
341
342
343def _build_fallback_llm(primary: LLM, cfg: HarryMotionsPromptConfig) -> LLM:
344    return LLM(
345        model=cfg.model,
346        system_prompt=primary.system_prompt,
347        cache_maxsize=0,
348        output_format=None,
349        json=cfg.model.provider in ("ollama", "local"),
350    )
351
352
353def _parse_label(output: LabelResponse | dict | str, label_names: list[str]) -> str:
354    allowed = {l.lower(): l for l in label_names}
355    raw = output.label if isinstance(output, LabelResponse) else output.get("label", "") if isinstance(output, dict) else str(output)
356    candidate = raw.strip()
357    # Handle responses like `H{"label":"neg"}` or any surrounding text.
358    json_like = re.search(r"\{[\s\S]*\}", candidate)
359    if json_like:
360        fragment = json_like.group(0)
361        try:
362            parsed = json.loads(fragment)
363            if isinstance(parsed, dict) and isinstance(parsed.get("label"), str):
364                candidate = parsed["label"].strip()
365        except Exception:
366            # Fallback: regex extraction from json-like fragment.
367            m = re.search(r'"label"\s*:\s*"([^"]+)"', fragment)
368            if m:
369                candidate = m.group(1).strip()
370
371    if candidate.lower() in allowed:
372        return allowed[candidate.lower()]
373    token = re.split(r"[\s,;:]+", candidate)[0].strip(" \"'`[]{}()")
374    if token.lower() in allowed:
375        return allowed[token.lower()]
376    for k, v in allowed.items():
377        if k in candidate.lower():
378            return v
379    raise ValueError(f"Could not parse a valid label from output: {raw!r}")
380
381
382def _parse_sentiment(output: SentimentScalarResponse | dict | str) -> float:
383    if isinstance(output, SentimentScalarResponse):
384        return max(-1.0, min(1.0, float(output.sentiment)))
385    if isinstance(output, dict):
386        v = output.get("sentiment", output.get("polarity", 0.0))
387        try:
388            return max(-1.0, min(1.0, float(v)))
389        except (TypeError, ValueError):
390            pass
391    candidate = str(output).strip()
392    json_like = re.search(r"\{[\s\S]*\}", candidate)
393    if json_like:
394        fragment = json_like.group(0)
395        try:
396            parsed = json.loads(fragment)
397            if isinstance(parsed, dict):
398                for key in ("sentiment", "polarity"):
399                    val = parsed.get(key)
400                    if isinstance(val, (int, float, str)):
401                        try:
402                            return max(-1.0, min(1.0, float(val)))
403                        except (TypeError, ValueError):
404                            continue
405        except Exception:
406            m = re.search(
407                r'"(?:sentiment|polarity)"\s*:\s*([-+]?\d*\.?\d+(?:[eE][-+]?\d+)?)',
408                fragment,
409            )
410            if m:
411                return max(-1.0, min(1.0, float(m.group(1))))
412    m = re.search(r"[-+]?\d*\.?\d+(?:[eE][-+]?\d+)?", candidate)
413    if m:
414        return max(-1.0, min(1.0, float(m.group(0))))
415    raise ValueError(f"Could not parse sentiment from output: {output!r}")
416
417
418def classify_harrymotions_continuous(
419    sample_text: str,
420    cfg: HarryMotionsPromptConfig,
421    *,
422    llm: LLM | None = None,
423    fallback_llm: LLM | None = None,
424) -> HarryMotionsPrediction:
425    primary = llm or _build_llm_continuous(cfg)
426    fallback = fallback_llm or _build_fallback_llm(primary, cfg)
427    try:
428        completion = _invoke_llm_resilient(primary, sample_text)
429        return HarryMotionsPrediction(
430            label=str(_parse_sentiment(completion.content)),
431            reason=completion.reasoning,
432        )
433    except Exception as exc:
434        logger.warning(
435            "Structured sentiment parsing failed; retrying with raw text output. "
436            f"Reason: {type(exc).__name__}: {exc}"
437        )
438        last_exc: Exception | None = None
439        for attempt in range(10):
440            prompt = sample_text if attempt == 0 else (
441                sample_text + '\n\nReturn only a JSON object like {"sentiment": <float between -1 and 1>}.'
442            )
443            try:
444                completion = _invoke_llm_resilient(fallback, prompt)
445                return HarryMotionsPrediction(
446                    label=str(_parse_sentiment(completion.content)),
447                    reason=completion.reasoning,
448                )
449            except Exception as inner_exc:
450                last_exc = inner_exc
451        raise RuntimeError(
452            "Could not classify HarryMotions continuous sample after retries."
453        ) from last_exc
454
455
456def classify_harrymotions_sample(
457    sample_text: str,
458    cfg: HarryMotionsPromptConfig,
459    *,
460    llm: LLM | None = None,
461    fallback_llm: LLM | None = None,
462) -> HarryMotionsPrediction:
463    if cfg.n_classes not in _supported_n_classes():
464        raise ValueError(f"Unsupported n_classes={cfg.n_classes}; expected one of {_supported_n_classes()}.")
465    _validate_sample_mode(sample_text, cfg.marker_mode)
466    if cfg.n_classes == 0:
467        return classify_harrymotions_continuous(
468            sample_text, cfg, llm=llm, fallback_llm=fallback_llm,
469        )
470    labels = _labels_for_setting(cfg.directed, cfg.n_classes)
471    primary = llm or _build_llm(cfg, labels)
472    fallback = fallback_llm or _build_fallback_llm(primary, cfg)
473    try:
474        completion = _invoke_llm_resilient(primary, sample_text)
475        return HarryMotionsPrediction(
476            label=_parse_label(completion.content, labels),
477            reason=completion.reasoning,
478        )
479    except Exception as exc:
480        logger.warning(
481            "Structured label parsing failed; retrying with raw text output. "
482            f"Reason: {type(exc).__name__}: {exc}"
483        )
484        last_exc: Exception | None = None
485        for attempt in range(10):
486            prompt = sample_text if attempt == 0 else (
487                sample_text
488                + "\n\nReturn only a JSON object like {\"label\":\"<one_allowed_label>\"}."
489            )
490            try:
491                completion = _invoke_llm_resilient(fallback, prompt)
492                return HarryMotionsPrediction(
493                    label=_parse_label(completion.content, labels),
494                    reason=completion.reasoning,
495                )
496            except Exception as inner_exc:
497                last_exc = inner_exc
498        raise RuntimeError(
499            "Could not classify HarryMotions discrete sample after retries."
500        ) from last_exc
501
502
503def _character_key(character: UIMACharacter | None) -> str | None:
504    if character is None:
505        return None
506    try:
507        return str(character.id)
508    except Exception:
509        return None
510
511
512def _ref_character_key(ref: UIMACharacterReference) -> str | None:
513    try:
514        return _character_key(ref.referred_entity)
515    except Exception:
516        return None
517
518
519def _reference_preference(ref: UIMACharacterReference) -> tuple[int, int]:
520    """
521    Lower tuple is better:
522    - prefer longest reference span in sentence
523    - tie-break by earlier mention
524    """
525    length_rank = -(int(ref.end) - int(ref.begin))
526    return length_rank, int(ref.begin)
527
528
529def _extract_window_text(doc: UIMADocument, left: UIMACharacterReference, right: UIMACharacterReference, window_size: int) -> tuple[str, int]:
530    try:
531        tokens = list(doc.tokens)
532        left_tok = int(left.token_begin_within)
533        right_tok = int(right.token_begin_within)
534        lo = max(0, min(left_tok, right_tok) - window_size)
535        hi = min(len(tokens) - 1, max(left_tok, right_tok) + window_size)
536        start = int(tokens[lo].begin)
537        end = int(tokens[hi].end)
538        return doc.text[start:end], start
539    except Exception:
540        start = max(0, min(int(left.begin), int(right.begin)) - 300)
541        end = min(len(doc.text), max(int(left.end), int(right.end)) + 300)
542        return doc.text[start:end], start
543
544
545def _insert_role_markers(sample: str, sample_offset: int, exp_ref: UIMACharacterReference, targ_ref: UIMACharacterReference) -> str:
546    edits = [
547        (int(exp_ref.begin) - sample_offset, int(exp_ref.end) - sample_offset, "EXP"),
548        (int(targ_ref.begin) - sample_offset, int(targ_ref.end) - sample_offset, "TARG"),
549    ]
550    out = sample
551    for b, e, tag in sorted(edits, key=lambda x: x[0], reverse=True):
552        if b < 0 or e > len(out) or b >= e:
553            raise ValueError("Mention boundaries are outside extracted sample window.")
554        mention = out[b:e]
555        out = out[:b] + f"{tag}-{mention}-{tag}" + out[e:]
556    return out
557
558
559def _pair_distance_tokens(a: UIMACharacterReference, b: UIMACharacterReference) -> int | None:
560    try:
561        return abs(int(a.token_begin_within) - int(b.token_begin_within))
562    except Exception:
563        return None
564
565
566def _append_interaction_pair(
567        pairs: list[tuple[UIMACharacterReference, UIMACharacterReference]],
568        left: UIMACharacterReference,
569        right: UIMACharacterReference,
570        *,
571        directed: bool,
572        max_pairs: int | None,
573        yielded: int,
574) -> tuple[list[tuple[UIMACharacterReference, UIMACharacterReference]], int, bool]:
575    if _ref_character_key(left) == _ref_character_key(right):
576        return pairs, yielded, False
577    if directed:
578        pairs.append((left, right))
579        yielded += 1
580        if max_pairs is not None and yielded >= max_pairs:
581            return pairs, yielded, True
582        pairs.append((right, left))
583        yielded += 1
584        if max_pairs is not None and yielded >= max_pairs:
585            return pairs, yielded, True
586    else:
587        pairs.append((left, right))
588        yielded += 1
589        if max_pairs is not None and yielded >= max_pairs:
590            return pairs, yielded, True
591    return pairs, yielded, False
592
593
594def _iter_interaction_pairs(
595        doc: UIMADocument,
596        sampling: HarryMotionsSamplingConfig,
597        directed: bool,
598        selected_characters: list[UIMACharacter] | None = None,
599) -> Iterable[tuple[UIMACharacterReference, UIMACharacterReference]]:
600    refs = [r for r in doc.character_references if _ref_character_key(r) is not None]
601    if not refs:
602        return []
603
604    counts = Counter(_ref_character_key(r) for r in refs)
605    counts.pop(None, None)
606    if selected_characters:
607        selected = {_character_key(c) for c in selected_characters}
608        selected.discard(None)
609    else:
610        selected = {key for key, _ in sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))[:sampling.top_k_characters]}
611
612    ranked_selected = sorted(selected, key=lambda k: (-counts.get(k, 0), k))
613    hub_key = ranked_selected[0] if ranked_selected else None
614
615    sentences = list(doc.sentences)
616    if not sentences:
617        return []
618
619    window_sentences = max(1, int(sampling.window_sentences))
620    step = 1 if sampling.overlapping_windows else window_sentences
621
622    yielded = 0
623    pairs: list[tuple[UIMACharacterReference, UIMACharacterReference]] = []
624    # Generate sentence windows; if text is shorter than window, use one full window.
625    starts = [0] if len(sentences) <= window_sentences else list(range(0, len(sentences) - window_sentences + 1, step))
626    for start_idx in starts:
627        end_idx = min(len(sentences), start_idx + window_sentences)
628        window_refs: list[UIMACharacterReference] = []
629        for sentence in sentences[start_idx:end_idx]:
630            window_refs.extend([r for r in sentence.covered(UIMACharacterReference) if _ref_character_key(r) in selected])
631
632        representative_ref: dict[str, UIMACharacterReference] = {}
633        for ref in window_refs:
634            cid = _ref_character_key(ref)
635            if cid is None:
636                continue
637            if cid not in representative_ref or _reference_preference(ref) < _reference_preference(representative_ref[cid]):
638                representative_ref[cid] = ref
639
640        if sampling.pairing_mode == HarryMotionsPairingMode.STAR:
641            if hub_key is None or hub_key not in representative_ref:
642                continue
643            hub_ref = representative_ref[hub_key]
644            for cid, other_ref in representative_ref.items():
645                if cid == hub_key:
646                    continue
647                pairs, yielded, stop = _append_interaction_pair(
648                    pairs, hub_ref, other_ref, directed=directed, max_pairs=sampling.max_pairs, yielded=yielded,
649                )
650                if stop:
651                    return pairs
652        else:
653            ordered_refs = sorted(representative_ref.values(), key=lambda r: int(r.begin))
654            for left, right in combinations(ordered_refs, 2):
655                pairs, yielded, stop = _append_interaction_pair(
656                    pairs, left, right, directed=directed, max_pairs=sampling.max_pairs, yielded=yielded,
657                )
658                if stop:
659                    return pairs
660    return pairs
661
662
663def _sentiment_for_2class(label: str) -> float | None:
664    if label.startswith("pos"):
665        return 1.0
666    if label.startswith("neg"):
667        return -1.0
668    return None
669
670
671def _setting_id(cfg: HarryMotionsPromptConfig) -> str:
672    return f"{cfg.n_classes}cl_{'dir' if cfg.directed else 'undir'}_{cfg.marker_mode.value}"
673
674
675def _setting_feature(cfg: HarryMotionsPromptConfig, suffix: str) -> str:
676    return f"harrymotions_{_setting_id(cfg)}_{suffix}"
677
678
679@dataclass
680class HarryMotionsPromptingProcessor(PipelineProcessor):
681    prompt_config: HarryMotionsPromptConfig = field(default_factory=lambda: HarryMotionsPromptConfig(False, 2, HarryMotionsMarkerMode.ENTITY))
682    sampling_config: HarryMotionsSamplingConfig = field(default_factory=HarryMotionsSamplingConfig)
683    _llm: LLM | None = field(default=None, init=False, repr=False)
684    _fallback_llm: LLM | None = field(default=None, init=False, repr=False)
685
686    def _ensure_llms(self) -> tuple[LLM, LLM]:
687        if self._llm is not None and self._fallback_llm is not None:
688            return self._llm, self._fallback_llm
689        if self.prompt_config.n_classes == 0:
690            self._llm = _build_llm_continuous(self.prompt_config)
691        else:
692            labels = _labels_for_setting(self.prompt_config.directed, self.prompt_config.n_classes)
693            self._llm = _build_llm(self.prompt_config, labels)
694        self._fallback_llm = _build_fallback_llm(self._llm, self.prompt_config)
695        return self._llm, self._fallback_llm
696
697    def __call__(self, doc: UIMADocument, unit_type: type[UIMASpan] | None, overwrite: bool = False, **kwargs) -> UIMADocument:
698        label_key = _setting_feature(self.prompt_config, "label")
699        sentiment_polarity_key = _setting_feature(self.prompt_config, "sentiment_polarity")
700        continuous = self.prompt_config.n_classes == 0
701        prediction_key = sentiment_polarity_key if continuous else label_key
702        setting_key = _setting_feature(self.prompt_config, "setting")
703        marker_mode_key = _setting_feature(self.prompt_config, "marker_mode")
704        directed_key = _setting_feature(self.prompt_config, "directed")
705        window_size_key = _setting_feature(self.prompt_config, "window_size")
706        model_key = _setting_feature(self.prompt_config, "model")
707        sentiment_bin_key = _setting_feature(self.prompt_config, "sentiment_bin")
708        reason_key = _setting_feature(self.prompt_config, "reason")
709
710        interactions = list(doc.interactions)
711        targets: list[tuple[UIMAInteraction, UIMACharacterReference, UIMACharacterReference]] = []
712        if interactions:
713            if not overwrite and all(prediction_key in inter.additional_features for inter in interactions):
714                logger.info(f"All existing interactions already contain '{prediction_key}'. Skipping.")
715                return doc
716            for inter in interactions:
717                if not overwrite and prediction_key in inter.additional_features:
718                    continue
719                try:
720                    left = inter.ne1
721                    right = inter.ne2
722                except Exception:
723                    logger.warning("Skipping interaction without NE1/NE2 references.")
724                    continue
725                targets.append((inter, left, right))
726        else:
727            selected_characters = kwargs.get("characters")
728            if selected_characters is not None and not isinstance(selected_characters, list):
729                raise ValueError("characters must be a list[UIMACharacter] when provided.")
730            if selected_characters is not None and any(not isinstance(c, UIMACharacter) for c in selected_characters):
731                raise ValueError("characters must contain only UIMACharacter instances.")
732            pairs = list(_iter_interaction_pairs(
733                doc,
734                self.sampling_config,
735                directed=self.prompt_config.directed,
736                selected_characters=selected_characters,
737            ))
738            if not pairs:
739                logger.warning("No eligible character mention pairs found for HarryMotions prompting.")
740                return doc
741            for left, right in pairs:
742                inter: UIMAInteraction = doc.create_anno(UIMAInteraction, min(int(left.begin), int(right.begin)),
743                                                         max(int(left.end), int(right.end)), add_to_document=True)
744                inter.ne1 = left
745                inter.ne2 = right
746                targets.append((inter, left, right))
747
748        llm, fallback_llm = self._ensure_llms()
749
750        def _classify_single(inter: UIMAInteraction, left: UIMACharacterReference, right: UIMACharacterReference) -> tuple[UIMAInteraction, HarryMotionsPrediction]:
751            base, offset = _extract_window_text(doc, left, right, self.sampling_config.window_size)
752            role_text = _insert_role_markers(base, offset, exp_ref=left, targ_ref=right)
753            sample_text = _apply_indicator(role_text, self.prompt_config.marker_mode)
754            try:
755                prediction = classify_harrymotions_sample(
756                    sample_text, self.prompt_config, llm=llm, fallback_llm=fallback_llm,
757                )
758            except ValueError as exc:
759                logger.warning(f"Skipping invalid sample for interaction due to marker validation/parsing issue: {exc}")
760                prediction = HarryMotionsPrediction(label="")
761            return inter, prediction
762
763        workers = max(1, int(self.sampling_config.workers))
764        classified: list[tuple[UIMAInteraction, HarryMotionsPrediction]] = []
765        if workers == 1 or len(targets) <= 1:
766            for inter, left, right in tqdm(targets):
767                classified.append(_classify_single(inter, left, right))
768        else:
769            with ThreadPoolExecutor(max_workers=workers) as executor:
770                futures = [executor.submit(_classify_single, inter, left, right) for inter, left, right in targets]
771                for future in tqdm(as_completed(futures), total=len(futures)):
772                    classified.append(future.result())
773
774        for inter, prediction in classified:
775            inter.additional_features[setting_key] = _setting_id(self.prompt_config)
776            inter.additional_features[marker_mode_key] = self.prompt_config.marker_mode.value
777            inter.additional_features[directed_key] = self.prompt_config.directed
778            inter.additional_features[window_size_key] = self.sampling_config.window_size
779            inter.additional_features[model_key] = self.prompt_config.model.name
780            if prediction.reason:
781                inter.additional_features[reason_key] = prediction.reason
782
783            if continuous:
784                try:
785                    polarity = float(prediction.label) if prediction.label else 0.0
786                except (TypeError, ValueError):
787                    polarity = 0.0
788                polarity = max(-1.0, min(1.0, polarity))
789                inter.additional_features[sentiment_polarity_key] = polarity
790                inter.sentiment = polarity
791                inter.additional_features[sentiment_bin_key] = polarity
792            else:
793                inter.additional_features[label_key] = prediction.label
794                if self.prompt_config.n_classes == 2:
795                    sentiment = _sentiment_for_2class(prediction.label)
796                    if sentiment is None:
797                        sentiment = 0.0
798                    inter.sentiment = sentiment
799                    inter.additional_features[sentiment_bin_key] = sentiment
800
801        return doc
802
803
804def HarryMotionsPromptingStep(
805        prompt_config: HarryMotionsPromptConfig | None = None,
806        sampling_config: HarryMotionsSamplingConfig | None = None,
807        name: str = "HarryMotions Prompting",
808) -> PipelineStep[HarryMotionsPromptingProcessor]:
809    if prompt_config is None:
810        prompt_config = HarryMotionsPromptConfig(False, 2, HarryMotionsMarkerMode.ENTITY)
811    processor = HarryMotionsPromptingProcessor(
812        prompt_config=prompt_config,
813        sampling_config=sampling_config or HarryMotionsSamplingConfig(),
814    )
815    return PipelineStep(
816        name=name,
817        processor=processor,
818        unit_type=None,
819        modified_types=[UIMAInteraction],
820        added_additional_features=[
821            "harrymotions_*",
822        ],
823    )
824
825
826def HarryMotionsEntity2clStep(
827        sampling_config: HarryMotionsSamplingConfig | None = None,
828        name: str = "HarryMotions Entity 2cl Prompting",
829) -> PipelineStep[HarryMotionsPromptingProcessor]:
830    return HarryMotionsPromptingStep(
831        prompt_config=HarryMotionsPromptConfig(False, 2, HarryMotionsMarkerMode.ENTITY),
832        sampling_config=sampling_config,
833        name=name,
834    )
835
836
837def HarryMotionsEntity5clStep(
838        sampling_config: HarryMotionsSamplingConfig | None = None,
839        name: str = "HarryMotions Entity 5cl Prompting",
840) -> PipelineStep[HarryMotionsPromptingProcessor]:
841    return HarryMotionsPromptingStep(
842        prompt_config=HarryMotionsPromptConfig(False, 5, HarryMotionsMarkerMode.ENTITY),
843        sampling_config=sampling_config,
844        name=name,
845    )
846
847
848def HarryMotionsEntity8clStep(
849        sampling_config: HarryMotionsSamplingConfig | None = None,
850        name: str = "HarryMotions Entity 8cl Prompting",
851) -> PipelineStep[HarryMotionsPromptingProcessor]:
852    return HarryMotionsPromptingStep(
853        prompt_config=HarryMotionsPromptConfig(False, 8, HarryMotionsMarkerMode.ENTITY),
854        sampling_config=sampling_config,
855        name=name,
856    )