diff --git a/src/twinkle_agentic/chunker/__init__.py b/src/twinkle_agentic/chunker/__init__.py deleted file mode 100644 index f826a6452..000000000 --- a/src/twinkle_agentic/chunker/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .base import Chunker -from .native import NativeChunker - -__all__ = ['Chunker', 'NativeChunker'] diff --git a/src/twinkle_agentic/chunker/base.py b/src/twinkle_agentic/chunker/base.py deleted file mode 100644 index 22beb8b88..000000000 --- a/src/twinkle_agentic/chunker/base.py +++ /dev/null @@ -1,14 +0,0 @@ -from abc import ABC, abstractmethod - -from twinkle.data_format import Trajectory -from twinkle_agentic.data_format import Chunks - - -class Chunker(ABC): - """ - TODO: Experimental feature, wait for testing - """ - - @abstractmethod - def __call__(self, trajectory: Trajectory) -> Chunks: - raise NotImplementedError diff --git a/src/twinkle_agentic/chunker/native.py b/src/twinkle_agentic/chunker/native.py deleted file mode 100644 index f5879f3c0..000000000 --- a/src/twinkle_agentic/chunker/native.py +++ /dev/null @@ -1,254 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -import re -from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence - -from twinkle.data_format import Trajectory -from twinkle_agentic.data_format import Chunk, Chunks -from .base import Chunker - -# Recursive separator list, coarsest → finest. The empty string at the -# end forces a hard character cut when nothing finer fits. -_DEFAULT_SEPARATORS: tuple = ( - '\n\n', - '\n', - '。', - '．', - '.', - '！', - '!', - '？', - '?', - '；', - ';', - '，', - ',', - ' ', - '', -) - -_MULTIMODAL_TYPES = ('image', 'video', 'audio') - -_SplitFn = Optional[Callable[[str], List[str]]] - - -class NativeChunker(Chunker): - """Character-level recursive chunker for trajectories. - TODO: Experimental feature, wait for testing - Args: - chunk_size: Soft upper bound (in characters) for every emitted - text chunk. Must be positive. - separators: Ordered separator list. The chunker tries each - separator in turn; any piece still larger than - ``chunk_size`` is re-split with the next one. A terminal - ``''`` (hard character cut) is appended automatically if - missing so the algorithm is guaranteed to terminate. - passage_boundary_re: Optional regex (compiled with - ``re.MULTILINE``) whose matches act as **hard, non-mergeable** - passage boundaries on the first user message. The regex - match is preserved at the start of the next piece (so - ``''.join(pieces) == text``). Pieces that are already - ``<= chunk_size`` are emitted as-is and are **never merged** - across boundaries; only pieces that still exceed - ``chunk_size`` fall back to the normal recursive split + merge. - This is how you keep e.g. HotpotQA passages atomic per - ````. - """ - - def __init__( - self, - chunk_size: int = 1024, - separators: Sequence[str] | None = None, - passage_boundary_re: str | None = None, - ): - if chunk_size <= 0: - raise ValueError(f'chunk_size must be positive, got {chunk_size}') - self.chunk_size = chunk_size - seps = tuple(separators) if separators is not None else _DEFAULT_SEPARATORS - if '' not in seps: - seps += ('', ) - self.separators = seps - self.passage_boundary_re: re.Pattern | None = ( - re.compile(passage_boundary_re, re.MULTILINE) if passage_boundary_re else None) - - # ------------------------------------------------------------------ - # public entry - # ------------------------------------------------------------------ - def __call__(self, trajectory: Trajectory) -> Chunks: - chunks: list[Chunk] = [] - first_user_done = False - # ``round`` is 1-indexed at the first user message. Any messages - # emitted before that (e.g., leading ``system``) carry round 0. - round_idx = 0 - for msg in trajectory.get('messages') or []: - is_user = msg.get('role') == 'user' - if is_user: - round_idx += 1 - split = (self._split_text if is_user and not first_user_done else None) - if is_user: - first_user_done = True - for chunk in self._parts(msg, split): - chunk['round'] = round_idx - chunks.append(chunk) - return Chunks(chunks=chunks) - - # ------------------------------------------------------------------ - # message → chunks decomposition - # ------------------------------------------------------------------ - def _parts(self, message: dict[str, Any], split: _SplitFn) -> Iterator[Chunk]: - role = message.get('role') or 'user' - tcid = message.get('tool_call_id') - - rc = message.get('reasoning_content') - if rc: - yield _text_chunk(role, rc, kind='reasoning_content', tool_call_id=tcid) - - content = message.get('content') - if isinstance(content, str): - yield from self._emit_text(role, content, split, tcid) - elif isinstance(content, list): - for part in content: - if not isinstance(part, dict): - continue - ptype = part.get('type') - if ptype == 'text': - yield from self._emit_text(role, part.get('text') or '', split, tcid) - elif ptype in _MULTIMODAL_TYPES: - # Keep raw part so Chunks.to_trajectory can rebuild - # the original OpenAI-style entry verbatim. - yield { # type: ignore[misc] - 'type': ptype, 'content': part.get(ptype), - 'raw': dict(part), 'role': role, - } - - for tc in message.get('tool_calls') or []: - yield _text_chunk(role, '', kind='tool_call', tool_call=tc, tool_call_id=tcid) - - def _emit_text(self, role: str, text: str, split: _SplitFn, tool_call_id: str | None) -> Iterator[Chunk]: - if not text: - return - pieces = split(text) if split is not None else [text] - for piece in pieces: - if piece: - yield _text_chunk(role, piece, tool_call_id=tool_call_id) - - # ------------------------------------------------------------------ - # recursive text splitter - # ------------------------------------------------------------------ - def _split_text(self, text: str) -> list[str]: - if not text: - return [] - if self.passage_boundary_re is None: - if len(text) <= self.chunk_size: - return [text] - return self._merge(self._recursive_split(text, list(self.separators))) - # Force-split first; each forced piece is kept intact when it is - # already short enough, and is recursively re-split (but NOT - # merged with sibling passages) when it exceeds ``chunk_size``. - out: list[str] = [] - for piece in self._force_split(text): - if not piece or not piece.strip(): - continue - if len(piece) <= self.chunk_size: - out.append(piece) - else: - out.extend(self._merge(self._recursive_split(piece, list(self.separators)))) - return out - - def _force_split(self, text: str) -> list[str]: - """Split ``text`` at every ``passage_boundary_re`` match; the - match itself sticks to the start of the **next** piece, so - ``''.join(_force_split(text)) == text``. - """ - assert self.passage_boundary_re is not None - matches = list(self.passage_boundary_re.finditer(text)) - if not matches: - return [text] - out: list[str] = [] - prev = 0 - for m in matches: - start = m.start() - if start > prev: - out.append(text[prev:start]) - prev = start - if prev < len(text): - out.append(text[prev:]) - return out - - def _recursive_split(self, text: str, separators: list[str]) -> list[str]: - if len(text) <= self.chunk_size: - return [text] if text else [] - # Terminal: no more separators, or next one is the hard-cut sentinel. - if not separators or separators[0] == '': - return _hard_cut(text, self.chunk_size) - - sep, *rest = separators - out: list[str] = [] - for piece in _split_keep(text, sep): - if not piece: - continue - if len(piece) <= self.chunk_size: - out.append(piece) - else: - out.extend(self._recursive_split(piece, rest)) - return out - - def _merge(self, pieces: list[str]) -> list[str]: - """Greedy concatenation: small fragments fuse up to ``chunk_size`` - without exceeding it. Relative order is preserved. - """ - merged: list[str] = [] - buf = '' - for p in pieces: - if not p: - continue - if buf and len(buf) + len(p) > self.chunk_size: - merged.append(buf) - buf = '' - buf += p - if buf: - merged.append(buf) - return merged - - -# ---------------------------------------------------------------------- -# helpers -# ---------------------------------------------------------------------- -def _split_keep(text: str, sep: str) -> list[str]: - """``str.split(sep)`` but the separator stays glued to the end of - each left-hand piece, so ``''.join(result) == text``. - """ - if not sep or sep not in text: - return [text] if text else [] - out: list[str] = [] - start, n = 0, len(sep) - while (i := text.find(sep, start)) != -1: - out.append(text[start:i + n]) - start = i + n - if start < len(text): - out.append(text[start:]) - return out - - -def _hard_cut(text: str, size: int) -> list[str]: - return [text[i:i + size] for i in range(0, len(text), size)] if text else [] - - -def _text_chunk( - role: str, - content: str, - *, - kind: str | None = None, - tool_call: Any = None, - tool_call_id: str | None = None, -) -> Chunk: - raw: dict[str, Any] = {} - if kind is not None: - raw['kind'] = kind - if tool_call is not None: - raw['tool_call'] = tool_call - if tool_call_id is not None: - raw['tool_call_id'] = tool_call_id - chunk: Chunk = {'type': 'text', 'content': content, 'role': role} # type: ignore[assignment] - if raw: - chunk['raw'] = raw - return chunk diff --git a/src/twinkle_agentic/classifier/__init__.py b/src/twinkle_agentic/classifier/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/twinkle_agentic/classifier/base.py b/src/twinkle_agentic/classifier/base.py new file mode 100644 index 000000000..d79f6c2e2 --- /dev/null +++ b/src/twinkle_agentic/classifier/base.py @@ -0,0 +1,11 @@ +from abc import ABC, abstractmethod + + +class Classifier(ABC): + + def __init__(self, model_path: str, **kwargs): + self.model_path = model_path + + @abstractmethod + def classify(self, text: str) -> str: + pass \ No newline at end of file diff --git a/src/twinkle_agentic/condenser/__init__.py b/src/twinkle_agentic/condenser/__init__.py index e78545002..a48fd73c1 100644 --- a/src/twinkle_agentic/condenser/__init__.py +++ b/src/twinkle_agentic/condenser/__init__.py @@ -1,5 +1,4 @@ from .base import Condenser -from .keyword import KeywordCondenser -from .model import ModelCondenser +from .facts import FactsCondenser -__all__ = ['Condenser', 'KeywordCondenser', 'ModelCondenser'] +__all__ = ['Condenser', 'FactsCondenser'] diff --git a/src/twinkle_agentic/condenser/base.py b/src/twinkle_agentic/condenser/base.py index 5e42dab17..8fd2acd61 100644 --- a/src/twinkle_agentic/condenser/base.py +++ b/src/twinkle_agentic/condenser/base.py @@ -1,13 +1,201 @@ -from abc import ABC, abstractmethod +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations -from twinkle_agentic.data_format import Chunks +import math +import re +from typing import TYPE_CHECKING, Any, Sequence +from twinkle_agentic.utils.llm_backup import llm_backup -class Condenser(ABC): - """ - TODO: Experimental feature, wait for testing +if TYPE_CHECKING: + from twinkle.data_format import SamplingParams, Trajectory # noqa: F401 + from twinkle.sampler.base import Sampler # noqa: F401 + + +DEFAULT_USER_PROMPT_TEMPLATE = """\ +Compress the following text as much as possible while preserving all key information. + +## Target length +HARD CEILING: {budget} chars. If core facts fit in far fewer chars, output fewer. + +## Text +{text}""" + + +class Condenser: + """Base condenser with progressive distillation via llm_backup. + + Subclasses customize compression behavior by providing their own + ``system_prompt``, ``user_prompt_template``, and ``lora_path``. + The shared ``_sample`` method (decorated with ``@llm_backup``) handles + the student-teacher routing transparently. + + Teacher is a global OpenAI-compatible API configured via env vars: + - LLM_BACKUP_MODEL: teacher model name + - LLM_BACKUP_API_KEY: API key + - LLM_BACKUP_BASE_URL: API endpoint + + Args: + sampler: Student model sampler (local inference, shared across types). + compression_ratio: Target compression factor (> 1). + model_path: Model identifier. + sampling_params: Default sampling params. + system_prompt: System prompt for this condenser type. + user_prompt_template: User prompt template. Must contain + ``{budget}`` and ``{text}``. May contain ``{query}``. + min_budget_chars: Floor for the character budget in the prompt. + template: Optional :class:`Template` for special token stripping. + lora_path: LoRA adapter path specific to this condenser type. + Each subclass can use a different LoRA for its task. """ - @abstractmethod - def __call__(self, chunks: Chunks, **kwargs) -> Chunks: - raise NotImplementedError + def __init__( + self, + sampler: Sampler, + compression_ratio: float = 2.0, + *, + model_path: str = '', + sampling_params: SamplingParams | None = None, + system_prompt: str = 'You are a text compression assistant.', + user_prompt_template: str | None = None, + min_budget_chars: int = 250, + template: Any | None = None, + lora_path: str | None = None, + ): + if sampler is None: + raise ValueError('sampler is required') + if compression_ratio <= 1.0: + raise ValueError(f'compression_ratio must be > 1, got {compression_ratio}') + if min_budget_chars < 1: + raise ValueError(f'min_budget_chars must be >= 1, got {min_budget_chars}') + + tpl = user_prompt_template or DEFAULT_USER_PROMPT_TEMPLATE + if '{budget}' not in tpl or '{text}' not in tpl: + raise ValueError('user_prompt_template must contain both {budget} and {text}') + + self.model_path = model_path + self.sampler = sampler + self.compression_ratio = float(compression_ratio) + self.sampling_params = sampling_params + self.system_prompt = system_prompt + self.user_prompt_template = tpl + self.min_budget_chars = int(min_budget_chars) + self.template = template + self.lora_path = lora_path if lora_path else None + self._special_tokens_cache: tuple[str, ...] | None = None + + # ------------------------------------------------------------------ + # public entry point (pre/post processing, NOT decorated) + # ------------------------------------------------------------------ + def __call__(self, text: str, system: str = None, query: str = None, + sampling_params: Any = None) -> str: + system = system or self.system_prompt + budget = max(self.min_budget_chars, math.ceil(len(text) / self.compression_ratio)) + if budget >= len(text): + return text + trajectory = self._make_trajectory(system, self.user_prompt_template, text, budget, query) + sp = sampling_params or self.sampling_params or self._default_sampling_params(budget) + + raw = self._sample(trajectory=trajectory, sampling_params=sp, query=query) + + result = self._postprocess(raw, text, self._get_special_tokens()) + return result if result is not None else text + + # ------------------------------------------------------------------ + # student sampling (decorated with llm_backup) + # ------------------------------------------------------------------ + @llm_backup(key_params=["query"]) + def _sample(self, trajectory, sampling_params, query: str = None) -> str: + """Student model: trajectory + sampling_params -> raw text.""" + sample_kwargs: dict[str, Any] = {'sampling_params': sampling_params} + if self.lora_path is None: + sample_kwargs['use_base_model'] = True + else: + sample_kwargs['adapter_path'] = self.lora_path + responses = self.sampler.sample([trajectory], **sample_kwargs) + return self._decoded(list(responses)[0]) if responses else '' + + # ------------------------------------------------------------------ + # internals + # ------------------------------------------------------------------ + def _get_special_tokens(self) -> tuple[str, ...]: + if self._special_tokens_cache is not None: + return self._special_tokens_cache + tpl = self.template or getattr(self.sampler, 'template', None) + tokenizer = getattr(tpl, 'tokenizer', None) if tpl is not None else None + tokens: list[str] = [] + if tokenizer is not None: + extras = getattr(tokenizer, 'all_special_tokens', None) or [] + if extras: + tokens.extend(t for t in extras if isinstance(t, str) and t and not t.isspace()) + else: + for attr in ('eos_token', 'pad_token', 'bos_token'): + t = getattr(tokenizer, attr, None) + if isinstance(t, str) and t: + tokens.append(t) + self._special_tokens_cache = tuple(dict.fromkeys(tokens)) + return self._special_tokens_cache + + + # ------------------------------------------------------------------ + # static helpers + # ------------------------------------------------------------------ + _CODE_FENCE_RE = re.compile(r'^```[a-zA-Z]*\s*\n(.*?)\n```\s*$', re.DOTALL) + + @staticmethod + def _make_trajectory(system: str, user_template: str, text: str, + budget: int, query: str | None = None) -> dict: + """Build a trajectory dict for sampler / API.""" + user = user_template.replace('{budget}', str(budget)) + user = user.replace('{text}', text) + if '{query}' in user: + q_text = ( + query.strip() if isinstance(query, str) and query and query.strip() else + '(no explicit query; compress by general salience)') + user = user.replace('{query}', q_text) + return { + 'messages': [ + {'role': 'system', 'content': system}, + {'role': 'user', 'content': user}, + ], + } + + @staticmethod + def _default_sampling_params(budget: int): + from twinkle.data_format.sampling import SamplingParams + max_new = max(512, budget * 3 + 128) + return SamplingParams(temperature=0.0, max_tokens=max_new) + + @staticmethod + def _postprocess(raw: str, original: str, special_tokens: tuple[str, ...]) -> str | None: + text = Condenser._strip_special_tokens( + Condenser._strip_code_fences(raw), special_tokens).strip() + if not text or not Condenser._has_alnum(text): + return None + if len(text) >= len(original): + return None + return text + + @staticmethod + def _decoded(response: Any) -> str: + seqs = getattr(response, 'sequences', None) or [] + if not seqs: + return '' + return getattr(seqs[0], 'decoded', None) or '' + + @staticmethod + def _strip_code_fences(text: str) -> str: + stripped = text.strip() + m = Condenser._CODE_FENCE_RE.match(stripped) + return m.group(1) if m else text + + @staticmethod + def _strip_special_tokens(text: str, tokens: Sequence[str]) -> str: + for tok in tokens: + if tok and tok in text: + text = text.replace(tok, '') + return text + + @staticmethod + def _has_alnum(text: str) -> bool: + return any(ch.isalnum() for ch in text) diff --git a/src/twinkle_agentic/condenser/facts.py b/src/twinkle_agentic/condenser/facts.py new file mode 100644 index 000000000..6a8313b9b --- /dev/null +++ b/src/twinkle_agentic/condenser/facts.py @@ -0,0 +1,83 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from twinkle_agentic.condenser.base import Condenser + +if TYPE_CHECKING: + from twinkle.data_format import SamplingParams # noqa: F401 + from twinkle.sampler.base import Sampler # noqa: F401 + +_SECTION_SCHEMA = """You are a text compression assistant. A downstream model will read your compressed output to decide whether the detail it needs is inside this block; if yes, it will fetch and read the original passage. + +Downstream model workflow: +Read your compressed output -> Decide whether needed info is in this block -> If yes -> Fetch original. + +Therefore your compression MUST NOT lose major information from the source. + +Output format: + +```text +## Summary +Overview plus facts STRONGLY RELATED to the Query, stated explicitly. + +## More +A collapsed index; expansion required to see specific information. +``` + +Rules: +1. Telegraphic style — drop function words ("the", "a", "is", "are", "of", ...); colons and commas mean "is" / "has". + * Exception: KEEP role-tagging verb+preposition phrases verbatim ("published by X", "written by X", "directed by X", "starring X", "founded by X", "created by X", "composed by X", "produced by X", "based on X", "adapted from X"). Collapsing these to a bare name loses the relation role (author vs publisher vs director) that the downstream question may hinge on. +2. Summary MUST contain the passage's primary topic + 2–4 concrete core facts drawn from the source (entities, numbers, dates, relations). If a Query is given, order Query-relevant facts first, but STILL include other core facts within the budget. A Query is an ORDERING HINT, NOT a filter. +3. Summary MUST NOT be meta-commentary about the Query. Forbidden patterns: "no X mention", "Query info: absent", "passage covers Y only", "does not contain ...", "no relevant info", or summaries that are only abstract category words like "structure/order/usage" with no facts. If the passage is unrelated to the Query, you still summarize the passage normally. +4. More is an INDEX of category keywords, NOT inline data. Enumerate what CAN be recovered from the source (e.g. "birthplace, death place, age"); do NOT paste dates/numbers/names inline. Make sure all category of useful facts are introduced here. +5. Output language MUST match the source language. +6. Do NOT fabricate. Do NOT omit major information. Any fact not in the source MUST NOT appear in your output. + +Now begin. +""" # noqa + +_SECTION_USER_TEMPLATE = """\ +Downstream model will read your compressed block to decide whether to \ +expand it. Compress faithfully: preserve the passage topic + core facts. \ +Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary \ +about the Query (never write "Query info: absent", "no X mention", etc.); \ +if the passage does not address the Query, still summarize the passage. + +## Query (ordering hint only — still summarize the whole passage) +{query} + +## Target length +Compress AS MUCH AS faithfully possible. HARD CEILING: {budget} chars. \ +If core facts fit in far fewer chars, output fewer. \ +Never exceed the ceiling. + +## Passage +{text}""" + + +class FactsCondenser(Condenser): + + def __init__( + self, + sampler: Sampler, + compression_ratio: float = 2.0, + *, + model_path: str = '', + sampling_params: SamplingParams | None = None, + min_budget_chars: int = 250, + template: Any | None = None, + lora_path: str | None = None, + ): + super().__init__( + sampler, + compression_ratio, + model_path=model_path, + sampling_params=sampling_params, + system_prompt=_SECTION_SCHEMA, + user_prompt_template=_SECTION_USER_TEMPLATE, + min_budget_chars=min_budget_chars, + template=template, + lora_path=lora_path, + ) diff --git a/src/twinkle_agentic/condenser/keyword.py b/src/twinkle_agentic/condenser/keyword.py deleted file mode 100644 index e17c3ca7c..000000000 --- a/src/twinkle_agentic/condenser/keyword.py +++ /dev/null @@ -1,486 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -import math -import re -import threading -from typing import Any, Dict, FrozenSet, List, Optional, Sequence, Tuple - -from twinkle_agentic.condenser.base import Condenser -from twinkle_agentic.data_format import Chunk, Chunks - -# --------------------------------------------------------------------------- -# spaCy lazy loader (one model per process, thread-safe) -# --------------------------------------------------------------------------- -_SPACY_MODELS: dict[str, Any] = {} -_SPACY_LOCK = threading.Lock() - - -def _load_spacy(name: str): - nlp = _SPACY_MODELS.get(name) - if nlp is not None: - return nlp - with _SPACY_LOCK: - nlp = _SPACY_MODELS.get(name) - if nlp is not None: - return nlp - try: - import spacy - except ImportError as e: - raise ImportError('KeywordCondenser requires spaCy. Install with: ' - '`pip install spacy && python -m spacy download en_core_web_sm`') from e - try: - nlp = spacy.load(name) - except OSError as e: - raise OSError(f'spaCy model {name!r} not found. Download with: ' - f'`python -m spacy download {name}`') from e - _SPACY_MODELS[name] = nlp - return nlp - - -# --------------------------------------------------------------------------- -# configuration-free constants -# --------------------------------------------------------------------------- -# Entity labels dropped from keyword candidates (low recall value). -_DROP_ENT_LABELS: frozenset[str] = frozenset({'CARDINAL', 'ORDINAL', 'PERCENT', 'QUANTITY'}) - -# Dependency labels that introduce sub-clauses / conjuncts we do NOT want -# to pull into a single noun-phrase span. -_DROP_NP_DEPS: frozenset[str] = frozenset( - {'relcl', 'acl', 'advcl', 'ccomp', 'xcomp', 'conj', 'cc', 'appos', 'parataxis'}) - -# Tokens stripped from NP boundaries. -_LEADING_STRIP_POS: frozenset[str] = frozenset({'DET', 'PUNCT'}) - -# Tuple-slot separator. ``|`` avoids confusion when a slot itself -# contains a comma (e.g. ``"London, England"``). -_SLOT_SEP = ' | ' -_TRIPLE_SEP = '; ' - -_WORD_RE = re.compile(r'\w+', flags=re.UNICODE) - - -# --------------------------------------------------------------------------- -# NP / verb surface helpers -# --------------------------------------------------------------------------- -def _np_text(head) -> str: - """Return the noun-phrase text headed by ``head``. - - Keeps the contiguous span from the leftmost to the rightmost kept - token so internal punctuation (hyphens, apostrophes, slashes) is - preserved verbatim. Drops clausal / conjunct sub-trees and trims - leading determiners / possessive pronouns. - """ - # Collect subtree tokens, cutting off whole clausal children. - collected: list = [] - - def _walk(tok): - if tok is not head and tok.dep_ in _DROP_NP_DEPS: - return - collected.append(tok) - for child in tok.children: - _walk(child) - - _walk(head) - if not collected: - return head.text - collected.sort(key=lambda t: t.i) - - # Strip leading det/punct and possessive pronouns. - while collected and (collected[0].pos_ in _LEADING_STRIP_POS or - (collected[0].pos_ == 'PRON' and collected[0].dep_ == 'poss')): - collected.pop(0) - while collected and collected[-1].pos_ == 'PUNCT': - collected.pop() - if not collected: - return head.text - - start, end = collected[0].i, collected[-1].i + 1 - # If the kept tokens form a contiguous span, use the original text - # (preserves hyphens etc.). Otherwise fall back to text_with_ws. - if end - start == len(collected): - return head.doc[start:end].text.strip() - return ''.join(t.text_with_ws for t in collected).strip() - - -def _verb_surface(verb_tok) -> str: - """Verb text including auxiliaries (``was born``, ``has been released``).""" - aux = [c for c in verb_tok.children if c.dep_ in ('aux', 'auxpass')] - if not aux: - return verb_tok.text - tokens = sorted(aux + [verb_tok], key=lambda t: t.i) - return ' '.join(t.text for t in tokens) - - -def _first_child(token, deps: Sequence[str]): - if token is None: - return None - for c in token.children: - if c.dep_ in deps: - return c - return None - - -def _strip_leading_nc(noun_chunk) -> str: - toks = list(noun_chunk) - while toks and (toks[0].pos_ in _LEADING_STRIP_POS or toks[0].pos_ == 'NUM' or - (toks[0].pos_ == 'PRON' and toks[0].tag_ in ('PRP$', 'WP$'))): - toks.pop(0) - while toks and toks[-1].pos_ == 'PUNCT': - toks.pop() - if not toks: - return '' - start, end = toks[0].i, toks[-1].i + 1 - if end - start == len(toks): - return noun_chunk.doc[start:end].text.strip() - return ''.join(t.text_with_ws for t in toks).strip() - - -def _word_tokens_lower(text: str) -> frozenset[str]: - return frozenset(m.group(0).lower() for m in _WORD_RE.finditer(text)) - - -def _word_boundary_truncate(text: str, limit: int) -> str: - """Truncate ``text`` to ``limit`` chars at the nearest space.""" - if len(text) <= limit: - return text - cut = text[:limit] - sp = cut.rfind(' ') - trimmed = cut[:sp] if sp >= limit // 2 else cut - return trimmed.rstrip() or cut - - -# --------------------------------------------------------------------------- -# extraction (pure functions on spaCy Doc) -# --------------------------------------------------------------------------- -def _extract_opening(doc, max_chars: int) -> str: - """First non-empty sentence, word-boundary-truncated to ``max_chars``.""" - if max_chars <= 0: - return '' - for sent in doc.sents: - text = sent.text.strip() - if text: - return _word_boundary_truncate(text, max_chars) - return '' - - -def _extract_triples(doc, n: int) -> list[tuple[str, ...]]: - """Subject-verb-object (+ optional prep-obj) triples. - - - Skips pronoun subjects (unresolved coreference is noise). - - Preserves verb surface form (``was born`` rather than ``bear``). - - Deduplicates on lemmas. - """ - if n <= 0: - return [] - out: list[tuple[str, ...]] = [] - seen: set = set() - for sent in doc.sents: - for verb in sent: - if verb.pos_ not in ('VERB', 'AUX'): - continue - subj = _first_child(verb, ('nsubj', 'nsubjpass', 'csubj')) - if subj is None or subj.pos_ == 'PRON': - continue - obj = _first_child(verb, ('dobj', 'attr', 'oprd')) - prep = _first_child(verb, ('prep', )) - prep_obj = _first_child(prep, ('pobj', 'pcomp')) if prep is not None else None - - subj_txt = _np_text(subj) - verb_txt = _verb_surface(verb) - - if obj is not None and prep_obj is not None: - triple = (subj_txt, verb_txt, _np_text(obj), f'{prep.text} {_np_text(prep_obj)}') - key = (subj.lemma_.lower(), verb.lemma_.lower(), obj.lemma_.lower(), - f'{prep.text.lower()} {prep_obj.lemma_.lower()}') - elif obj is not None: - triple = (subj_txt, verb_txt, _np_text(obj)) - key = (subj.lemma_.lower(), verb.lemma_.lower(), obj.lemma_.lower()) - elif prep_obj is not None: - triple = (subj_txt, f'{verb_txt} {prep.text}', _np_text(prep_obj)) - key = (subj.lemma_.lower(), f'{verb.lemma_.lower()} {prep.text.lower()}', prep_obj.lemma_.lower()) - else: - continue - if key in seen: - continue - seen.add(key) - out.append(triple) - if len(out) >= n: - return out - return out - - -def _extract_keywords(doc, k: int, excluded_tokens: frozenset[str]) -> list[str]: - """Rank keyword candidates by (entity-weighted) frequency. - - - Drops pure-numeric entities (CARDINAL / ORDINAL / PERCENT / QUANTITY). - - Skips any term whose words are all already in ``excluded_tokens`` - (so we don't repeat what the opening already says). - - Subsumption dedup: drops a shorter form if a longer form - containing it is already kept (``"Nolan"`` dropped when - ``"Christopher Nolan"`` is present). - """ - if k <= 0: - return [] - counts: dict[str, float] = {} - order: dict[str, int] = {} - idx = 0 - - def _add(term: str, weight: float) -> None: - nonlocal idx - t = term.strip() - if len(t) < 2: - return - words = [w.lower() for w in _WORD_RE.findall(t)] - if not words: - return - if all(w in excluded_tokens for w in words): - return - if t not in order: - order[t] = idx - idx += 1 - counts[t] = counts.get(t, 0.0) + weight - - for ent in doc.ents: - if ent.label_ in _DROP_ENT_LABELS: - continue - _add(ent.text, weight=10.0) - for nc in doc.noun_chunks: - _add(_strip_leading_nc(nc), weight=1.0) - for tok in doc: - if tok.pos_ == 'PROPN' and not tok.is_stop: - _add(tok.text, weight=2.0) - - ranked = sorted(counts.keys(), key=lambda t: (-counts[t], order[t])) - - kept: list[str] = [] - kept_word_sets: list[frozenset[str]] = [] - for term in ranked: - words = frozenset(_WORD_RE.findall(term.lower())) - # Subsumed by any already-kept term (identical or proper subset). - if any(words == ws or words < ws for ws in kept_word_sets): - continue - # Also drop earlier-kept strict subsets of the current term. - to_remove = [i for i, ws in enumerate(kept_word_sets) if ws < words] - for i in reversed(to_remove): - kept.pop(i) - kept_word_sets.pop(i) - kept.append(term) - kept_word_sets.append(words) - if len(kept) >= k: - break - return kept - - -# --------------------------------------------------------------------------- -# budget-aware formatting (pure strings) -# --------------------------------------------------------------------------- -def _format_triple(triple: tuple[str, ...]) -> str: - return '(' + _SLOT_SEP.join(triple) + ')' - - -def _compose(opening: str, rel: str, kw: str) -> str: - parts: list[str] = [] - if opening: - parts.append(f'Open: {opening}') - if rel: - parts.append(f'Rel: {rel}') - if kw: - parts.append(f'More: {kw}') - return '\n'.join(parts) - - -def _fit_under_budget( - opening: str, - triples: list[tuple[str, ...]], - keywords: list[str], - budget: int, - *, - fallback_text: str = '', -) -> str: - """Pack as many triples + keywords as possible under ``budget``. - - Strategy: - 1. If opening alone is already too long, word-boundary truncate it. - 2. Greedily append triples one-by-one, keeping a running string. - 3. Greedily append keywords one-by-one on top of whatever fits. - 4. Never exceed ``budget`` — final safety clamp applies. - """ - # ----- opening ----- - if opening and len(f'Open: {opening}') > budget: - max_open = max(0, budget - len('Open: ')) - opening = _word_boundary_truncate(opening, max_open) if max_open else '' - - if not opening and not triples and not keywords: - # Nothing extractable — fall back to raw text, strict-truncated. - base = fallback_text[:budget] if fallback_text else '' - return _word_boundary_truncate(base, budget) if base else base - - current = _compose(opening, '', '') - if len(current) > budget: - return current[:budget] - - # ----- triples ----- - kept_triples: list[tuple[str, ...]] = [] - for t in triples: - trial_rel = _TRIPLE_SEP.join(_format_triple(x) for x in kept_triples + [t]) - trial = _compose(opening, trial_rel, '') - if len(trial) <= budget: - kept_triples.append(t) - else: - break - - rel_str = _TRIPLE_SEP.join(_format_triple(x) for x in kept_triples) - - # ----- keywords ----- - kept_kws: list[str] = [] - for k in keywords: - trial_kw = ', '.join(kept_kws + [k]) - trial = _compose(opening, rel_str, trial_kw) - if len(trial) <= budget: - kept_kws.append(k) - else: - break - - kw_str = ', '.join(kept_kws) - result = _compose(opening, rel_str, kw_str) - if not result: - # Budget too tight for any extracted slot — fall back to raw - # text truncated at a word boundary. - base = fallback_text[:budget] if fallback_text else '' - return _word_boundary_truncate(base, budget) if base else base - # Belt-and-braces: budget is strict. - return result if len(result) <= budget else result[:budget] - - -# --------------------------------------------------------------------------- -# KeywordCondenser -# --------------------------------------------------------------------------- -class KeywordCondenser(Condenser): - """Extractive, spaCy-driven passage condenser. - TODO: Experimental feature, wait for testing - - Args: - num_relations: Max number of - ``(subject, verb, object[, prep-obj])`` tuples per chunk. - Set to ``0`` to disable the ``Rel:`` slot. - max_first_sentence_chars: Hard cap for the opening slot, applied - before the global compression budget. - num_keywords: Max keyword items per chunk. ``0`` disables ``More:``. - compression_ratio: Target compression factor. Must be ``> 1``. - ``len(output) <= ceil(len(input) / compression_ratio)`` is - strictly enforced for every chunk that passes ``min_chars``. - spacy_model: spaCy pipeline name (default ``en_core_web_sm``). - min_chars: Pre-filter. Chunks shorter than this are passed - through **unchanged**; the ratio contract does not apply to - them. Set to ``0`` to always compress. - skip_roles: Roles whose chunks are never compressed. - rounds: Optional set/list of conversation-turn numbers to - compress. ``None`` (default) = no round-based filtering; - when provided, chunks whose ``round`` is not in this set - are passed through unchanged. Chunks that lack a ``round`` - field are also skipped when this filter is active. - - Every produced chunk is marked with ``raw.condensed=True`` so - :meth:`Chunks.to_trajectory` wraps it in ``...``. - - Example: - >>> from twinkle_agentic.chunker import NativeChunker - >>> from twinkle_agentic.condenser.keyword import KeywordCondenser - >>> chunker = NativeChunker(chunk_size=1024) - >>> cond = KeywordCondenser( - ... num_relations=3, max_first_sentence_chars=160, - ... num_keywords=8, compression_ratio=4.0) - >>> traj = {'messages': [{'role': 'user', 'content': long_passage}]} - >>> chunks = cond(chunker(traj)) - >>> traj_compressed = chunks.to_trajectory() - """ - - def __init__( - self, - num_relations: int = 3, - max_first_sentence_chars: int = 160, - num_keywords: int = 8, - compression_ratio: float = 4.0, - spacy_model: str = 'en_core_web_sm', - min_chars: int = 200, - skip_roles: Sequence[str] = ('system', 'tool', 'assistant'), - rounds: Sequence[int] | None = None, - ): - if num_relations < 0: - raise ValueError(f'num_relations must be >= 0, got {num_relations}') - if num_keywords < 0: - raise ValueError(f'num_keywords must be >= 0, got {num_keywords}') - if max_first_sentence_chars < 0: - raise ValueError(f'max_first_sentence_chars must be >= 0, got {max_first_sentence_chars}') - if compression_ratio <= 1.0: - raise ValueError(f'compression_ratio must be > 1, got {compression_ratio}') - if min_chars < 0: - raise ValueError(f'min_chars must be >= 0, got {min_chars}') - - self.num_relations = num_relations - self.max_first_sentence_chars = max_first_sentence_chars - self.num_keywords = num_keywords - self.compression_ratio = float(compression_ratio) - self.spacy_model = spacy_model - self.min_chars = min_chars - self.skip_roles = tuple(skip_roles) - self.rounds = set(rounds) if rounds is not None else None - - # ------------------------------------------------------------------ - def __call__(self, chunks: Chunks, **kwargs) -> Chunks: - nlp = _load_spacy(self.spacy_model) - out: list[Chunk] = [] - for c in chunks.chunks: - if not self._should_condense(c): - out.append(c) - continue - compressed = self._condense(c['content'], nlp) - out.append(self._mark_condensed(c, compressed)) - return Chunks(chunks=out) - - # ------------------------------------------------------------------ - # selection policy - # ------------------------------------------------------------------ - def _should_condense(self, chunk: Chunk) -> bool: - if chunk.get('type') != 'text': - return False - if chunk.get('role') in self.skip_roles: - return False - if self.rounds is not None and chunk.get('round') not in self.rounds: - return False - content = chunk.get('content') - if not isinstance(content, str) or not content: - return False - if len(content) < self.min_chars: - return False - raw = chunk.get('raw') or {} - if isinstance(raw, dict): - # Chunker-emitted reasoning / tool-call text chunks carry a - # non-empty ``kind`` marker; leave them alone. - if raw.get('kind'): - return False - # Idempotency — don't re-condense already condensed chunks. - if raw.get('condensed'): - return False - return True - - @staticmethod - def _mark_condensed(chunk: Chunk, content: str) -> Chunk: - new: dict[str, Any] = dict(chunk) - raw = dict(new.get('raw') or {}) - raw.setdefault('original', new.get('content', '')) - new['content'] = content - raw['condensed'] = True - new['raw'] = raw - return new # type: ignore[return-value] - - # ------------------------------------------------------------------ - # core extractive compression - # ------------------------------------------------------------------ - def _condense(self, text: str, nlp) -> str: - budget = max(1, math.ceil(len(text) / self.compression_ratio)) - doc = nlp(text) - opening = _extract_opening(doc, self.max_first_sentence_chars) - excluded = _word_tokens_lower(opening) - triples = _extract_triples(doc, self.num_relations) - keywords = _extract_keywords(doc, self.num_keywords, excluded) - return _fit_under_budget(opening, triples, keywords, budget, fallback_text=text) diff --git a/src/twinkle_agentic/condenser/model.py b/src/twinkle_agentic/condenser/model.py deleted file mode 100644 index 521d38063..000000000 --- a/src/twinkle_agentic/condenser/model.py +++ /dev/null @@ -1,508 +0,0 @@ -# Copyright (c) ModelScope Contributors. All rights reserved. -from __future__ import annotations - -import math -import re -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Tuple - -from twinkle_agentic.condenser.base import Condenser -from twinkle_agentic.data_format import Chunk, Chunks - -if TYPE_CHECKING: - from twinkle.data_format import SamplingParams, Trajectory # noqa: F401 - from twinkle.sampler.base import Sampler # noqa: F401 - -_SECTION_SCHEMA = """You are a text compression assistant. A downstream model will read your compressed output to decide whether the detail it needs is inside this block; if yes, it will fetch and read the original passage. - -Downstream model workflow: -Read your compressed output -> Decide whether needed info is in this block -> If yes -> Fetch original. - -Therefore your compression MUST NOT lose major information from the source. - -Output format: - -```text -## Summary -Overview plus facts STRONGLY RELATED to the Query, stated explicitly. - -## More -A collapsed index; expansion required to see specific information. -``` - -Rules: -1. Telegraphic style — drop function words ("the", "a", "is", "are", "of", ...); colons and commas mean "is" / "has". - * Exception: KEEP role-tagging verb+preposition phrases verbatim ("published by X", "written by X", "directed by X", "starring X", "founded by X", "created by X", "composed by X", "produced by X", "based on X", "adapted from X"). Collapsing these to a bare name loses the relation role (author vs publisher vs director) that the downstream question may hinge on. -2. Summary MUST contain the passage's primary topic + 2–4 concrete core facts drawn from the source (entities, numbers, dates, relations). If a Query is given, order Query-relevant facts first, but STILL include other core facts within the budget. A Query is an ORDERING HINT, NOT a filter. -3. Summary MUST NOT be meta-commentary about the Query. Forbidden patterns: "no X mention", "Query info: absent", "passage covers Y only", "does not contain ...", "no relevant info", or summaries that are only abstract category words like "structure/order/usage" with no facts. If the passage is unrelated to the Query, you still summarize the passage normally. -4. More is an INDEX of category keywords, NOT inline data. Enumerate what CAN be recovered from the source (e.g. "birthplace, death place, age"); do NOT paste dates/numbers/names inline. Make sure all category of useful facts are introduced here. -5. Output language MUST match the source language. -6. Do NOT fabricate. Do NOT omit major information. Any fact not in the source MUST NOT appear in your output. - -Example: - -Source: -```text -Marie Curie (7 Nov 1867 – 4 Jul 1934), born Maria Sklodowska in Warsaw (then Russian Poland); parents were teachers. Barred from Polish universities, she and her sister agreed to take turns funding each other's overseas study. - -In 1891 Marie reached Paris and enrolled at the Sorbonne, earning a physics degree (1893) and a mathematics degree (1894), becoming the school's first female physics lecturer. In 1895 she married French physicist Pierre Curie; they spent the rest of their lives on radioactivity research. - -In July 1898 she discovered polonium, named after her homeland Poland; in December she and Pierre announced the discovery of radium. She coined "radioactivity" and showed it is an atomic property, not a chemical reaction. - -In 1903 she shared the Nobel Prize in Physics with Pierre and Henri Becquerel. In 1911 she alone won the Nobel Prize in Chemistry for polonium and radium. She is the first woman to win a Nobel, and the only person to win Nobels in two different sciences. After Pierre died in a carriage accident in 1906, Marie took his chair and became the first female professor at the Sorbonne. - -During World War I she developed mobile X-ray units, called "Petites Curies" in French; about 20 were deployed to the front, examining over 1,000,000 wounded soldiers. - -She died of aplastic anaemia from radiation exposure on 4 July 1934 in Passy, Haute-Savoie, France, aged 66. Her notebooks remain highly radioactive, kept in lead boxes; researchers must wear protective gear to consult them. -``` - -Compressed: -```text -## Summary -Marie Curie: French-Polish physicist/chemist, founder of radioactivity research, first female Sorbonne professor. -- Nobel x2 (Physics + Chemistry); first woman Nobel laureate; only person with Nobels in two sciences. -- Discovered polonium + radium; coined "radioactivity"; proved it is an atomic property. - -## More -- birthplace, death place, age, cause of death -- degree years, in-school firsts x2 -- element naming origin, collaborators, full timeline -- Nobel year per prize, co-laureates, citation -- device name, deployment scale, patients treated -- notebook radioactivity, storage, access conditions -``` - -Now begin. -""" # noqa - -DEFAULT_SYSTEM_PROMPT = _SECTION_SCHEMA - -DEFAULT_USER_PROMPT_TEMPLATE = """\ -Downstream model will read your compressed block to decide whether to \ -expand it. Compress faithfully: preserve the passage topic + core facts. \ -Do NOT invent facts. Do NOT drop major facts. Do NOT write meta-commentary \ -about the Query (never write "Query info: absent", "no X mention", etc.); \ -if the passage does not address the Query, still summarize the passage. - -## Query (ordering hint only — still summarize the whole passage) -{query} - -## Target length -Compress AS MUCH AS faithfully possible. HARD CEILING: {budget} chars. \ -If core facts fit in far fewer chars, output fewer. \ -Never exceed the ceiling. - -## Passage -{text}""" - -# A (chunk_index, chunk, char_budget) triple marking one compression job. -_Job = Tuple[int, Chunk, int] - - -# --------------------------------------------------------------------------- -# ModelCondenser -# --------------------------------------------------------------------------- -class ModelCondenser(Condenser): - """Compressor that delegates summarization to an LLM via a :class:`Sampler`. - TODO: Experimental feature, wait for testing - Args: - sampler: Configured :class:`Sampler` with a template set. - compression_ratio: Target factor (> 1). Used only to derive a - soft character budget passed into the prompt and to size - ``SamplingParams.max_tokens``. Model output is NOT hard - truncated; a chunk whose decoded output is not strictly - shorter than the original passage is left unchanged (and - not flagged ``raw.condensed``). - sampling_params: Override for per-call sampling; when ``None`` a - greedy config is derived from the max budget in the batch. - system_prompt: Override for the system prompt. Used verbatim. - user_prompt_template: Override the user prompt. Must contain - ``{budget}`` and ``{text}``. ``{query}`` is optional and is - replaced with the trajectory's question extracted by the - ``related_query`` callback (see below); jobs without a - detected query get a neutral placeholder. - min_chars: Pre-filter; chunks shorter than this pass through. - min_budget_chars: Floor for the soft character budget exposed - to the prompt. When ``ceil(len / compression_ratio)`` falls - below this, the budget is raised to this floor so short - passages keep room for all three sections in the model's - plan. Since the condenser no longer hard-clips output, - this only influences prompt wording and sampling token - limits; pass ``1`` to use the raw ratio everywhere. - template: Optional :class:`Template`. When provided, its - ``tokenizer.all_special_tokens`` are stripped from every - decoded response before length-clamping, preventing - protocol tokens (``<|im_end|>``, ``<|eot_id|>``, ````, - ...) from leaking into the compressed output. When - omitted, falls back to ``sampler.template`` if available. - skip_roles: Roles whose chunks are never compressed. - skip_pattern: Optional regex (compiled with ``re.MULTILINE``). - Any chunk whose ``content`` has a match for this pattern - is passed through unchanged, regardless of length / ratio. - Uses :func:`re.search` semantics, so anchor with ``^`` / - start-of-string if you want boundary-matching only (e.g. - ``r'^Question:'`` to preserve the question prefix in a - HotpotQA-style user message). ``None`` disables the filter. - This flag is purely a compression-skip filter; query - extraction is the orthogonal job of ``related_query``. - related_query: Optional ``(chunk) -> Optional[str]`` callback - that returns the query string carried by ``chunk`` (e.g. - the user's HotpotQA question), or ``None`` if the chunk - is not a query carrier. Walked in chunk order; the most - recently returned non-``None`` query is broadcast to all - subsequent condense-eligible chunks until the next hit. - Because :class:`MultiTurnCondenseRollout` may merge - multiple trajectories into one chunk list, each - trajectory's question chunk must precede its passages so - this rolling state correctly partitions queries - per-trajectory. ``None`` disables query injection (the - ``{query}`` slot collapses to a neutral placeholder). - rounds: Optional set of conversation turn indices to compress. - ``None`` = no round-based filter; chunks lacking a ``round`` - field are skipped when this filter is active. - batch_size: Max chunks per sampler call. Partial batches are - padded with a duplicate of the last trajectory so that - distributed samplers (DP slice) always receive a full batch. - lora_path: Optional LoRA adapter to use for compression. - - ``None`` (default): forwards ``use_base_model=True`` to - :meth:`Sampler.sample` so compression bypasses any - currently-synced LoRA — strongly recommended when the - sampler is also the training policy. - - ``str``: forwards ``adapter_path=lora_path`` so a - dedicated condenser LoRA (e.g. a ModelScope slug or - local directory) is loaded and used instead of the base. - - Compressed chunks are flagged ``raw.condensed=True``; a subsequent - :meth:`Chunks.to_trajectory` call wraps them in ````. - - Example:: - - >>> from twinkle.sampler import vLLMSampler - >>> sampler = vLLMSampler(model_id='Qwen/Qwen2.5-3B-Instruct', - ... engine_args={'dtype': 'bfloat16'}) - >>> sampler.set_template('qwen2_5') - >>> cond = ModelCondenser(sampler, compression_ratio=2.0) - >>> compressed = cond(chunks) - """ - - def __init__( - self, - sampler: Sampler, - compression_ratio: float = 2.0, - *, - sampling_params: SamplingParams | None = None, - system_prompt: str | None = None, - user_prompt_template: str | None = None, - min_chars: int = 200, - min_budget_chars: int = 250, - template: Any | None = None, - skip_roles: Sequence[str] = ('system', 'tool', 'assistant'), - skip_pattern: str | None = None, - related_query: Callable[[Chunk], str | None] | None = None, - rounds: Sequence[int] | None = None, - batch_size: int = None, - lora_path: str | None = None, - ): - if sampler is None: - raise ValueError('sampler is required') - if compression_ratio <= 1.0: - raise ValueError(f'compression_ratio must be > 1, got {compression_ratio}') - if min_chars < 0: - raise ValueError(f'min_chars must be >= 0, got {min_chars}') - if min_budget_chars < 1: - raise ValueError(f'min_budget_chars must be >= 1, got {min_budget_chars}') - if batch_size is not None and batch_size <= 0: - raise ValueError(f'batch_size must be >= 1, got {batch_size}') - - tpl = user_prompt_template or DEFAULT_USER_PROMPT_TEMPLATE - if '{budget}' not in tpl or '{text}' not in tpl: - raise ValueError('user_prompt_template must contain both {budget} and {text}') - - self.sampler = sampler - self.compression_ratio = float(compression_ratio) - self.sampling_params = sampling_params - self.system_prompt = system_prompt or DEFAULT_SYSTEM_PROMPT - self.user_prompt_template = tpl - self.min_chars = min_chars - self.min_budget_chars = int(min_budget_chars) - self.template = template - self.skip_roles = tuple(skip_roles) - # ``^`` must anchor to start-of-string, not start-of-line: a passage - # whose body contains a ``Question:`` line would otherwise skip compression. - self.skip_re: re.Pattern | None = (re.compile(skip_pattern) if skip_pattern else None) - self.related_query = related_query - self.rounds = set(rounds) if rounds is not None else None - self.batch_size = batch_size - self.lora_path = lora_path if lora_path else None - self._special_tokens_cache: tuple[str, ...] | None = None - - # ------------------------------------------------------------------ - # entry point - # ------------------------------------------------------------------ - def __call__(self, chunks: Chunks, **_kwargs: Any) -> Chunks: - out: list[Chunk] = list(chunks.chunks) - items = self._collect_jobs(out) - if not items: - return Chunks(chunks=out) - - batch_size = self.batch_size or len(items) - for start in range(0, len(items), batch_size): - sub = items[start:start + batch_size] - batch = [job for job, _q in sub] - queries = [q for _job, q in sub] - responses = self._sample_batch(batch, queries=queries) - for (idx, chunk, _budget), resp in zip(batch, responses): - text = self._postprocess(_decoded(resp), chunk['content']) - if text is None: - continue - out[idx] = _mark_condensed(chunk, text) - return Chunks(chunks=out) - - # ------------------------------------------------------------------ - # eligibility + job collection - # ------------------------------------------------------------------ - def _collect_jobs( - self, - chunks: Sequence[Chunk], - ) -> list[tuple[_Job, str | None]]: - """Collect compression jobs, tagging each with its trajectory's query. - - Walks ``chunks`` in order and maintains a rolling - ``current_query`` state driven by the ``related_query`` - callback: every chunk for which the callback returns a - non-``None`` string updates the state, and every subsequent - condense-eligible chunk picks up the most recent query. - Because the chunker emits each trajectory's question chunk - before its passages, this walk correctly partitions queries - per-trajectory even when ``MultiTurnCondenseRollout`` merges - multiple trajectories into a single chunk list — A's - passages only ever see A's question, B's only B's. - """ - items: list[tuple[_Job, str | None]] = [] - current_query: str | None = None - extract = self.related_query - for i, c in enumerate(chunks): - content = c.get('content') - if extract is not None: - q = extract(c) - if isinstance(q, str) and q: - current_query = q - if not self._should_condense(c): - continue - budget = max(self.min_budget_chars, math.ceil(len(content) / self.compression_ratio)) - if budget >= len(content): - continue - items.append(((i, c, max(1, budget)), current_query)) - return items - - def _should_condense(self, chunk: Chunk) -> bool: - if chunk.get('type') != 'text': - return False - if chunk.get('role') in self.skip_roles: - return False - if self.rounds is not None and chunk.get('round') not in self.rounds: - return False - content = chunk.get('content') - if not isinstance(content, str) or len(content) < self.min_chars: - return False - if self.skip_re is not None and self.skip_re.search(content): - return False - raw = chunk.get('raw') or {} - if isinstance(raw, dict): - # Skip chunker-emitted reasoning / tool_call text chunks. - if raw.get('kind'): - return False - # Idempotent — never re-compress something already compressed. - if raw.get('condensed'): - return False - return True - - # ------------------------------------------------------------------ - # batched sampling - # ------------------------------------------------------------------ - def _sample_batch( - self, - batch: Sequence[_Job], - *, - queries: Sequence[str | None] = (), - ) -> list[Any]: - """Dispatch one batch to the sampler, padded to ``batch_size``. - - Distributed samplers slice inputs across DP workers and can - mis-behave when the final batch is smaller than ``batch_size``; - we pad with a duplicate of the last trajectory and trim the - matching extra responses here. - - ``queries`` is aligned 1:1 with ``batch``; each per-job query - is injected into the user prompt's ``{query}`` slot. When - empty or ``None`` at an index, a neutral placeholder is used. - """ - qs: list[str | None] = list(queries) if queries else [None] * len(batch) - if len(qs) != len(batch): - raise ValueError(f'queries length ({len(qs)}) must match batch length ' - f'({len(batch)})') - trajectories = [ - self._build_trajectory(chunk['content'], budget, query=q) for (_, chunk, budget), q in zip(batch, qs) - ] - actual = len(trajectories) - device_mesh = getattr(self.sampler, 'device_mesh', None) - min_batch_size = (device_mesh.data_world_size if device_mesh is not None else 1) - if actual < min_batch_size: - trajectories.extend([trajectories[-1]] * (min_batch_size - actual)) - - sp = self._sampling_params_for(max(b for _, _, b in batch)) - kwargs: dict[str, Any] = {'sampling_params': sp} - if self.lora_path is None: - kwargs['use_base_model'] = True - else: - kwargs['adapter_path'] = self.lora_path - responses = self.sampler.sample(trajectories, **kwargs) - # Coerce to list (some samplers may return tuples) and drop - # padding responses so downstream ``zip`` aligns with ``batch``. - return list(responses)[:actual] - - def _build_trajectory( - self, - text: str, - budget: int, - *, - query: str | None = None, - ) -> Trajectory: - system = self.system_prompt - user = self.user_prompt_template.replace('{budget}', str(budget)) - user = user.replace('{text}', text) - q_text = ( - query.strip() if isinstance(query, str) and query and query.strip() else - '(no explicit query; compress by general salience)') - user = user.replace('{query}', q_text) - return { # type: ignore[return-value] - 'messages': [ - {'role': 'system', 'content': system}, - {'role': 'user', 'content': user}, - ], - } - - def _sampling_params_for(self, budget: int) -> SamplingParams: - if self.sampling_params is not None: - return self.sampling_params - from twinkle.data_format.sampling import SamplingParams - - # CJK worst case ~2 tokens/char; budget is a soft char ceiling, not output truth. - max_new = max(512, budget * 3 + 128) - return SamplingParams(temperature=0.0, max_tokens=max_new) - - # ------------------------------------------------------------------ - # postprocess - # ------------------------------------------------------------------ - def _postprocess(self, raw: str, original: str) -> str | None: - """Return compressed text, or ``None`` to signal passthrough. - - ``None`` is returned when the decoded output is empty, - degenerate (markdown markers only, no alphanumerics), or its - character length is **not strictly shorter** than ``original`` - — in which case the model failed to produce a useful - compression and the caller should keep the original passage - verbatim (no ```` wrap, not marked ``raw.condensed``). - """ - text = _strip_special_tokens(_strip_code_fences(raw), self._get_special_tokens()).strip() - if not text or not _has_alnum(text): - return None - if len(text) >= len(original): - return None - return text - - def _get_special_tokens(self) -> tuple[str, ...]: - """Return protocol tokens to strip from decoded output (cached). - - Resolution order: - - 1. ``self.template.tokenizer`` — explicit template passed to - ``__init__``. Preferred in distributed setups where - ``sampler.template`` on the driver is a proxy and may be - ``None``. - 2. ``self.sampler.template.tokenizer`` — best-effort fallback - for single-process use. - 3. Empty tuple — no stripping (safe no-op). - - Uses ``tokenizer.all_special_tokens`` when available so the - full eos/bos/pad/unk/sep/cls/mask/additional set is covered - in one shot; this means ChatML (``<|im_end|>``), Llama - (``<|eot_id|>``), T5 (````) etc. are all handled without - per-model hard-coding. - """ - if self._special_tokens_cache is not None: - return self._special_tokens_cache - tpl = self.template or getattr(self.sampler, 'template', None) - tokenizer = getattr(tpl, 'tokenizer', None) if tpl is not None else None - tokens: list[str] = [] - if tokenizer is not None: - extras = getattr(tokenizer, 'all_special_tokens', None) or [] - if extras: - tokens.extend(t for t in extras if isinstance(t, str) and t and not t.isspace()) - else: - for attr in ('eos_token', 'pad_token', 'bos_token'): - t = getattr(tokenizer, attr, None) - if isinstance(t, str) and t: - tokens.append(t) - # Order-preserving dedupe. - self._special_tokens_cache = tuple(dict.fromkeys(tokens)) - return self._special_tokens_cache - - -# --------------------------------------------------------------------------- -# pure helpers -# --------------------------------------------------------------------------- -_CODE_FENCE_RE = re.compile(r'^```[a-zA-Z]*\s*\n(.*?)\n```\s*$', re.DOTALL) - - -def _decoded(response: Any) -> str: - """Extract the first decoded sequence, or ``''`` on empty/malformed input.""" - seqs = getattr(response, 'sequences', None) or [] - if not seqs: - return '' - return getattr(seqs[0], 'decoded', None) or '' - - -def _mark_condensed(chunk: Chunk, content: str) -> Chunk: - """Return a shallow copy of ``chunk`` with compressed ``content`` - and ``raw.condensed=True`` (preserving any original content under - ``raw.original`` so a future :class:`ExtractCondensed` call can - recover the full text). - """ - new: dict[str, Any] = dict(chunk) - raw = dict(new.get('raw') or {}) - raw.setdefault('original', new.get('content', '')) - raw['condensed'] = True - new['content'] = content - new['raw'] = raw - return new # type: ignore[return-value] - - -def _strip_code_fences(text: str) -> str: - """Unwrap a leading/trailing triple-backtick fence if present.""" - stripped = text.strip() - m = _CODE_FENCE_RE.match(stripped) - return m.group(1) if m else text - - -def _strip_special_tokens(text: str, tokens: Sequence[str]) -> str: - """Remove tokenizer special tokens that leaked through decode. - - ``tokens`` is typically ``tokenizer.all_special_tokens`` from the - template's tokenizer (see :meth:`ModelCondenser._get_special_tokens`). - Uses literal :meth:`str.replace` rather than a regex so we only - strip registered protocol markers and never legitimate passage - content that happens to look like ``<|...|>``. - """ - for tok in tokens: - if tok and tok in text: - text = text.replace(tok, '') - return text - - -def _has_alnum(text: str) -> bool: - """True iff ``text`` contains at least one alphanumeric character. - - Used to detect degenerate model outputs like ``'##'`` or ``'- '`` - that are pure markdown markers with no actual words. - """ - return any(ch.isalnum() for ch in text) diff --git a/src/twinkle_agentic/train/__init__.py b/src/twinkle_agentic/train/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/twinkle_agentic/train/base.py b/src/twinkle_agentic/train/base.py new file mode 100644 index 000000000..a0fd281e3 --- /dev/null +++ b/src/twinkle_agentic/train/base.py @@ -0,0 +1,6 @@ + + +class Trainer: + + def train(self): + pass \ No newline at end of file diff --git a/src/twinkle_agentic/train/cron.py b/src/twinkle_agentic/train/cron.py new file mode 100644 index 000000000..d118df8bc --- /dev/null +++ b/src/twinkle_agentic/train/cron.py @@ -0,0 +1,5 @@ + + +class CronTrainManager: + + pass \ No newline at end of file diff --git a/src/twinkle_agentic/utils/llm_backup.py b/src/twinkle_agentic/utils/llm_backup.py new file mode 100644 index 000000000..ce5199932 --- /dev/null +++ b/src/twinkle_agentic/utils/llm_backup.py @@ -0,0 +1,314 @@ +import functools +import hashlib +import inspect +import json +import os +import random +import threading +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + + +@dataclass +class EvalRecord: + """A single evaluation record comparing student and teacher outputs.""" + student_result: Any + teacher_result: Any + match: bool + + +@dataclass +class DistillationState: + """Per-key state tracking confidence, dataset, and call count.""" + confidence: Optional[float] = None + dataset: List[EvalRecord] = field(default_factory=list) + call_count: int = 0 + + +class DistillationRegistry: + """Thread-safe registry maintaining distillation state per unique key.""" + + def __init__(self): + self._states: Dict[str, DistillationState] = defaultdict(DistillationState) + self._lock = threading.Lock() + + def get_confidence(self, key: str) -> float: + """Get confidence for a key. Returns 0.0 if no data available.""" + with self._lock: + state = self._states[key] + if state.confidence is not None: + return state.confidence + if state.dataset: + state.confidence = self._compute_confidence(state.dataset) + return state.confidence + return 0.0 + + def increment_call(self, key: str) -> int: + with self._lock: + state = self._states[key] + state.call_count += 1 + return state.call_count + + def add_record(self, key: str, student_result: Any, teacher_result: Any, match: bool): + with self._lock: + state = self._states[key] + state.dataset.append(EvalRecord( + student_result=student_result, + teacher_result=teacher_result, + match=match, + )) + + def refresh_confidence(self, key: str) -> float: + with self._lock: + state = self._states[key] + if state.dataset: + state.confidence = self._compute_confidence(state.dataset) + else: + state.confidence = 0.0 + return state.confidence + + @staticmethod + def _compute_confidence(dataset: List[EvalRecord]) -> float: + if not dataset: + return 0.0 + matches = sum(1 for r in dataset if r.match) + return matches / len(dataset) + + +# --------------------------------------------------------------------------- +# Global state +# --------------------------------------------------------------------------- +_registry = DistillationRegistry() +_teacher_api = None +_teacher_lock = threading.Lock() + + +def _get_teacher_api(): + """Lazy-init global teacher API from environment variables. + + Env vars: + LLM_BACKUP_MODEL: Model name (default: "gpt-4o") + LLM_BACKUP_API_KEY: API key + LLM_BACKUP_BASE_URL: Base URL for OpenAI-compatible endpoint + """ + global _teacher_api + if _teacher_api is not None: + return _teacher_api + with _teacher_lock: + if _teacher_api is not None: + return _teacher_api + from twinkle_agentic.protocol.openai import OpenAI + _teacher_api = OpenAI( + model=os.environ.get('LLM_BACKUP_MODEL', 'gpt-4o'), + api_key=os.environ.get('LLM_BACKUP_API_KEY'), + base_url=os.environ.get('LLM_BACKUP_BASE_URL'), + ) + return _teacher_api + + +def _call_teacher(trajectory, sampling_params) -> str: + """Call teacher API and extract raw content string.""" + api = _get_teacher_api() + message = api(trajectory, sampling_params) + if isinstance(message, list): + message = message[0] + return message.get('content', '') if isinstance(message, dict) else '' + + +# --------------------------------------------------------------------------- +# Key building +# --------------------------------------------------------------------------- +def _build_key(func_name: str, args: tuple, kwargs: dict, + param_names: List[str], key_params: Sequence[str]) -> str: + """Build a unique key from specified parameter values.""" + key_parts = [func_name] + for i, name in enumerate(param_names): + if name in key_params: + if i < len(args): + key_parts.append(f"{name}={_serialize_value(args[i])}") + elif name in kwargs: + key_parts.append(f"{name}={_serialize_value(kwargs[name])}") + for name in key_params: + if name not in param_names[:len(args)] and name in kwargs: + if f"{name}={_serialize_value(kwargs[name])}" not in key_parts: + key_parts.append(f"{name}={_serialize_value(kwargs[name])}") + raw_key = "|".join(key_parts) + return hashlib.md5(raw_key.encode()).hexdigest() + + +def _serialize_value(value: Any) -> str: + try: + return json.dumps(value, sort_keys=True, default=str) + except (TypeError, ValueError): + return str(value) + + +def _extract_param(args: tuple, kwargs: dict, param_names: List[str], name: str) -> Any: + """Extract a named parameter from args/kwargs given the signature's param_names.""" + if name in kwargs: + return kwargs[name] + for i, pname in enumerate(param_names): + if pname == name and i < len(args): + return args[i] + return None + + +# --------------------------------------------------------------------------- +# Decorator +# --------------------------------------------------------------------------- +def llm_backup( + key_params: Sequence[str], + comparator: Optional[Callable[[Any, Any], bool]] = None, + sample_rate: float = 0.2, + refresh_env_var: str = "LLM_BACKUP_REFRESH_INTERVAL", + default_refresh_interval: int = 50, +): + """Decorator for progressive distillation from teacher API to student model. + + The decorated function is the STUDENT (local model sampling). The TEACHER + is a global OpenAI-compatible API constructed from environment variables. + + The decorated function MUST accept ``trajectory`` and ``sampling_params`` + as parameters (by name) and return a raw string. This ensures: + - Teacher and student receive identical inputs + - The dataset contains raw (trajectory, student_output, teacher_output) tuples + - No pre/post processing is included, making data directly trainable + + Routing logic: + - confidence% -> use student (decorated fn) + - Of those, sample_rate% also call teacher for comparison + - (1 - confidence)% -> use teacher API + - Always also call student for comparison + + Every N calls the confidence is recalculated from the comparison dataset. + + Environment variables: + LLM_BACKUP_MODEL: Teacher model name (default "gpt-4o") + LLM_BACKUP_API_KEY: Teacher API key + LLM_BACKUP_BASE_URL: Teacher API base URL + LLM_BACKUP_REFRESH_INTERVAL: Confidence refresh interval N (default 50) + + Args: + key_params: Parameter names for unique confidence key (e.g. ["query"]). + comparator: function(student, teacher) -> bool. Default: equality. + sample_rate: Probability of teacher verification when student is used. + refresh_env_var: Env var name for refresh interval. + default_refresh_interval: Default refresh interval. + + Example: + >>> @llm_backup(key_params=["query"]) + ... def _sample(self, trajectory, sampling_params, query=None) -> str: + ... responses = self.sampler.sample([trajectory], ...) + ... return decode(responses[0]) + """ + if comparator is None: + comparator = lambda a, b: a == b # noqa: E731 + + def decorator(fn: Callable) -> Callable: + sig = inspect.signature(fn) + param_names = list(sig.parameters.keys()) + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + key = _build_key(fn.__qualname__, args, kwargs, param_names, key_params) + confidence = _registry.get_confidence(key) + + try: + refresh_interval = int(os.environ.get(refresh_env_var, default_refresh_interval)) + except (ValueError, TypeError): + refresh_interval = default_refresh_interval + + # Extract trajectory and sampling_params for teacher call + trajectory = _extract_param(args, kwargs, param_names, 'trajectory') + sampling_params = _extract_param(args, kwargs, param_names, 'sampling_params') + + roll = random.random() + use_student = roll < confidence + + if use_student: + # High confidence: trust student + result = fn(*args, **kwargs) + # Occasionally verify against teacher + if random.random() < sample_rate: + teacher_result = _call_teacher(trajectory, sampling_params) + match = comparator(result, teacher_result) + _registry.add_record(key, result, teacher_result, match) + if not match: + result = teacher_result + else: + # Low confidence: use teacher + teacher_result = _call_teacher(trajectory, sampling_params) + student_result = fn(*args, **kwargs) + match = comparator(student_result, teacher_result) + _registry.add_record(key, student_result, teacher_result, match) + result = teacher_result + + call_count = _registry.increment_call(key) + if refresh_interval > 0 and call_count % refresh_interval == 0: + _registry.refresh_confidence(key) + + return result + + wrapper._registry = _registry + return wrapper + + return decorator + + +def llm_backup_async( + key_params: Sequence[str], + comparator: Optional[Callable[[Any, Any], bool]] = None, + sample_rate: float = 0.2, + refresh_env_var: str = "LLM_BACKUP_REFRESH_INTERVAL", + default_refresh_interval: int = 50, +): + """Async version of llm_backup. Same semantics.""" + if comparator is None: + comparator = lambda a, b: a == b # noqa: E731 + + def decorator(fn: Callable) -> Callable: + sig = inspect.signature(fn) + param_names = list(sig.parameters.keys()) + + @functools.wraps(fn) + async def wrapper(*args, **kwargs): + key = _build_key(fn.__qualname__, args, kwargs, param_names, key_params) + confidence = _registry.get_confidence(key) + + try: + refresh_interval = int(os.environ.get(refresh_env_var, default_refresh_interval)) + except (ValueError, TypeError): + refresh_interval = default_refresh_interval + + trajectory = _extract_param(args, kwargs, param_names, 'trajectory') + sampling_params = _extract_param(args, kwargs, param_names, 'sampling_params') + + roll = random.random() + use_student = roll < confidence + + if use_student: + result = await fn(*args, **kwargs) + if random.random() < sample_rate: + teacher_result = _call_teacher(trajectory, sampling_params) + match = comparator(result, teacher_result) + _registry.add_record(key, result, teacher_result, match) + if not match: + result = teacher_result + else: + teacher_result = _call_teacher(trajectory, sampling_params) + student_result = await fn(*args, **kwargs) + match = comparator(student_result, teacher_result) + _registry.add_record(key, student_result, teacher_result, match) + result = teacher_result + + call_count = _registry.increment_call(key) + if refresh_interval > 0 and call_count % refresh_interval == 0: + _registry.refresh_confidence(key) + + return result + + wrapper._registry = _registry + return wrapper + + return decorator