openai-gabriel 1.0.1__py3-none-any.whl

gabriel/utils/mapmaker.py

@@ -0,0 +1,304 @@
+"""
+mapmaker.py
+~~~~~~~~~~~~
+
+This module implements a self‑contained helper class for generating
+choropleth maps from ranked GABRIEL outputs.  The ``MapMaker`` can
+produce county‑, state‑ or country‑level maps depending on the
+``map_type`` parameter, and it optionally normalises values as
+z‑scores.  All mapping logic resides in this file, replacing the need
+for the separate ``create_county_choropleth`` function.
+"""
+
+from __future__ import annotations
+
+import os
+import json
+import requests
+from typing import Iterable, Optional
+
+import numpy as np
+import pandas as pd
+import plotly.express as px
+
+
+class MapMaker:
+    """Utility for generating geographic choropleth maps from a data frame.
+
+    The input data frame should contain at least one column of numeric
+    scores and one or more columns identifying the geographic unit
+    (county FIPS codes, two‑letter US state abbreviations or ISO‑3
+    country codes).  Individual maps are rendered using Plotly and
+    written to ``save_dir`` with names derived from the value column.
+
+    Parameters
+    ----------
+    df:
+        DataFrame containing the data to plot.  Each row should
+        correspond to a geographic region.
+    fips_col:
+        Name of the column containing five‑digit county FIPS codes.
+    state_col:
+        Name of the column containing two‑letter US state abbreviations.
+    country_col:
+        Name of the column containing ISO‑3 country codes.
+    save_dir:
+        Directory to which map files will be written.  If ``None``,
+        a ``maps`` subdirectory in the current working directory is used.
+    z_score:
+        Whether to convert values to z‑scores before plotting.
+    color_scale:
+        Name of the Plotly colour scale to apply.  Defaults to
+        ``"RdBu"`` (diverging) when z‑scores are enabled and
+        ``"Viridis"`` otherwise.
+    map_type:
+        Determines the map produced: ``"county"``, ``"state"``
+        or ``"country"`` (with ``"global"`` as an alias).
+    """
+
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        *,
+        fips_col: Optional[str] = None,
+        state_col: Optional[str] = None,
+        country_col: Optional[str] = None,
+        save_dir: Optional[str] = None,
+        z_score: bool = True,
+        color_scale: str = "RdBu",
+        map_type: str = "county",
+    ) -> None:
+        self.df = df.copy()
+        self.fips_col = fips_col
+        self.state_col = state_col
+        self.country_col = country_col
+
+        # normalise map_type and validate
+        map_type = map_type.lower()
+        if map_type not in {"county", "state", "country", "global"}:
+            raise ValueError(
+                "map_type must be one of 'county', 'state', 'country' or 'global'"
+            )
+        self.map_type = "country" if map_type == "global" else map_type
+
+        # choose save directory
+        if save_dir is None:
+            save_dir = os.path.join(os.getcwd(), "maps")
+        save_dir = os.path.expandvars(os.path.expanduser(save_dir))
+        os.makedirs(save_dir, exist_ok=True)
+        self.save_dir = save_dir
+
+        self.z_score = z_score
+        self.color_scale = color_scale
+
+    def _compute_zscore(self, values: np.ndarray) -> np.ndarray:
+        """Compute z‑scores with safe handling of NaNs and constant arrays."""
+        vals = values.astype(float)
+        if len(vals) > 1 and np.nanstd(vals) > 0:
+            return (vals - np.nanmean(vals)) / np.nanstd(vals)
+        return np.zeros_like(vals)
+
+    def _create_state_choropleth(
+        self,
+        df: pd.DataFrame,
+        state_col: str,
+        value_col: str,
+        title: str,
+        save_path: str,
+    ) -> None:
+        """Create and save a state‑level choropleth."""
+        plot_col = value_col
+        colour_scale = self.color_scale
+        df_local = df.copy()
+        if self.z_score:
+            zs = self._compute_zscore(df_local[value_col].values)
+            plot_col = f"_zscore_{value_col}"
+            df_local[plot_col] = zs
+            colour_scale = "RdBu" if self.color_scale == "RdBu" else "PuOr"
+        fig = px.choropleth(
+            df_local,
+            locations=state_col,
+            locationmode="USA-states",
+            color=plot_col,
+            color_continuous_scale=colour_scale,
+            scope="usa",
+            hover_data={state_col: True, value_col: True},
+        )
+        if self.z_score:
+            fig.update_coloraxes(cmid=0)
+        ext = os.path.splitext(save_path)[1].lower()
+        if ext in {".png", ".jpg", ".jpeg"}:
+            fig.write_image(save_path, scale=3)
+        else:
+            fig.write_html(save_path)
+
+    def _create_country_choropleth(
+        self,
+        df: pd.DataFrame,
+        country_col: str,
+        value_col: str,
+        title: str,
+        save_path: str,
+    ) -> None:
+        """Create and save a global choropleth using ISO‑3 codes."""
+        plot_col = value_col
+        colour_scale = self.color_scale
+        df_local = df.copy()
+        if self.z_score:
+            zs = self._compute_zscore(df_local[value_col].values)
+            plot_col = f"_zscore_{value_col}"
+            df_local[plot_col] = zs
+            colour_scale = "RdBu" if self.color_scale == "RdBu" else "PuOr"
+        fig = px.choropleth(
+            df_local,
+            locations=country_col,
+            locationmode="ISO-3",
+            color=plot_col,
+            color_continuous_scale=colour_scale,
+            scope="world",
+            hover_data={country_col: True, value_col: True},
+        )
+        if self.z_score:
+            fig.update_coloraxes(cmid=0)
+        ext = os.path.splitext(save_path)[1].lower()
+        if ext in {".png", ".jpg", ".jpeg"}:
+            fig.write_image(save_path, scale=3)
+        else:
+            fig.write_html(save_path)
+
+    def _create_county_choropleth(
+        self,
+        df: pd.DataFrame,
+        fips_col: str,
+        value_col: str,
+        title: str,
+        save_path: str,
+    ) -> None:
+        """Create and save a county‑level choropleth with FIPS codes.
+
+        This method inlines the logic of the old ``create_county_choropleth``
+        function to avoid external dependencies.  It downloads a GeoJSON of
+        U.S. counties on first use and caches it in ``~/.cache/county_geo.json``.
+        """
+        # pad FIPS codes to five digits
+        df_local = df.copy()
+        df_local[fips_col] = df_local[fips_col].astype(str).str.zfill(5)
+
+        # find a county name column for hover text
+        county_col = None
+        for cand in ["county", "County", "region", "Region"]:
+            if cand in df_local.columns:
+                county_col = cand
+                break
+        if county_col is None:
+            county_col = "_county_name"
+            df_local[county_col] = ""
+
+        # load or download county GeoJSON
+        geojson_url = (
+            "https://raw.githubusercontent.com/plotly/datasets/master/geojson-counties-fips.json"
+        )
+        cache_path = os.path.join(os.path.expanduser("~"), ".cache", "county_geo.json")
+        if not os.path.exists(cache_path):
+            os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+            resp = requests.get(geojson_url, timeout=30)
+            resp.raise_for_status()
+            with open(cache_path, "w", encoding="utf-8") as f:
+                f.write(resp.text)
+        with open(cache_path, encoding="utf-8") as f:
+            counties = json.load(f)
+
+        # prepare value column (with optional z‑score)
+        plot_col = value_col
+        colour_scale = self.color_scale
+        if self.z_score:
+            zs = self._compute_zscore(df_local[value_col].values)
+            plot_col = f"_zscore_{value_col}"
+            df_local[plot_col] = zs
+            colour_scale = "RdBu" if self.color_scale == "RdBu" else "PuOr"
+
+        hover_data = {county_col: True, fips_col: True, value_col: True}
+        fig = px.choropleth(
+            df_local,
+            geojson=counties,
+            locations=fips_col,
+            color=plot_col,
+            color_continuous_scale=colour_scale,
+            scope="usa",
+            hover_data=hover_data,
+        )
+        if self.z_score:
+            fig.update_coloraxes(cmid=0)
+        ext = os.path.splitext(save_path)[1].lower()
+        if ext in {".png", ".jpg", ".jpeg"}:
+            fig.write_image(save_path, scale=3)
+        else:
+            fig.write_html(save_path)
+
+    def make_maps(self, value_cols: Iterable[str]) -> None:
+        """Generate and save maps for each specified numeric column.
+
+        The map type is determined by ``self.map_type``: ``"county"``
+        uses FIPS codes; ``"state"`` uses two‑letter abbreviations; and
+        ``"country"`` uses ISO‑3 codes.
+        """
+        for value_col in value_cols:
+            if self.map_type == "county":
+                if not self.fips_col:
+                    raise ValueError("fips_col must be provided for county maps")
+                fname = f"county_map_{value_col}.html"
+                save_path = os.path.join(self.save_dir, fname)
+                self._create_county_choropleth(
+                    self.df,
+                    self.fips_col,
+                    value_col,
+                    title=f"County Map for {value_col}",
+                    save_path=save_path,
+                )
+            elif self.map_type == "state":
+                if not self.state_col:
+                    raise ValueError("state_col must be provided for state maps")
+                fname = f"state_map_{value_col}.html"
+                save_path = os.path.join(self.save_dir, fname)
+                self._create_state_choropleth(
+                    self.df,
+                    self.state_col,
+                    value_col,
+                    title=f"State Map for {value_col}",
+                    save_path=save_path,
+                )
+            elif self.map_type == "country":
+                if not self.country_col:
+                    raise ValueError("country_col must be provided for country maps")
+                fname = f"country_map_{value_col}.html"
+                save_path = os.path.join(self.save_dir, fname)
+                self._create_country_choropleth(
+                    self.df,
+                    self.country_col,
+                    value_col,
+                    title=f"Country Map for {value_col}",
+                    save_path=save_path,
+                )
+            else:
+                # should not happen due to validation in __init__
+                raise ValueError(f"Unsupported map type: {self.map_type}")
+
+
+def create_county_choropleth(
+    df: pd.DataFrame,
+    *,
+    fips_col: str,
+    value_col: str,
+    title: str,
+    save_path: str,
+    z_score: bool = True,
+) -> None:
+    """Backward compatible helper to generate a county-level choropleth.
+
+    This thin wrapper instantiates :class:`MapMaker` and delegates to its
+    internal implementation.  It mirrors the signature of the legacy
+    ``create_county_choropleth`` function used elsewhere in the codebase.
+    """
+
+    mm = MapMaker(df, fips_col=fips_col, z_score=z_score, save_dir=None, map_type="county")
+    mm._create_county_choropleth(df, fips_col, value_col, title, save_path)

gabriel/utils/media_utils.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import os
+from typing import Any, Dict, List
+
+from .image_utils import encode_image
+from .audio_utils import encode_audio
+from .pdf_utils import encode_pdf
+
+
+def load_image_inputs(val: Any) -> List[str]:
+    """Return a list of base64-encoded images from a DataFrame cell.
+
+    ``val`` may be a single file path, a list of file paths, or a list of
+    pre-encoded base64 strings. Non-existing paths are ignored.
+    """
+    if not val:
+        return []
+    imgs = val if isinstance(val, list) else [val]
+    encoded: List[str] = []
+    for img in imgs:
+        if isinstance(img, str) and os.path.exists(img):
+            enc = encode_image(img)
+            if enc:
+                encoded.append(enc)
+        elif isinstance(img, str):
+            encoded.append(img)
+    return encoded
+
+
+def load_audio_inputs(val: Any) -> List[Dict[str, str]]:
+    """Return a list of audio dicts from a DataFrame cell.
+
+    ``val`` may be a single file path, a list of file paths, or a list of
+    already-encoded dicts. Non-existing paths are ignored.
+    """
+    if not val:
+        return []
+    auds = val if isinstance(val, list) else [val]
+    encoded: List[Dict[str, str]] = []
+    for aud in auds:
+        if isinstance(aud, str) and os.path.exists(aud):
+            enc = encode_audio(aud)
+            if enc:
+                encoded.append(enc)
+        elif isinstance(aud, dict):
+            encoded.append(aud)
+    return encoded
+
+
+def load_pdf_inputs(val: Any) -> List[Dict[str, str]]:
+    """Return a list of PDF dicts from a DataFrame cell."""
+    if not val:
+        return []
+    pdfs = val if isinstance(val, list) else [val]
+    encoded: List[Dict[str, str]] = []
+    for pdf in pdfs:
+        if isinstance(pdf, str) and os.path.exists(pdf):
+            enc = encode_pdf(pdf)
+            if enc:
+                encoded.append(enc)
+        elif isinstance(pdf, dict):
+            encoded.append(pdf)
+        elif isinstance(pdf, str):
+            lowered = pdf.lower()
+            if lowered.startswith("data:application/pdf"):
+                encoded.append({"filename": "document.pdf", "file_data": pdf})
+            elif lowered.startswith("http://") or lowered.startswith("https://"):
+                if lowered.endswith(".pdf"):
+                    encoded.append({"file_url": pdf})
+            else:
+                encoded.append(
+                    {
+                        "filename": "document.pdf",
+                        "file_data": f"data:application/pdf;base64,{pdf}",
+                    }
+                )
+    return encoded

gabriel/utils/modality_utils.py

@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+import os
+import re
+from typing import Any, Iterable, List, Tuple
+
+from .file_utils import AUDIO_EXTENSIONS, IMAGE_EXTENSIONS, PDF_EXTENSIONS
+
+IMAGE_SUFFIXES = {ext.lower() for ext in IMAGE_EXTENSIONS}
+AUDIO_SUFFIXES = {ext.lower() for ext in AUDIO_EXTENSIONS}
+PDF_SUFFIXES = {ext.lower() for ext in PDF_EXTENSIONS}
+
+
+def warn_if_modality_mismatch(
+    values: Iterable[Any],
+    modality: str,
+    *,
+    column_name: str,
+    sample_size: int = 100,
+) -> None:
+    """Inspect a sample of values and warn if modality likely mismatches."""
+
+    sample = list(values)[:sample_size]
+    if not sample:
+        return
+
+    pdf_hits = 0
+    image_hits = 0
+    audio_hits = 0
+    text_word_counts: List[int] = []
+
+    for value in sample:
+        candidate = _coerce_first_value(value)
+        if candidate is None:
+            continue
+        kind = _detect_media_kind(candidate)
+        if kind == "pdf":
+            pdf_hits += 1
+            continue
+        if kind == "image":
+            image_hits += 1
+            continue
+        if kind == "audio":
+            audio_hits += 1
+            continue
+        if isinstance(candidate, str):
+            words = re.findall(r"\\b\\w+\\b", candidate)
+            text_word_counts.append(len(words))
+
+    total = len(sample)
+    if total == 0:
+        return
+
+    if pdf_hits == total and modality != "pdf":
+        print(
+            f"[gabriel] Detected PDFs in column '{column_name}'. "
+            "Set modality='pdf' to attach files directly, or set modality='text' "
+            "(or 'entity'/'web') to extract plain text instead."
+        )
+    if image_hits == total and modality != "image":
+        print(
+            f"[gabriel] Detected image-like inputs in column '{column_name}'. "
+            "Set modality='image' to attach image files correctly."
+        )
+    if audio_hits == total and modality != "audio":
+        print(
+            f"[gabriel] Detected audio-like inputs in column '{column_name}'. "
+            "Set modality='audio' to attach audio files correctly."
+        )
+    if modality == "pdf" and pdf_hits == 0:
+        print(
+            f"[gabriel] Column '{column_name}' doesn't look like PDF inputs. "
+            "If this is text, consider modality='text' (or 'entity'/'web')."
+        )
+    if modality == "image" and image_hits == 0:
+        print(
+            f"[gabriel] Column '{column_name}' doesn't look like image inputs. "
+            "If this is text, consider modality='text' (or 'entity'/'web')."
+        )
+    if modality == "audio" and audio_hits == 0:
+        print(
+            f"[gabriel] Column '{column_name}' doesn't look like audio inputs. "
+            "If this is text, consider modality='text' (or 'entity'/'web')."
+        )
+
+    if modality == "text" and text_word_counts:
+        avg_words = sum(text_word_counts) / max(1, len(text_word_counts))
+        if avg_words < 10:
+            print(
+                f"[gabriel] Average word count in column '{column_name}' is {avg_words:.1f} "
+                "words. Confirm you intended modality='text' (rather than 'entity' or 'web')."
+            )
+    if modality in {"entity", "web"} and text_word_counts:
+        avg_words = sum(text_word_counts) / max(1, len(text_word_counts))
+        if avg_words > 30:
+            print(
+                f"[gabriel] Average word count in column '{column_name}' is {avg_words:.1f} "
+                "words. Consider modality='text' for long-form passages."
+            )
+
+
+def _coerce_first_value(value: Any) -> Any:
+    if value is None:
+        return None
+    if isinstance(value, (list, tuple)):
+        return value[0] if value else None
+    return value
+
+
+def _detect_media_kind(value: Any) -> str:
+    if isinstance(value, dict):
+        file_data = value.get("file_data")
+        file_url = value.get("file_url")
+        if isinstance(file_data, str) and file_data.startswith("data:application/pdf"):
+            return "pdf"
+        if isinstance(file_url, str) and file_url.lower().endswith(".pdf"):
+            return "pdf"
+        if value.get("format"):
+            return "audio"
+    if not isinstance(value, str):
+        return ""
+    lowered = value.lower().strip()
+    if lowered.startswith("data:application/pdf"):
+        return "pdf"
+    if lowered.startswith("data:image/"):
+        return "image"
+    if lowered.startswith("data:audio/"):
+        return "audio"
+    if _looks_like_url(lowered):
+        if lowered.endswith(tuple(PDF_SUFFIXES)):
+            return "pdf"
+        if lowered.endswith(tuple(IMAGE_SUFFIXES)):
+            return "image"
+        if lowered.endswith(tuple(AUDIO_SUFFIXES)):
+            return "audio"
+    if os.path.exists(value):
+        ext = os.path.splitext(value)[1].lower()
+        if ext in PDF_SUFFIXES:
+            return "pdf"
+        if ext in IMAGE_SUFFIXES:
+            return "image"
+        if ext in AUDIO_SUFFIXES:
+            return "audio"
+    return ""
+
+
+def _looks_like_url(value: str) -> bool:
+    return value.startswith("http://") or value.startswith("https://")