lionagi 0.6.1__py3-none-any.whl → 0.7.1__py3-none-any.whl

lionagi/libs/token_transform/llmlingua.py

@@ -0,0 +1 @@
+# TODO

lionagi/libs/token_transform/perplexity.py

@@ -0,0 +1,439 @@
+import asyncio
+from dataclasses import dataclass
+from timeit import default_timer as timer
+
+import numpy as np
+from pydantic import BaseModel
+
+from lionagi.protocols.generic.event import EventStatus
+from lionagi.protocols.generic.log import Log
+from lionagi.service.endpoints.base import APICalling
+from lionagi.service.imodel import iModel
+from lionagi.utils import alcall, lcall, to_dict, to_list
+
+
+@dataclass
+class PerplexityScores:
+    """
+    Stores logprobs, tokens, and derived perplexity from a completion response.
+    """
+
+    completion_response: BaseModel | dict
+    original_tokens: list[str]
+    n_samples: int
+
+    @property
+    def logprobs(self) -> list[float]:
+        """Return list of logprobs extracted from the model response."""
+        return [i["logprob"] for i in self.perplexity_scores]
+
+    @property
+    def perplexity(self) -> float:
+        """
+        e^(mean logprob), if logprobs exist. Fallback to 1.0 if empty.
+        """
+        if not self.logprobs:
+            return 1.0
+        return np.exp(np.mean(self.logprobs))
+
+    @property
+    def perplexity_scores(self) -> list[dict]:
+        """
+        Return [{'token': ..., 'logprob': ...}, ...].
+        Handles two possible logprob structures:
+          - "tokens" + "token_logprobs"
+          - "content" (older style)
+        """
+        outs = []
+        try:
+            if isinstance(self.completion_response, dict):
+                log_prob = self.completion_response["choices"][0]["logprobs"]
+            else:
+                # Pydantic or other object
+                log_prob = self.completion_response.choices[0].logprobs
+        except Exception:
+            return outs
+
+        if not log_prob:
+            return outs
+
+        if "tokens" in log_prob and "token_logprobs" in log_prob:
+            # OpenAI style logprobs
+            for token, lp in zip(
+                log_prob["tokens"], log_prob["token_logprobs"]
+            ):
+                outs.append({"token": token, "logprob": lp})
+        elif "content" in log_prob:
+            # Old style logprobs
+            for item in log_prob["content"]:
+                outs.append(
+                    {"token": item["token"], "logprob": item["logprob"]}
+                )
+        return outs
+
+    def to_dict(self) -> dict:
+        """
+        Construct a dictionary representation, including perplexity, usage, etc.
+        """
+        # usage info
+        usage = {}
+        if isinstance(self.completion_response, dict):
+            usage = self.completion_response.get("usage", {})
+        else:
+            usage = to_dict(self.completion_response.usage)
+
+        return {
+            "perplexity": self.perplexity,
+            "original_tokens": self.original_tokens,
+            "prompt_tokens": usage.get("prompt_tokens", 0),
+            "completion_tokens": usage.get("completion_tokens", 0),
+            "total_tokens": usage.get("total_tokens", 0),
+        }
+
+    def to_log(self) -> Log:
+        """
+        Return a Log object for convenience.
+        """
+        return Log(content=self.to_dict())
+
+
+async def compute_perplexity(
+    chat_model: iModel,
+    initial_context: str = None,
+    tokens: list[str] = None,
+    system_msg: str = None,
+    n_samples: int = 1,
+    use_residue: bool = True,
+    **kwargs,
+) -> list[PerplexityScores]:
+    """
+    Splits tokens into n_samples chunks, calls the model with logprobs=True,
+    and returns PerplexityScores for each chunk.
+    """
+    context = initial_context or ""
+    n_samples = n_samples or len(tokens)
+
+    sample_token_len, residue = divmod(len(tokens), n_samples)
+    if n_samples == 1:
+        samples = [tokens]
+    else:
+        samples = [
+            tokens[: (i + 1) * sample_token_len] for i in range(n_samples)
+        ]
+        if use_residue and residue != 0:
+            samples.append(tokens[-residue:])
+
+    # Build text for each chunk
+    sampless = [context + " " + " ".join(s) for s in samples]
+    kwargs["logprobs"] = True
+
+    async def _inner(api_call: APICalling):
+        await api_call.invoke()
+        elapsed = 0
+        while (
+            api_call.status not in [EventStatus.COMPLETED, EventStatus.FAILED]
+            and elapsed < 5
+        ):
+            await asyncio.sleep(0.1)
+            elapsed += 0.1
+        return api_call.response
+
+    # Create and schedule calls
+    api_calls = []
+    for sample_txt in sampless:
+        messages = []
+        if system_msg:
+            if not chat_model.sequential_exchange:
+                messages.append({"role": "system", "content": system_msg})
+            messages.append({"role": "user", "content": sample_txt})
+        else:
+            messages.append({"role": "user", "content": sample_txt})
+
+        api_calls.append(
+            chat_model.create_api_calling(messages=messages, **kwargs)
+        )
+
+    results = await alcall(api_calls, _inner, max_concurrent=50)
+
+    def _pplx_score(input_):
+        idx, resp = input_
+        return PerplexityScores(resp, samples[idx], n_samples)
+
+    return lcall(enumerate(results), _pplx_score)
+
+
+class LLMCompressor:
+    """
+    Compress text by selecting segments with highest perplexity tokens
+    (or in practice, rank segments by logprob).
+    """
+
+    def __init__(
+        self,
+        chat_model: iModel,
+        system_msg=None,
+        tokenizer=None,
+        splitter=None,
+        target_ratio=0.2,
+        n_samples=5,
+        chunk_size=64,
+        max_tokens_per_sample=80,
+        min_pplx=0,
+        split_overlap=0,
+        split_threshold=0,
+        verbose=True,
+    ):
+        # Must have "logprobs" support
+        if "logprobs" not in chat_model.endpoint.acceptable_kwargs:
+            raise ValueError(
+                f"Model {chat_model.model_name} does not support logprobs. "
+                "Please use a model that supports logprobs."
+            )
+
+        self.chat_model = chat_model
+        self.tokenizer = tokenizer
+        self.splitter = splitter
+        self.system_msg = (
+            system_msg or "Concisely summarize content for storage:"
+        )
+        self.target_ratio = target_ratio
+        self.n_samples = n_samples
+        self.chunk_size = chunk_size
+        self.max_tokens_per_sample = max_tokens_per_sample
+        self.min_pplx = min_pplx
+        self.verbose = verbose
+        self.split_overlap = split_overlap
+        self.split_threshold = split_threshold
+
+    def tokenize(self, text: str, **kwargs) -> list[str]:
+        """
+        Tokenize text. If no custom tokenizer, use the default from lionagi.
+        """
+        if not self.tokenizer:
+            from lionagi.service.endpoints.token_calculator import (
+                TokenCalculator,
+            )
+
+            return TokenCalculator.tokenize(
+                text,
+                encoding_name=self.chat_model.model_name,
+                return_tokens=True,
+            )
+        if hasattr(self.tokenizer, "tokenize"):
+            return self.tokenizer.tokenize(text, **kwargs)
+        return self.tokenizer(text, **kwargs)
+
+    def split(
+        self,
+        text: str,
+        chunk_size=None,
+        overlap=None,
+        threshold=None,
+        by_chars=False,
+        return_tokens=False,
+        **kwargs,
+    ) -> list:
+        """
+        Split text into segments. If no custom splitter, default to chunk_content from lionagi.
+        """
+        if not self.splitter:
+            from lionagi.libs.file.chunk import chunk_content
+
+            contents = chunk_content(
+                content=text,
+                chunk_size=chunk_size or self.chunk_size,
+                overlap=overlap or self.split_overlap,
+                threshold=threshold or self.split_threshold,
+                return_tokens=return_tokens,
+                chunk_by="chars" if by_chars else "tokens",
+            )
+            return [i["chunk_content"] for i in contents]
+
+        # If user provided an object with .split or .chunk or .segment
+        for meth in ["split", "chunk", "segment"]:
+            if hasattr(self.splitter, meth):
+                return getattr(self.splitter, meth)(text, **kwargs)
+        raise ValueError(
+            "No valid method found in splitter: must have .split/.chunk/.segment"
+        )
+
+    async def rank_by_pplex(
+        self,
+        items: list,
+        initial_text=None,
+        cumulative=False,
+        n_samples=None,
+        use_residue=True,
+        **kwargs,
+    ) -> list:
+        """
+        Rank items (token lists or strings) by perplexity descending.
+        If cumulative=True, each item is appended to the context.
+        """
+
+        async def _get_item_perplexity(item):
+            # Ensure item is a list of tokens
+            item_toks = item if isinstance(item, list) else [item]
+            if len(item_toks) > self.max_tokens_per_sample:
+                item_toks = item_toks[: self.max_tokens_per_sample]
+            pplex_scores = await compute_perplexity(
+                chat_model=self.chat_model,
+                initial_context=initial_text,
+                tokens=item_toks,
+                n_samples=n_samples or self.n_samples,
+                system_msg=self.system_msg,
+                use_residue=use_residue,
+                **kwargs,
+            )
+            # Usually we only look at pplex_scores[0], as there's one chunk
+            return pplex_scores
+
+        # If user passed a single string, tokenize it
+        if isinstance(items, str):
+            items = self.tokenize(items)
+
+        if len(items) == 1:
+            single_scores = await _get_item_perplexity(items[0])
+            return [(items[0], single_scores[0])]
+
+        segments = []
+        if cumulative:
+            ctx = initial_text or ""
+            for i in items:
+                seg_toks = i if isinstance(i, list) else [i]
+                joined = " ".join(seg_toks)
+                ctx += " " + joined
+                segments.append(ctx)
+        else:
+            for i in items:
+                seg_toks = i if isinstance(i, list) else [i]
+                segments.append(" ".join(seg_toks))
+
+        tasks = [
+            asyncio.create_task(_get_item_perplexity(seg)) for seg in segments
+        ]
+        results = await asyncio.gather(*tasks)
+        # Pair each item with the first pplex (p[0]) if multiple were returned
+        pairs = [(itm, pplex[0]) for itm, pplex in zip(items, results)]
+
+        # Sort descending by perplexity
+        return sorted(pairs, key=lambda x: x[1].perplexity, reverse=True)
+
+    async def compress(
+        self,
+        text: str,
+        compression_ratio=None,
+        initial_text=None,
+        cumulative=False,
+        split_kwargs=None,
+        min_pplx=None,
+        **kwargs,
+    ) -> str:
+        """
+        Main method to compress text:
+        1) Split text
+        2) Rank by perplexity
+        3) Select best segments until reaching target ratio
+        """
+        start = timer()
+        if split_kwargs is None:
+            split_kwargs = {
+                "chunk_size": self.max_tokens_per_sample,
+                "overlap": self.split_overlap,
+                "threshold": self.split_threshold,
+                "return_tokens": True,
+            }
+
+        # Tokenize once to get total length
+        all_tokens = self.tokenize(text)
+        original_len = len(all_tokens)
+
+        # Split text
+        items = self.split(text, **split_kwargs)
+        # items -> list of token-lists
+
+        # Rank
+        ranked = await self.rank_by_pplex(
+            items=items,
+            initial_text=initial_text,
+            cumulative=cumulative,
+            **kwargs,
+        )
+
+        # Select
+        selected = self.select_by_pplex(
+            ranked_items=ranked,
+            target_compression_ratio=compression_ratio or self.target_ratio,
+            original_length=original_len,
+            min_pplx=min_pplx or self.min_pplx,
+        )
+
+        if self.verbose:
+            compressed_len = sum(
+                len(to_list(self.tokenize(x), dropna=True, flatten=True))
+                for x in selected
+            )
+            ratio = compressed_len / original_len if original_len else 1
+            print(
+                f"Original tokens: {original_len}\n"
+                f"Selected tokens: {compressed_len}\n"
+                f"Compression ratio: {ratio:.3f}\n"
+                f"Time: {timer() - start:.3f}s\n"
+            )
+
+        # Join final
+        out_str = " ".join(selected)
+        return out_str.strip()
+
+    def select_by_pplex(
+        self,
+        ranked_items: list,
+        target_compression_ratio: float,
+        original_length: int,
+        min_pplx=0,
+    ) -> list[str]:
+        """
+        From highest perplexity to lowest, pick items until we reach the desired ratio.
+        Items below min_pplx are skipped.
+        """
+        desired_len = int(original_length * target_compression_ratio)
+
+        chosen = []
+        current_len = 0
+        for item, info in ranked_items:
+            if info.perplexity > min_pplx:
+                if isinstance(item, list):
+                    item_toks = to_list(item, dropna=True, flatten=True)
+                else:
+                    item_toks = self.tokenize(item)
+                if current_len + len(item_toks) > desired_len:
+                    break
+                chosen.append(" ".join(item_toks))
+                current_len += len(item_toks)
+
+        return chosen
+
+
+# Helper function to quickly compress text using perplexity
+# (If you don't want to manually create LLMCompressor instance everywhere)
+async def compress_text(
+    text: str,
+    chat_model: iModel,
+    system_msg: str = None,
+    target_ratio: float = 0.2,
+    n_samples: int = 5,
+    max_tokens_per_sample=80,
+    verbose=True,
+) -> str:
+    """
+    Convenience function that instantiates LLMCompressor and compresses text.
+    """
+    compressor = LLMCompressor(
+        chat_model=chat_model,
+        system_msg=system_msg,
+        target_ratio=target_ratio,
+        n_samples=n_samples,
+        max_tokens_per_sample=max_tokens_per_sample,
+        verbose=verbose,
+    )
+    return await compressor.compress(text)