retab 0.0.42__py3-none-any.whl → 0.0.43__py3-none-any.whl

retab/_utils/display.py

@@ -1,440 +0,0 @@
-import base64
-import json
-from io import BytesIO
-from math import ceil
-from pathlib import Path
-from typing import List, Literal, Optional, TypedDict
-
-import numpy as np
-import requests
-import tiktoken  # For text tokenization
-from PIL import Image
-from rich.console import Console
-from rich.table import Table
-
-
-class TokenStats(TypedDict):
-    min: float
-    max: float
-    mean: float
-    median: float
-    p5: float
-    p95: float
-
-
-class TokenCounts(TypedDict):
-    input_text_tokens: int
-    output_text_tokens: int
-    input_image_tokens: int
-    output_image_tokens: int
-
-
-class MetricCategory(TypedDict):
-    num_examples: int
-    total_tokens: TokenStats
-    input_tokens: TokenStats
-    output_tokens: TokenStats
-    sum_total_tokens: float
-    sum_input_tokens: float
-    sum_output_tokens: float
-    num_examples_over_token_limit: int
-
-
-class Metrics(TypedDict):
-    Text: MetricCategory
-    Image: MetricCategory
-    Total: MetricCategory
-
-
-def count_text_tokens(content: str, encoding_name: str = "cl100k_base") -> int:
-    """
-    Count the number of tokens in a given text content using the specified encoding.
-    """
-    enc = tiktoken.get_encoding(encoding_name)
-    return len(enc.encode(content))
-
-
-def count_image_tokens(image_url: str, detail: Literal["low", "high", "auto"] = "high") -> int:
-    base_token_cost = 85  # cost for all images
-    token_per_tile = 170  # cost per 512×512 tile in high detail
-
-    # 1. Decide detail=low or detail=high
-    #    If detail=auto, figure out from user input or some heuristic
-    if detail == "low":
-        # 2. Low detail => always 85 tokens
-        return base_token_cost
-    else:
-        assert detail == "high" or detail == "auto"
-        # 3. High detail => 2-step scaling + tile-based cost
-
-        # (a) Get the raw image dimensions
-        try:
-            if image_url.startswith("data:image"):
-                header, encoded_data = image_url.split(",", 1)
-                image_data = base64.b64decode(encoded_data)
-                img = Image.open(BytesIO(image_data))
-            else:
-                # HTTP URL or local path
-                response = requests.get(image_url, timeout=5)
-                response.raise_for_status()
-                img = Image.open(BytesIO(response.content))
-
-            width, height = img.size
-        except Exception:
-            # If we fail to decode or fetch, maybe return the base cost
-            # plus one tile as a fallback
-            return base_token_cost + token_per_tile
-
-        # (b) Scale so neither dimension exceeds 2048
-        max_side = max(width, height)
-        if max_side > 2048:
-            scale_factor = 2048.0 / max_side
-            width = int(width * scale_factor)
-            height = int(height * scale_factor)
-
-        # (c) Upscale if shortest side < 768
-        min_side = min(width, height)
-        if min_side < 768:
-            upscale_factor = 768.0 / min_side
-            width = int(width * upscale_factor)
-            height = int(height * upscale_factor)
-
-        # (d) Count 512×512 tiles in the scaled image
-        tiles_wide = ceil(width / 512)
-        tiles_high = ceil(height / 512)
-        total_tiles = tiles_wide * tiles_high
-
-        return base_token_cost + (token_per_tile * total_tiles)
-
-
-def process_jsonl_file(jsonl_path: str) -> List[TokenCounts]:
-    """
-    Process a JSONL file and calculate the text and image tokens for each example.
-    Returns a list of dictionaries with token counts for system, user, assistant, and images.
-    """
-    results = []
-
-    with open(jsonl_path, "r", encoding="utf-8") as file:
-        for line in file:
-            example = json.loads(line)
-            input_text_tokens = 0
-            output_text_tokens = 0
-            input_image_tokens = 0
-            output_image_tokens = 0
-
-            for message in example.get("messages", []):
-                role = message.get("role")
-                content = message.get("content")
-
-                if isinstance(content, str):
-                    # Count text tokens based on role
-                    if role in ["developer", "system", "user"]:
-                        input_text_tokens += count_text_tokens(content)
-                    elif role == "assistant":
-                        output_text_tokens += count_text_tokens(content)
-
-                elif isinstance(content, list):  # Check for images in content
-                    for item in content:
-                        if item.get("type") == "image_url" and "image_url" in item:
-                            image_url = item["image_url"]["url"]
-                            tokens = count_image_tokens(image_url)
-                            if role in ["developer", "system", "user"]:
-                                input_image_tokens += tokens
-                            elif role == "assistant":
-                                output_image_tokens += tokens
-
-            results.append(
-                TokenCounts(
-                    input_text_tokens=input_text_tokens, output_text_tokens=output_text_tokens, input_image_tokens=input_image_tokens, output_image_tokens=output_image_tokens
-                )
-            )
-
-    return results
-
-
-def calculate_statistics(data: List[int]) -> TokenStats:
-    """
-    Calculate statistics for a list of numbers.
-    """
-    if not data:
-        return {"min": 0, "max": 0, "mean": 0, "median": 0, "p5": 0, "p95": 0}
-
-    return {
-        "min": float(min(data)),
-        "max": float(max(data)),
-        "mean": float(np.mean(data)),
-        "median": float(np.median(data)),
-        "p5": float(np.percentile(data, 5)),
-        "p95": float(np.percentile(data, 95)),
-    }
-
-
-def process_dataset_and_compute_metrics(jsonl_path: Path | str, token_limit: int = 128000) -> Metrics:
-    """
-    Process the dataset to compute metrics for Text, Image, and Total tokens.
-    """
-    # Initialize metrics
-    metrics: Metrics = {
-        "Text": MetricCategory(
-            num_examples=0,
-            total_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            input_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            output_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            sum_total_tokens=0,
-            sum_input_tokens=0,
-            sum_output_tokens=0,
-            num_examples_over_token_limit=0,
-        ),
-        "Image": MetricCategory(
-            num_examples=0,
-            total_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            input_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            output_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            sum_total_tokens=0,
-            sum_input_tokens=0,
-            sum_output_tokens=0,
-            num_examples_over_token_limit=0,
-        ),
-        "Total": MetricCategory(
-            num_examples=0,
-            total_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            input_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            output_tokens=TokenStats(min=0, max=0, mean=0, median=0, p5=0, p95=0),
-            sum_total_tokens=0,
-            sum_input_tokens=0,
-            sum_output_tokens=0,
-            num_examples_over_token_limit=0,
-        ),
-    }
-
-    # Accumulate token counts
-    input_text_tokens = []
-    output_text_tokens = []
-    messages_text_tokens = []
-
-    input_image_tokens = []
-    output_image_tokens = []
-    messages_image_tokens = []
-
-    input_total_tokens = []
-    output_total_tokens = []
-    messages_total_tokens = []
-
-    with open(jsonl_path, "r", encoding="utf-8") as file:
-        for line in file:
-            example = json.loads(line)
-
-            input_text_tokens_example = 0
-            output_text_tokens_example = 0
-
-            input_image_tokens_example = 0
-            output_image_tokens_example = 0
-
-            for message in example.get("messages", []):
-                role = message.get("role")
-                content = message.get("content")
-
-                if isinstance(content, str):
-                    if role in ["developer", "system", "user"]:
-                        input_text_tokens_example += count_text_tokens(content)
-                    elif role == "assistant":
-                        output_text_tokens_example += count_text_tokens(content)
-                elif isinstance(content, list):  # Handle images
-                    for item in content:
-                        if item.get("type") == "image_url" and "image_url" in item:
-                            image_url = item["image_url"]["url"]
-                            detail = item["image_url"]["detail"]
-                            tokens = count_image_tokens(image_url, detail)
-                            if role in ["developer", "system", "user"]:
-                                input_image_tokens_example += tokens
-                            elif role == "assistant":
-                                output_image_tokens_example += tokens
-
-                        elif item.get("type") == "text":
-                            if role in ["developer", "system", "user"]:
-                                input_text_tokens_example += count_text_tokens(item["text"])
-                            elif role == "assistant":
-                                output_text_tokens_example += count_text_tokens(item["text"])
-
-            # Calculate totals for the example
-            example_total_tokens = input_text_tokens_example + output_text_tokens_example + input_image_tokens_example + output_image_tokens_example
-
-            # Add to accumulators
-            input_text_tokens.append(input_text_tokens_example)
-            output_text_tokens.append(output_text_tokens_example)
-            messages_text_tokens.append(input_text_tokens_example + output_text_tokens_example)
-
-            input_image_tokens.append(input_image_tokens_example)
-            output_image_tokens.append(output_image_tokens_example)
-            messages_image_tokens.append(input_image_tokens_example + output_image_tokens_example)
-
-            input_total_tokens.append(input_text_tokens_example + input_image_tokens_example)
-            output_total_tokens.append(output_text_tokens_example + output_image_tokens_example)
-            messages_total_tokens.append(input_text_tokens_example + output_text_tokens_example + input_image_tokens_example + output_image_tokens_example)
-
-            # Count examples over token limit
-            if input_text_tokens_example > token_limit:
-                metrics["Text"]["num_examples_over_token_limit"] += 1
-            if input_image_tokens_example > token_limit:
-                metrics["Image"]["num_examples_over_token_limit"] += 1
-            if example_total_tokens > token_limit:
-                metrics["Total"]["num_examples_over_token_limit"] += 1
-                # print(example_total_tokens, token_limit)
-
-    # Update metrics for Text, Image, and Total
-    metrics["Text"]["num_examples"] = len(input_text_tokens)
-    metrics["Text"]["total_tokens"] = calculate_statistics(messages_text_tokens)
-    metrics["Text"]["input_tokens"] = calculate_statistics(input_text_tokens)
-    metrics["Text"]["output_tokens"] = calculate_statistics(output_text_tokens)
-    metrics["Text"]["sum_input_tokens"] = sum(input_text_tokens)
-    metrics["Text"]["sum_output_tokens"] = sum(output_text_tokens)
-    metrics["Text"]["sum_total_tokens"] = sum(messages_text_tokens)
-
-    metrics["Image"]["num_examples"] = len(input_image_tokens)
-    metrics["Image"]["total_tokens"] = calculate_statistics(messages_image_tokens)
-    metrics["Image"]["input_tokens"] = calculate_statistics(input_image_tokens)
-    metrics["Image"]["output_tokens"] = calculate_statistics(output_image_tokens)
-    metrics["Image"]["sum_input_tokens"] = sum(input_image_tokens)
-    metrics["Image"]["sum_output_tokens"] = sum(output_image_tokens)
-    metrics["Image"]["sum_total_tokens"] = sum(messages_image_tokens)
-
-    metrics["Total"]["num_examples"] = len(input_total_tokens)
-    metrics["Total"]["total_tokens"] = calculate_statistics(messages_total_tokens)
-    metrics["Total"]["input_tokens"] = calculate_statistics(input_total_tokens)
-    metrics["Total"]["output_tokens"] = calculate_statistics(output_total_tokens)
-    metrics["Total"]["sum_input_tokens"] = sum(input_total_tokens)
-    metrics["Total"]["sum_output_tokens"] = sum(output_total_tokens)
-    metrics["Total"]["sum_total_tokens"] = sum(messages_total_tokens)
-
-    return metrics
-
-
-def display_metrics(metrics: Metrics, input_token_price: Optional[float] = None, output_token_price: Optional[float] = None) -> None:
-    """
-    Display the metrics dictionary in a compact table with min/max, mean/median, and p5/p95 on the same row.
-    """
-    console = Console(style="on grey23")
-    table = Table(title="Dataset Metrics", show_lines=True)
-
-    # Add columns
-    table.add_column("Metric", justify="left", style="#BDE8F6", no_wrap=True)
-    table.add_column("Text", justify="right", style="#C2BDF6")
-    table.add_column("Image", justify="right", style="#F6BDBD")
-    table.add_column("Total", justify="right", style="#F6E4BD")
-
-    # Add rows
-    table.add_row("Num Examples", str(metrics["Text"]["num_examples"]), str(metrics["Image"]["num_examples"]), str(metrics["Total"]["num_examples"]))
-
-    table.add_row(
-        "Examples Over Limit",
-        str(metrics["Text"]["num_examples_over_token_limit"]),
-        str(metrics["Image"]["num_examples_over_token_limit"]),
-        str(metrics["Total"]["num_examples_over_token_limit"]),
-    )
-
-    table.add_row("")
-
-    # Rows for input tokens
-    table.add_row(
-        "Min / Max Input Tokens",
-        f"{metrics['Text']['input_tokens']['min']:.0f} / {metrics['Text']['input_tokens']['max']:.0f}",
-        f"{metrics['Image']['input_tokens']['min']:.0f} / {metrics['Image']['input_tokens']['max']:.0f}",
-        f"{metrics['Total']['input_tokens']['min']:.0f} / {metrics['Total']['input_tokens']['max']:.0f}",
-    )
-
-    table.add_row(
-        "Mean / Median Input Tokens",
-        f"{metrics['Text']['input_tokens']['mean']:.0f} / {metrics['Text']['input_tokens']['median']:.0f}",
-        f"{metrics['Image']['input_tokens']['mean']:.0f} / {metrics['Image']['input_tokens']['median']:.0f}",
-        f"{metrics['Total']['input_tokens']['mean']:.0f} / {metrics['Total']['input_tokens']['median']:.0f}",
-    )
-
-    table.add_row(
-        "P5 / P95 Input Tokens",
-        f"{metrics['Text']['input_tokens']['p5']:.0f} / {metrics['Text']['input_tokens']['p95']:.0f}",
-        f"{metrics['Image']['input_tokens']['p5']:.0f} / {metrics['Image']['input_tokens']['p95']:.0f}",
-        f"{metrics['Total']['input_tokens']['p5']:.0f} / {metrics['Total']['input_tokens']['p95']:.0f}",
-    )
-
-    table.add_row("Sum Input Tokens", f"{metrics['Text']['sum_input_tokens']}", f"{metrics['Image']['sum_input_tokens']}", f"{metrics['Total']['sum_input_tokens']}")
-
-    table.add_row("")  # Empty row for spacing
-
-    # Rows for output tokens
-    table.add_row(
-        "Min / Max Output Tokens",
-        f"{metrics['Text']['output_tokens']['min']:.0f} / {metrics['Text']['output_tokens']['max']:.0f}",
-        f"{metrics['Image']['output_tokens']['min']:.0f} / {metrics['Image']['output_tokens']['max']:.0f}",
-        f"{metrics['Total']['output_tokens']['min']:.0f} / {metrics['Total']['output_tokens']['max']:.0f}",
-    )
-
-    table.add_row(
-        "Mean / Median Output Tokens",
-        f"{metrics['Text']['output_tokens']['mean']:.0f} / {metrics['Text']['output_tokens']['median']:.0f}",
-        f"{metrics['Image']['output_tokens']['mean']:.0f} / {metrics['Image']['output_tokens']['median']:.0f}",
-        f"{metrics['Total']['output_tokens']['mean']:.0f} / {metrics['Total']['output_tokens']['median']:.0f}",
-    )
-
-    table.add_row(
-        "P5 / P95 Output Tokens",
-        f"{metrics['Text']['output_tokens']['p5']:.0f} / {metrics['Text']['output_tokens']['p95']:.0f}",
-        f"{metrics['Image']['output_tokens']['p5']:.0f} / {metrics['Image']['output_tokens']['p95']:.0f}",
-        f"{metrics['Total']['output_tokens']['p5']:.0f} / {metrics['Total']['output_tokens']['p95']:.0f}",
-    )
-
-    table.add_row("Sum Output Tokens", f"{metrics['Text']['sum_output_tokens']}", f"{metrics['Image']['sum_output_tokens']}", f"{metrics['Total']['sum_output_tokens']}")
-
-    table.add_row("")  # Empty row for spacing
-
-    # Total tokens
-    table.add_row(
-        "Min / Max Tokens",
-        f"{metrics['Text']['input_tokens']['min']:.0f} / {metrics['Text']['input_tokens']['max']:.0f}",
-        f"{metrics['Image']['input_tokens']['min']:.0f} / {metrics['Image']['input_tokens']['max']:.0f}",
-        f"{metrics['Total']['input_tokens']['min']:.0f} / {metrics['Total']['input_tokens']['max']:.0f}",
-    )
-
-    table.add_row(
-        "Mean / Median Tokens",
-        f"{metrics['Text']['input_tokens']['mean']:.0f} / {metrics['Text']['input_tokens']['median']:.0f}",
-        f"{metrics['Image']['input_tokens']['mean']:.0f} / {metrics['Image']['input_tokens']['median']:.0f}",
-        f"{metrics['Total']['input_tokens']['mean']:.0f} / {metrics['Total']['input_tokens']['median']:.0f}",
-    )
-
-    table.add_row(
-        "P5 / P95 Tokens",
-        f"{metrics['Text']['input_tokens']['p5']:.0f} / {metrics['Text']['input_tokens']['p95']:.0f}",
-        f"{metrics['Image']['input_tokens']['p5']:.0f} / {metrics['Image']['input_tokens']['p95']:.0f}",
-        f"{metrics['Total']['input_tokens']['p5']:.0f} / {metrics['Total']['input_tokens']['p95']:.0f}",
-    )
-
-    table.add_row("Sum Total Tokens", f"{metrics['Text']['sum_input_tokens']}", f"{metrics['Image']['sum_input_tokens']}", f"{metrics['Total']['sum_input_tokens']}")
-
-    table.add_row("")  # Empty row for spacing
-
-    if input_token_price is not None:
-        table.add_row(
-            "Input Cost",
-            f"{metrics['Text']['sum_input_tokens'] * input_token_price:.2f} USD",
-            f"{metrics['Image']['sum_input_tokens'] * input_token_price:.2f} USD",
-            f"{metrics['Total']['sum_input_tokens'] * input_token_price:.2f} USD",
-        )
-
-    if output_token_price is not None:
-        table.add_row(
-            "Output Cost",
-            f"{metrics['Text']['sum_output_tokens'] * output_token_price:.2f} USD",
-            f"{metrics['Image']['sum_output_tokens'] * output_token_price:.2f} USD",
-            f"{metrics['Total']['sum_output_tokens'] * output_token_price:.2f} USD",
-        )
-
-    if input_token_price is not None and output_token_price is not None:
-        table.add_row(
-            "Total Cost",
-            f"{metrics['Text']['sum_total_tokens'] * input_token_price:.2f} USD",
-            f"{metrics['Image']['sum_total_tokens'] * input_token_price:.2f} USD",
-            f"{metrics['Total']['sum_total_tokens'] * input_token_price:.2f} USD",
-        )
-
-    # Print the table
-    console.print(table)