from __future__ import annotations import argparse import json import time from dataclasses import dataclass from pathlib import Path from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Union from json_repair import repair_json from langchain_huggingface import HuggingFacePipeline from utils.components import Lexicon, get_punctuation, join_tokens, strip_punctuation, normalize_candidates from utils.controlled_mode import ControlledModeService, read_controlled_input as read_controlled_input_file from utils.distractor_utils import ( ensure_distractor_count, extract_forced_target, reattach_punctuation_to_output, sorted_distractor_keys, ) from utils.runtime_config import build_llm_runtime_config from utils.surprisal import SurprisalScorer, SurprisalThresholdConfig, SurprisalThresholdFilter @dataclass class AgentConfig: model_id: str lexicon_path: str punctuations: Sequence[str] # generation params max_new_tokens: int = 128 temperature: float = 0.2 top_p: float = 0.9 do_sample: bool = True return_full_text: bool = False # lexicon params min_candidates: int = 10 max_candidates: int = 20 num_distractors: int = 3 num_workers: int = 1 # surprisal threshold params apply_surprisal_threshold: bool = False min_abs: Optional[float] = None min_delta: float = 0.0 absolute_threshold_only: bool = False surprisal_device: Optional[str] = None # token joining strategy (e.g., "" for CJK char tokens, " " for space-separated) token_joiner: str = "" class LLMAgent: """ For each sentence (list of tokens), iterate i=1..len-1: - prefix_sentence = cfg.token_joiner.join(tokens[:i+1]) (ends at target word) - target_word = tokens[i] - candidates = lexicon neighbors for target_word (punct preserved) - LLM chooses distractors from candidates and returns schema-valid JSON - optional: enforce surprisal thresholds over chosen distractors """ def __init__(self, prompt: Any, config: AgentConfig, llm: Optional[HuggingFacePipeline] = None) -> None: self.prompt = prompt self.cfg = config if not 1 <= int(self.cfg.num_distractors) <= 10: raise ValueError("num_distractors must be in [1, 10].") if int(self.cfg.num_workers) < 1: raise ValueError("num_workers must be >= 1.") self.lexicon = Lexicon(self.cfg.lexicon_path) self.puncts = set(self.cfg.punctuations) self._neighbor_cache: dict[str, tuple[str, ...]] = {} self.llm = llm or HuggingFacePipeline.from_model_id( model_id=self.cfg.model_id, task="text-generation", pipeline_kwargs={ "max_new_tokens": self.cfg.max_new_tokens, "temperature": self.cfg.temperature, "top_p": self.cfg.top_p, "do_sample": self.cfg.do_sample, "return_full_text": self.cfg.return_full_text, }, ) self._configure_generation_padding() self._threshold_filter = SurprisalThresholdFilter( config=SurprisalThresholdConfig(enabled=False) ) if self.cfg.apply_surprisal_threshold: scorer = SurprisalScorer( model_id=self.cfg.model_id, device=self.cfg.surprisal_device, ) self._threshold_filter = SurprisalThresholdFilter( config=SurprisalThresholdConfig( enabled=True, min_abs=self.cfg.min_abs, min_delta=self.cfg.min_delta, absolute_threshold_only=self.cfg.absolute_threshold_only, ), scorer=scorer, ) def _configure_generation_padding(self) -> None: """Avoid right-padding warnings for decoder-only batched generation.""" pipe = getattr(self.llm, "pipeline", None) if pipe is None: return tokenizer = getattr(pipe, "tokenizer", None) model = getattr(pipe, "model", None) if tokenizer is None: return model_cfg = getattr(model, "config", None) is_decoder_only = not bool(getattr(model_cfg, "is_encoder_decoder", False)) if not is_decoder_only: return if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = "left" if model_cfg is not None and getattr(model_cfg, "pad_token_id", None) is None: model_cfg.pad_token_id = tokenizer.pad_token_id # ---------- helper function ---------- def _get_cached_neighbors(self, core: str) -> tuple[str, ...]: if core not in self._neighbor_cache: raw = self.lexicon.get_neighbor(core, min_size=self.cfg.min_candidates, max_size=self.cfg.max_candidates) clean = normalize_candidates(raw) clean = [w for w in clean if w != core] self._neighbor_cache[core] = tuple(clean) return self._neighbor_cache[core] def _reattach_punct(self, out: Dict[str, Any], prefix: str, suffix: str) -> Dict[str, Any]: return reattach_punctuation_to_output(out, prefix=prefix, suffix=suffix) @staticmethod def _sorted_distractor_keys(out: Dict[str, Any]) -> List[str]: return sorted_distractor_keys(out) def _ensure_distractor_count( self, out: Dict[str, Any], *, target_word: str, target_core: str, candidate_pool: Sequence[str], allow_candidate_fill: bool = True, ) -> Dict[str, Any]: return ensure_distractor_count( out, num_distractors=self.cfg.num_distractors, target_word=target_word, target_core=target_core, candidate_pool=candidate_pool, puncts=self.puncts, allow_candidate_fill=allow_candidate_fill, ) @staticmethod def _extract_forced_target(core: str) -> Optional[str]: return extract_forced_target(core) # ---------- candidate generation ---------- def get_candidates(self, word: str): prefix, core, suffix = strip_punctuation(word, self.puncts) if not core: return prefix, core, suffix, ("X" * len(word),), None forced_target = self._extract_forced_target(core) if forced_target is not None: forced_len = max(1, len(forced_target)) # Do not query lexicon / LLM for explicitly marked target tokens. return prefix, forced_target, suffix, tuple(), forced_len neighbors = self._get_cached_neighbors(core) return prefix, core, suffix, neighbors, None # ---------- LLM + parsing ---------- def _invoke_raw(self, sentence_prefix: str, target_word: str, candidates: Sequence[str]) -> str: prompt_text = self.prompt.render_text( sentence_prefix=sentence_prefix, word=target_word, candidates=candidates, ) resp = self.llm.invoke(prompt_text) return self._extract_generated_text(resp) @staticmethod def _extract_generated_text(resp: Any) -> str: if isinstance(resp, str): return resp if isinstance(resp, dict): for key in ("generated_text", "text", "content"): if key in resp and resp[key] is not None: return str(resp[key]) if isinstance(resp, list): if not resp: return "" first = resp[0] if isinstance(first, dict): return LLMAgent._extract_generated_text(first) if isinstance(first, list): return LLMAgent._extract_generated_text(first[0] if first else "") return str(resp) def _invoke_raw_batch( self, prompts: Sequence[str], *, progress_callback: Optional[Callable[[int, int], None]] = None, ) -> List[str]: """Batch invoke HF pipeline; prefer dataset-style streaming on GPU.""" prompt_list = list(prompts) total = len(prompt_list) if total == 0: return [] def _report(done: int) -> None: if progress_callback is not None: progress_callback(done, total) pipe = getattr(self.llm, "pipeline", None) if pipe is not None: try: from transformers.pipelines.pt_utils import KeyDataset class _PromptDataset: def __init__(self, rows: Sequence[str]) -> None: self.rows = rows def __len__(self) -> int: return len(self.rows) def __getitem__(self, idx: int) -> Dict[str, str]: return {"text": self.rows[idx]} dataset = _PromptDataset(prompt_list) streamed = pipe( KeyDataset(dataset, "text"), batch_size=max(1, int(self.cfg.num_workers)), ) out: List[str] = [] for i, item in enumerate(streamed, start=1): out.append(self._extract_generated_text(item)) _report(i) return out except Exception: pass try: batch_out = pipe( prompt_list, batch_size=max(1, int(self.cfg.num_workers)), ) out = [self._extract_generated_text(item) for item in batch_out] for i in range(1, len(out) + 1): _report(i) return out except Exception: pass out: List[str] = [] for i, p in enumerate(prompt_list, start=1): out.append(self._extract_generated_text(self.llm.invoke(p))) _report(i) return out def _parse_to_dict(self, raw: str, repair: bool = True) -> Dict[str, Any]: try: obj = self.prompt.parser.parse(raw) return obj.model_dump() except Exception: if not repair: raise try: fixed = repair_json(raw) data = json.loads(fixed) except Exception: return {} if isinstance(data, list): for x in reversed(data): if isinstance(x, dict): data = x break if not isinstance(data, dict): return {} data.setdefault("source", "") data.setdefault("distractor1", "X") data.setdefault("distractor2", "X") data.setdefault("distractor3", None) try: obj = self.prompt.parser.pydantic_object.model_validate(data) return obj.model_dump() except Exception: return {} # ---------- public API ---------- def iter_jobs(self, input_data: Sequence[Sequence[str]]) -> Iterable[Dict[str, Any]]: for s_idx, tokens in enumerate(input_data): for i in range(len(tokens)): target = tokens[i] sentence_prefix = join_tokens(tokens[: i + 1], join_with=self.cfg.token_joiner, puncts=self.puncts) context_prefix = join_tokens(tokens[:i], join_with=self.cfg.token_joiner, puncts=self.puncts) pfx, core, sfx, candidates, forced_len = self.get_candidates(target) yield { "sentence_index": s_idx, "word_index": i, "sentence_prefix": sentence_prefix, "context_prefix": context_prefix, "target_word": target, "target_core": core, "punct_prefix": pfx, "punct_suffix": sfx, "candidates": candidates, "forced_dummy_len": forced_len, } def read_controlled_input(self, path: Union[str, Path], split_on: Optional[str]) -> List[Dict[str, Any]]: return read_controlled_input_file(path, split_on) def _process_job(self, job: Dict[str, Any], *, repair: bool = True) -> Dict[str, Any]: forced_len = job.get("forced_dummy_len") if job["word_index"] == 0 or forced_len is not None: dummy = "X" * int(forced_len if forced_len is not None else len(job["target_word"])) out = {"source": job["target_core"], "distractor1": dummy} out = self._ensure_distractor_count( out, target_word=job["target_word"], target_core=job["target_core"], candidate_pool=job["candidates"], allow_candidate_fill=False, ) if forced_len is not None: out = self._reattach_punct(out, job["punct_prefix"], job["punct_suffix"]) else: raw = self._invoke_raw(job["sentence_prefix"], job["target_core"], job["candidates"]) out = self._parse_to_dict(raw, repair=repair) # Fallback if parsing totally failed if not out: dummy = "X" * len(job["target_word"]) out = { "source": job["target_core"], "distractor1": dummy, "distractor2": dummy, "distractor3": None, } out = self._ensure_distractor_count( out, target_word=job["target_word"], target_core=job["target_core"], candidate_pool=job["candidates"], allow_candidate_fill=True, ) out = self._threshold_filter.enforce_output_threshold( out=out, context_prefix=job["context_prefix"], target_core=job["target_core"], candidate_pool=job["candidates"], puncts=self.puncts, token_joiner=self.cfg.token_joiner, ) out = self._reattach_punct(out, job["punct_prefix"], job["punct_suffix"]) # attach metadata (useful for reconstructing maze items) out.update( sentence_index=job["sentence_index"], word_index=job["word_index"], sentence_prefix=job["sentence_prefix"], target_word=job["target_word"], ) return out def _process_job_with_raw(self, job: Dict[str, Any], raw: str, *, repair: bool = True) -> Dict[str, Any]: out = self._parse_to_dict(raw, repair=repair) if not out: dummy = "X" * len(job["target_word"]) out = { "source": job["target_core"], "distractor1": dummy, "distractor2": dummy, "distractor3": None, } out = self._ensure_distractor_count( out, target_word=job["target_word"], target_core=job["target_core"], candidate_pool=job["candidates"], allow_candidate_fill=True, ) out = self._threshold_filter.enforce_output_threshold( out=out, context_prefix=job["context_prefix"], target_core=job["target_core"], candidate_pool=job["candidates"], puncts=self.puncts, token_joiner=self.cfg.token_joiner, ) out = self._reattach_punct(out, job["punct_prefix"], job["punct_suffix"]) out.update( sentence_index=job["sentence_index"], word_index=job["word_index"], sentence_prefix=job["sentence_prefix"], target_word=job["target_word"], ) return out def run( self, input_data: Sequence[Sequence[str]], repair: bool = True, limit: Optional[int] = None, ) -> List[Dict[str, Any]]: """Run over input_data and return a sentence-grouped list of JSON dict outputs.""" run_start = time.perf_counter() jobs = list(self.iter_jobs(input_data)) if limit is not None: jobs = jobs[:limit] if not jobs: return [] total_sentences = len(input_data) print(f"[LLMAgent] Processing {len(jobs)} word-jobs from {total_sentences} sentences...") print(f"[LLMAgent] Generation batch size: {self.cfg.num_workers}") outputs: List[Dict[str, Any]] = [] llm_jobs: List[Dict[str, Any]] = [] llm_prompts: List[str] = [] total_llm_seconds = 0.0 for job in jobs: if job["word_index"] == 0 or job.get("forced_dummy_len") is not None: outputs.append(self._process_job(job, repair=repair)) else: llm_jobs.append(job) llm_prompts.append( self.prompt.render_text( sentence_prefix=job["sentence_prefix"], word=job["target_core"], candidates=job["candidates"], ) ) if llm_jobs: batch_size = max(1, int(self.cfg.num_workers)) total_llm_jobs = len(llm_jobs) print( f"[LLMAgent] LLM generation jobs: {total_llm_jobs} " f"(dataset-style, batch_size={batch_size})" ) llm_start = time.perf_counter() report_every = max(1, total_llm_jobs // 20) def _progress(done: int, total: int) -> None: if done == 1 or done == total or done % report_every == 0: elapsed = time.perf_counter() - llm_start avg = elapsed / max(1, done) eta = avg * (total - done) print( f"[LLMAgent] Progress: {done}/{total} LLM jobs done, " f"elapsed={elapsed:.2f}s, avg={avg:.2f}s/job, eta={eta:.2f}s" ) raw_outputs = self._invoke_raw_batch(llm_prompts, progress_callback=_progress) if len(raw_outputs) != len(llm_jobs): raw_outputs = [ self._invoke_raw(j["sentence_prefix"], j["target_core"], j["candidates"]) for j in llm_jobs ] total_llm_seconds = time.perf_counter() - llm_start for job, raw in zip(llm_jobs, raw_outputs): outputs.append(self._process_job_with_raw(job, raw, repair=repair)) outputs.sort(key=lambda x: (x["sentence_index"], x["word_index"])) grouped: Dict[int, List[Dict[str, Any]]] = {} for out in outputs: grouped.setdefault(int(out["sentence_index"]), []).append(out) run_seconds = time.perf_counter() - run_start if llm_jobs: print( f"[LLMAgent] Timing summary: total={run_seconds:.2f}s, " f"llm_calls={total_llm_seconds:.2f}s, " f"avg_llm_call={total_llm_seconds / max(1, len(llm_jobs)):.2f}s" ) else: print(f"[LLMAgent] Timing summary: total={run_seconds:.2f}s") return [{"sentence_index": s_idx, "words": grouped[s_idx]} for s_idx in sorted(grouped.keys())] def run_controlled_experiment( self, controlled_items: Sequence[Dict[str, Any]], repair: bool = True, limit: Optional[int] = None, ) -> List[Dict[str, Any]]: """ Controlled mode: - one shared distractor set per (item_id, word_index) - all conditions of the item reuse this shared set """ service = ControlledModeService( lexicon=self.lexicon, puncts=self.puncts, num_distractors=self.cfg.num_distractors, token_joiner=self.cfg.token_joiner, min_candidates=self.cfg.min_candidates, max_candidates=self.cfg.max_candidates, apply_surprisal_threshold=self.cfg.apply_surprisal_threshold, threshold_filter=self._threshold_filter, get_candidates=self.get_candidates, invoke_raw=self._invoke_raw, parse_to_dict=self._parse_to_dict, ensure_distractor_count=self._ensure_distractor_count, reattach_punct=self._reattach_punct, sorted_distractor_keys=self._sorted_distractor_keys, ) return service.run(controlled_items=controlled_items, repair=repair, limit=limit) @staticmethod def save_jsonl(records: Sequence[Dict[str, Any]], path: Union[str, Path]) -> None: """Machine-friendly JSONL (one object per line).""" path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as f: for r in records: f.write(json.dumps(r, ensure_ascii=False) + "\n") @staticmethod def save_pretty_json(records: Sequence[Dict[str, Any]], path: Union[str, Path]) -> None: """Human-friendly JSON (pretty-printed).""" path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as f: json.dump(records, f, ensure_ascii=False, indent=2) # ------------------------- # minimal usage example # ------------------------- if __name__ == "__main__": parser = argparse.ArgumentParser(description="LLM Maze distractor generation (lexicon + LLM selector).") parser.add_argument("--config-path", default="config.yaml", help="Path to YAML config file.") parser.add_argument("--language-code", default=None, help="Override LANGUAGE_CODE from config.") parser.add_argument("--min-abs", type=float, default=None, help="Minimum absolute distractor surprisal.") parser.add_argument("--min-delta", type=float, default=None, help="Minimum surprisal delta over target.") parser.add_argument( "--absolute-threshold-only", action=argparse.BooleanOptionalAction, default=None, help="Use only min_abs as hard threshold (ignore target + min_delta).", ) parser.add_argument("--surprisal-device", default=None, help='Device for surprisal scorer, e.g. "cuda" or "cpu".') parser.add_argument("--limit", type=int, default=None, help="Optional max number of word-jobs to run.") parser.add_argument("--output-path", default=None, help="Optional explicit output JSON path.") parser.add_argument("--model-id", default=None, help="Override MODEL_ID from config.") parser.add_argument("--input-data-path", default=None, help="Optional explicit input data path override.") parser.add_argument("--lexicon-path", default=None, help="Optional explicit lexicon path override.") parser.add_argument("--template-dir", default=None, help="Optional explicit template directory override.") parser.add_argument( "--processing-mode", default=None, choices=["naturalistic_reading", "controlled_experiment"], help="Processing mode: naturalistic_reading or controlled_experiment.", ) parser.add_argument("--word-separator", default=None, help="Override WORD_SEPARATOR from config.") parser.add_argument("--num-distractors", type=int, default=None, help="Number of distractors per word (1-10).") parser.add_argument("--num-workers", type=int, default=None, help="Parallel workers for word-level processing (>=1).") args = parser.parse_args() from utils.components import load_config, read_sentences_input from utils.maze_prompt import MazeLLMPrompt yaml_cfg = load_config(args.config_path) runtime_cfg = build_llm_runtime_config( yaml_cfg=yaml_cfg, overrides={ "AGENT_TYPE": "llm", "LANGUAGE_CODE": args.language_code, "MODEL_ID": args.model_id, "INPUT_DATA_PATH": args.input_data_path, "LEXICON_PATH": args.lexicon_path, "TEMPLATE_DIR": args.template_dir, "OUTPUT_PATH": args.output_path, "PROCESSING_MODE": args.processing_mode, "WORD_SEPARATOR": args.word_separator, "NUM_DISTRACTORS": args.num_distractors, "NUM_WORKERS": args.num_workers, "SURPRISAL_MIN_ABS": args.min_abs, "SURPRISAL_MIN_DELTA": args.min_delta, "SURPRISAL_ABSOLUTE_THRESHOLD_ONLY": args.absolute_threshold_only, "SURPRISAL_DEVICE": args.surprisal_device, }, ) language_code = runtime_cfg.language_code agent_type = runtime_cfg.agent_type model_id = runtime_cfg.model_id processing_mode = runtime_cfg.processing_mode main_path = Path(runtime_cfg.template_dir) word_separator = runtime_cfg.word_separator puncts = get_punctuation(language_code) if not main_path.exists(): raise ValueError(f"Template directory not found for LANGUAGE_CODE='{language_code}': {main_path}") if not Path(runtime_cfg.lexicon_path).exists(): raise ValueError(f"Lexicon file not found: {runtime_cfg.lexicon_path}") if not Path(runtime_cfg.input_data_path).exists(): raise ValueError(f"Input data file not found: {runtime_cfg.input_data_path}") maze_prompt = MazeLLMPrompt( path_to_user_template=main_path / "base.txt", path_to_extension_template=main_path / "extension.txt", path_to_system_template=main_path / "system.txt", ) agent_cfg = AgentConfig( model_id=model_id , lexicon_path=runtime_cfg.lexicon_path, punctuations=puncts, num_distractors=runtime_cfg.num_distractors, num_workers=runtime_cfg.num_workers, apply_surprisal_threshold=runtime_cfg.apply_surprisal_threshold, min_abs=runtime_cfg.min_abs, min_delta=runtime_cfg.min_delta, absolute_threshold_only=runtime_cfg.absolute_threshold_only, surprisal_device=runtime_cfg.surprisal_device, token_joiner="" if str(word_separator).strip() == "" else str(word_separator), ) agent = LLMAgent(prompt=maze_prompt, config=agent_cfg) main_start = time.perf_counter() if processing_mode == "controlled_experiment": controlled_items = agent.read_controlled_input( runtime_cfg.input_data_path, split_on=word_separator, ) outputs = agent.run_controlled_experiment(controlled_items, limit=args.limit) else: input_data = read_sentences_input(runtime_cfg.input_data_path, split_on=word_separator) outputs = agent.run(input_data, limit=args.limit) generation_elapsed = time.perf_counter() - main_start print( f"[LLMAgent] Generation finished in {generation_elapsed:.2f}s " f"(mode={processing_mode}, records={len(outputs)})." ) output_path = args.output_path or runtime_cfg.output_path agent.save_pretty_json(outputs, output_path) total_elapsed = time.perf_counter() - main_start print(f"Saved output to: {output_path}") print(f"[LLMAgent] Total elapsed (generation + save): {total_elapsed:.2f}s")