granuscore 1.0.0__py3-none-any.whl

granuscore/__init__.py

@@ -0,0 +1,13 @@
+from .granularity_predictor import HitGranularityPredictor
+from .pipeline import GranuScore
+
+def granuscore(
+    texts,
+    **kwargs,
+):
+    return GranuScore()(
+        texts,
+        **kwargs,
+    )
+
+__all__ = ["GranuScore", "granuscore", "HitGranularityPredictor"]

granuscore/artifcats.py

@@ -0,0 +1,116 @@
+import hashlib
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+
+import requests
+from importlib.resources import files
+from platformdirs import user_cache_dir
+
+
+def _sha256(path: Path) -> str:
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+@dataclass(frozen=True)
+class ArtifactSpec:
+    name: str
+    subdir: str              # "faiss" or "models"
+    url: str | None = None   # if None: must be bundled or already cached
+    sha256: str | None = None
+
+
+class ArtifactManager:
+    """
+    Resolution order:
+      1) cache (and sha256 if provided)
+      2) bundled assets inside the package -> copied to cache
+      3) download from spec.url (public) -> cached
+    """
+
+    def __init__(self, app_name: str = "granuscore"):
+        self.cache_root = Path(user_cache_dir(app_name))
+        self.cache_root.mkdir(parents=True, exist_ok=True)
+
+    def ensure(self, spec: ArtifactSpec, force: bool = False) -> Path:
+        dst_dir = self.cache_root / spec.subdir
+        dst_dir.mkdir(parents=True, exist_ok=True)
+        dst = dst_dir / spec.name
+
+        # 1) cache hit
+        if dst.exists() and not force:
+            if spec.sha256 is None or _sha256(dst) == spec.sha256:
+                return dst
+            dst.unlink(missing_ok=True)
+
+        # 2) bundled asset
+        bundled = files("granuscore") / "../../assets" / spec.subdir / spec.name
+        if bundled.is_file() and not force:
+            self._copy_to_cache(bundled, dst)
+            self._verify_hash_if_needed(dst, spec)
+            return dst
+
+        # 3) public download
+        if not spec.url:
+            raise FileNotFoundError(
+                f"{spec.name} not found in cache and not bundled, and no public download URL provided."
+            )
+
+        self._download_public(spec.url, dst)
+        self._verify_hash_if_needed(dst, spec)
+        return dst
+
+    def _verify_hash_if_needed(self, path: Path, spec: ArtifactSpec) -> None:
+        if spec.sha256 is None:
+            return
+        got = _sha256(path)
+        if got != spec.sha256:
+            path.unlink(missing_ok=True)
+            raise ValueError(
+                f"Hash mismatch for {spec.name}. Expected {spec.sha256}, got {got}."
+            )
+
+    def _copy_to_cache(self, src_path: Path, dst: Path) -> None:
+        tmp = dst.with_suffix(dst.suffix + ".part")
+        with src_path.open("rb") as src, tmp.open("wb") as out:
+            shutil.copyfileobj(src, out)
+        tmp.replace(dst)
+
+    def _download_public(self, url: str, dst: Path) -> None:
+        tmp = dst.with_suffix(dst.suffix + ".part")
+
+        with requests.get(url, stream=True, timeout=60, allow_redirects=True) as r:
+            # Common mistake: GitHub "blob" links download HTML
+            if r.status_code >= 400 and "github.com/" in url and "/blob/" in url:
+                raise requests.HTTPError(
+                    "You are using a GitHub 'blob' URL (HTML). Use a release asset URL "
+                    "(.../releases/download/...) or a raw.githubusercontent.com URL."
+                )
+
+            r.raise_for_status()
+            total = int(r.headers.get("Content-Length", 0)) or None
+
+            try:
+                from tqdm import tqdm  # type: ignore
+                bar = tqdm(total=total, unit="B", unit_scale=True, desc=f"Downloading {dst.name}")
+            except Exception:
+                bar = None
+
+            try:
+                with tmp.open("wb") as f:
+                    for chunk in r.iter_content(chunk_size=1024 * 1024):
+                        if not chunk:
+                            continue
+                        f.write(chunk)
+                        if bar is not None:
+                            bar.update(len(chunk))
+            finally:
+                if bar is not None:
+                    bar.close()
+
+        os.replace(tmp, dst)

granuscore/bucket_output.py

@@ -0,0 +1,53 @@
+import numpy as np
+from scipy.stats import percentileofscore
+
+
+class PercentileOutput:
+    def __init__(self, path):
+        self.reference_scores = np.load(path)
+        self.percentiles: dict[float, float] | None = None
+
+    def _compute_percentiles(self):
+        self.percentiles = {
+            p: float(np.percentile(self.reference_scores, p))
+            for p in [33.33, 66.66]
+        }
+
+    def score_to_percentile(self, scores: float | list[float] | np.ndarray):
+        scores = np.asarray(scores)
+        return percentileofscore(self.reference_scores, scores)
+
+    def percentile_to_score(self, percentiles: float | list[float] | np.ndarray):
+        is_scalar = np.isscalar(percentiles)
+        arr = np.asarray([percentiles] if is_scalar else percentiles, dtype=float)
+
+        scores = np.percentile(self.reference_scores, arr)
+
+        if is_scalar:
+            return float(scores[0])
+        return scores
+
+    def bucket(self, scores: float | list[float] | np.ndarray) -> str:
+        """
+        Assign a granularity bucket based on score-space percentiles.
+
+        Args:
+            scores:
+                Granuscore percentile value.
+        Returns:
+            One of: "low", "medium", "high"
+        """
+        if self.percentiles is None:
+            self._compute_percentiles()
+
+        q1 = 33.33
+        q2 = 66.66
+
+        is_scalar = np.isscalar(scores)
+        arr = np.asarray([scores] if is_scalar else scores, dtype=float)
+
+        labels = np.where(arr < q1, "low", np.where(arr < q2, "medium", "high"))
+
+        if is_scalar:
+            return str(labels[0])
+        return labels.tolist()

granuscore/cache.py

@@ -0,0 +1,189 @@
+from typing import List, Sequence, Protocol
+from pathlib import Path
+import sqlite3
+import hashlib
+import os
+import threading
+import numpy as np
+
+__all__ = ["GranuscoreCache"]
+
+
+def _md5(s: str) -> str:
+    return hashlib.md5(s.encode()).hexdigest()
+
+
+class _Scorer(Protocol):
+    def score_wo_cache(self, answers: Sequence[str], encoding_batch_size: int | None = None) -> np.ndarray | list: ...
+
+
+class GranuscoreCache:
+    """Thin wrapper around the Granuscore with SQLite caching."""
+
+    def __init__(
+            self,
+            scorer: "_Scorer",
+            hierarchy_model: str,
+            faiss_index: str,
+            lgb_model: str,
+            cache_dir: str | os.PathLike = None,
+    ) -> None:
+        self.scorer = scorer
+        self.hierarchy_model = hierarchy_model
+        self.faiss_index = faiss_index
+        self.lgb_model = lgb_model
+
+        if cache_dir is None:
+            cache_dir = os.getenv("GRANUSCORE_CACHE_DIR", f"{Path(__file__).parent}/.cache") # default is .cache dir next to this file
+        cache_path = Path(cache_dir)
+        cache_path.mkdir(parents=True, exist_ok=True)
+        self.db_path = cache_path / "granuscores.db"
+
+        self._init_db()
+
+        self._w_lock = threading.Lock()
+
+    def get(self, text: str) -> float:
+        """Return the granuscore for *text*, fetching from cache or scoring."""
+        text_hash = _md5(text)
+
+        # 1. try cache ------------------------------------------------------
+        row = self._select(text_hash)
+        if row is not None:
+            return row
+
+        score = self.scorer.score_wo_cache([text])[0]
+        self._insert(text_hash, text, score)
+        return score
+
+    def batch(self, texts: Sequence[str], encoding_batch_size: int | None = None) -> List[float]:
+        texts = list(texts)
+        if not texts:
+            return []
+
+        # Precompute hashes and preserve order
+        hashes = [_md5(t) for t in texts]
+
+        # 1) Bulk fetch cache hits
+        hit_map = self._select_many(hashes)  # {text_hash: granuscore}
+
+        # 2) Identify misses (preserve order)
+        miss_texts = [t for t, h in zip(texts, hashes) if h not in hit_map]
+        miss_hashes = [h for h in hashes if h not in hit_map]
+
+        # 3) Score misses in one model call
+        if miss_texts:
+            miss_scores = self.scorer.score_wo_cache(miss_texts, encoding_batch_size)
+            miss_scores = np.asarray(miss_scores, dtype=float)
+
+            # 4) Bulk insert misses
+            self._insert_many(miss_hashes, miss_texts, miss_scores)
+
+            # update hit_map so we can fill outputs
+            hit_map.update({h: float(s) for h, s in zip(miss_hashes, miss_scores)})
+
+        # 5) Build output aligned to input order
+        return [float(hit_map[h]) for h in hashes]
+
+    # ------------------------------------------------------------------
+    # private db helpers
+    # ------------------------------------------------------------------
+    def _open_conn(self) -> sqlite3.Connection:
+        conn = sqlite3.connect(
+            self.db_path,
+            timeout=30,
+            check_same_thread=False,
+            isolation_level=None  # autocommit
+        )
+        conn.execute("PRAGMA journal_mode=WAL;")
+        conn.execute("PRAGMA synchronous=NORMAL;")
+        return conn
+
+    def _init_db(self) -> None:
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute(
+                """
+                CREATE TABLE IF NOT EXISTS granuscores (
+                    text_hash       TEXT,
+                    hierarchy_model TEXT,
+                    faiss_index TEXT,
+                    lgb_model  TEXT,
+                    granuscore REAL NOT NULL,
+                    text       TEXT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    PRIMARY KEY(text_hash, hierarchy_model, faiss_index, lgb_model)
+                )
+                """
+            )
+
+            conn.execute("CREATE INDEX IF NOT EXISTS idx_model ON granuscores(hierarchy_model, faiss_index, lgb_model);")
+
+    def _select(self, text_hash: str) -> float | None:
+        with sqlite3.connect(self.db_path) as conn:
+            row = conn.execute(
+                "SELECT granuscore FROM granuscores WHERE text_hash=? AND hierarchy_model=? AND faiss_index=? AND lgb_model=?",
+                (text_hash, self.hierarchy_model, self.faiss_index, self.lgb_model),
+            ).fetchone()
+            if row:
+                return row[0]
+            return None
+
+    def _select_many(self, text_hashes: Sequence[str], chunk_size: int = 900) -> dict[str, float]:
+        """
+        Fetch many hashes at once. Chunked because SQLite has a limit on bound variables
+        (commonly 999).
+        """
+        out: dict[str, float] = {}
+        if not text_hashes:
+            return out
+
+        with sqlite3.connect(self.db_path) as conn:
+            for i in range(0, len(text_hashes), chunk_size):
+                chunk = text_hashes[i:i + chunk_size]
+                placeholders = ",".join(["?"] * len(chunk))
+
+                rows = conn.execute(
+                    f"""
+                    SELECT text_hash, granuscore
+                    FROM granuscores
+                    WHERE hierarchy_model = ?
+                      AND faiss_index = ?
+                      AND lgb_model = ?
+                      AND text_hash IN ({placeholders})
+                    """,
+                    (self.hierarchy_model, self.faiss_index, self.lgb_model, *chunk),
+                ).fetchall()
+
+                out.update({h: float(s) for h, s in rows})
+
+        return out
+
+    def _insert(self, text_hash: str, text: str, granusccore: float) -> None:
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute(
+                "INSERT OR REPLACE INTO granuscores (text_hash, hierarchy_model, faiss_index, lgb_model, granuscore, text) VALUES (?,?,?,?,?,?)",
+                (text_hash, self.hierarchy_model, self.faiss_index, self.lgb_model, granusccore, text),
+            )
+
+    def _insert_many(self, hashes: Sequence[str], texts: Sequence[str], scores: Sequence[float]) -> None:
+        """
+        Bulk insert inside one transaction. Much faster than opening a connection per row.
+        """
+        if not hashes:
+            return
+
+        rows = [
+            (h, self.hierarchy_model, self.faiss_index, self.lgb_model, float(s), t)
+            for h, t, s in zip(hashes, texts, scores)
+        ]
+
+        with self._w_lock:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.executemany(
+                    """
+                    INSERT OR REPLACE INTO granuscores
+                    (text_hash, hierarchy_model, faiss_index, lgb_model, granuscore, text)
+                    VALUES (?, ?, ?, ?, ?, ?)
+                    """,
+                    rows,
+                )

granuscore/claim_splitter.py

@@ -0,0 +1,300 @@
+"""
+Claim splitting module for extracting atomic facts from text.
+
+This module provides various strategies for splitting text into atomic claims,
+from simple word-based splitting to sophisticated LLM-based extraction.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Sequence, Literal, Dict, Any
+
+import spacy
+from spacy.matcher import Matcher
+
+from granuscore.text_cleaner import TextCleaner
+from granuscore.utils import ensure_spacy_model
+
+
+SplitScope = Literal["document", "sentence"]
+
+
+class UnitSplitter(ABC):
+    """
+    Abstract base class for unit splitting strategies.
+
+    Defines the interface for extracting referential units from text,
+    supporting both single-text and batch processing.
+    """
+
+    def __call__(self, batch: List[str], split_scope: SplitScope = "document") -> list[list[dict[str, Any]]]:
+        """Allow callable interface for batch processing."""
+        return self.get_referential_units_batch(batch, split_scope)
+
+    @abstractmethod
+    def get_referential_units(self, text: str, split_scope: SplitScope = "document") -> list[dict[str, Any]]:
+        """
+        Extract referential units from a single text.
+
+        Args:
+            text: Input text to process
+            split_scope: Defines the structural scope over inside which referential units should be split.
+            "document" splits the whole text.
+            "sentence" applies splitting within each sentence.
+
+        Returns:
+            List[Dict[str, Any]]
+                One entry per group (document or sentence) with:
+                - text: str
+                - referential_units: List[str]
+            """
+        pass
+
+    @abstractmethod
+    def get_referential_units_batch(self, texts: Sequence[str], split_scope: SplitScope = "document") -> list[list[dict[str, Any]]]:
+        """
+        Extract referential units from multiple texts.
+
+        Args:
+            texts: List of input texts to process
+            split_scope: Defines the structural scope over inside which referential units should be split.
+            "document" splits the whole text.
+            "sentence" applies splitting within each sentence.
+
+        Returns:
+            List[List[Dict[str, Any]]]:
+                A nested list structure where:
+                - the outer list indexes input texts,
+                - the second level indexes document- or sentence-level groups
+                  (depending on ``split_scope``),
+                - each group is represented as a dictionary with the following keys:
+                    - ``"text"`` (str):
+                        The text span corresponding to the group.
+                    - ``"referential_units"`` (List[str]):
+                        The extracted referential units for the group.
+        """
+        pass
+
+
+class SpacyNounPhraseSplitter(UnitSplitter):
+    """
+    Spacy-based noun phrase extractor.
+
+    Extracts noun phrases and significant tokens using spaCy's NLP pipeline.
+    Includes cleaning options for academic text (LaTeX, figures, etc.).
+    """
+
+    def __init__(
+            self,
+            model: str = "en_core_web_sm",
+            skip_brackets: bool = False,
+            skip_latex: bool = False,
+            skip_arxiv_meta: bool = False,
+            skip_figures: bool = False,
+            skip_pdf_noise: bool = False,
+    ):
+        """
+        Initialize with spaCy model and cleaning options.
+
+        Args:
+            model: spaCy model name to load
+            skip_brackets: Remove bracketed content
+            skip_latex: Remove LaTeX blocks
+            skip_arxiv_meta: Remove arXiv metadata
+            skip_figures: Remove figure captions
+            skip_pdf_noise: Remove PDF noise
+        """
+        ensure_spacy_model("en_core_web_sm")
+        self.nlp = spacy.load(model)
+        self._add_et_al_merging()
+        self.clean_kwargs = {
+            'skip_brackets': skip_brackets,
+            'skip_latex': skip_latex,
+            'skip_arxiv_meta': skip_arxiv_meta,
+            'skip_figures': skip_figures,
+            'skip_pdf_noise': skip_pdf_noise,
+        }
+
+    def _add_et_al_merging(self):
+        """Add custom component to merge 'et al' into single tokens."""
+        matcher = Matcher(self.nlp.vocab)
+        matcher.add(
+            "ET_AL",
+            [[
+                {"IS_TITLE": True},
+                {"LOWER": "et"},
+                {"LOWER": "al"},
+            ]]
+        )
+
+        @spacy.Language.component("merge_et_al")
+        def merge_et_al(doc):
+            matches = matcher(doc)
+            with doc.retokenize() as retokenizer:
+                for _, start, end in matches:
+                    retokenizer.merge(doc[start:end])
+            return doc
+
+        self.nlp.add_pipe("merge_et_al", before="parser")
+
+    @staticmethod
+    def _is_valid_noun_chunk(chunk) -> bool:
+        """Check if a noun chunk is valid for extraction."""
+        # Must be nominal
+        if chunk.root.pos_ not in {"NOUN", "PROPN", "PRON"}:
+            return False
+
+        # Reject punctuation-only
+        if all(tok.is_punct for tok in chunk):
+            return False
+
+        # Reject single hyphen or fragments
+        if chunk.text.strip() in {"-", "–", "—"}:
+            return False
+
+        return True
+
+
+    def _extract_from_doclike(self, text) -> list[str]:
+        """
+        Extract noun phrases and significant tokens from a spaCy Doc or Span.
+        Returns a flat list of units for that segment.
+        """
+        doclike = self.nlp(text)
+
+        spans = {
+            (chunk.start, chunk.end)
+            for chunk in doclike.noun_chunks
+            if self._is_valid_noun_chunk(chunk)
+        }
+
+        results: List[str] = []
+        i = 0
+        while i < len(doclike):
+            for start, end in spans:
+                if i == start:
+                    results.append(doclike[start:end].text)
+                    i = end
+                    break
+            else:
+                tok = doclike[i]
+                if (not tok.is_stop or tok.like_num) and any(c.isalnum() for c in tok.text):
+                    results.append(tok.text)
+                i += 1
+
+        return results
+
+    def get_referential_units(self, text: str, split_scope: SplitScope = "document") -> list[dict[str, Any]]:
+        """
+        Extract noun phrases and significant tokens, optionally grouped by sentence.
+
+        Returns
+        -------
+        List[Dict[str, Any]]
+            One entry per group (document or sentence) with:
+            - referential_units: List[str]
+            - cleaned_text: str
+            - original_text: str
+        """
+        text = TextCleaner.clean(text, **self.clean_kwargs)
+
+        if split_scope == "document":
+            return [{
+                "referential_units": self._extract_from_doclike(text),
+                "text": text,
+            }]
+
+        if split_scope == "sentence":
+            doc = self.nlp(text)
+            groups: List[Dict[str, Any]] = []
+            for sent in doc.sents:
+                sent_text = sent.text.strip()
+                if not sent_text:
+                    continue
+                groups.append({
+                    "referential_units": self._extract_from_doclike(sent_text),
+                    "text": sent_text,
+                })
+            if not groups:
+                groups.append({
+                    "referential_units": [],
+                    "text": text,
+                })
+            return groups
+
+        raise ValueError('split_scope must be either "document" or "sentence"')
+
+    def get_referential_units_batch(self, texts: Sequence[str], split_scope: SplitScope = 'document') -> list[list[dict[str, Any]]]:
+        """Process multiple texts."""
+        return [self.get_referential_units(text, split_scope) for text in texts]
+
+
+class SpacyObjectSplitter(SpacyNounPhraseSplitter):
+    DEP_KEEP = {
+        "nsubj", "nsubjpass",  # subjects
+        "dobj", "obj", "iobj",  # objects
+        "obl",  # oblique nominal (where/when/how)
+    }
+
+    def _extract_from_doclike(self, text: str) -> list[str]:
+        """
+        Extract units in document order:
+          1) noun chunks (preferred)
+          2) otherwise, subject/object-like tokens (dep-based)
+        """
+        doc = self.nlp(text)
+
+        # spans for noun chunks
+        spans = {
+            (chunk.start, chunk.end)
+            for chunk in doc.noun_chunks
+            if self._is_valid_noun_chunk(chunk)
+        }
+        start_to_end = {start: end for start, end in spans}
+
+        results = []
+        seen = set()
+
+        i = 0
+        while i < len(doc):
+            end = start_to_end.get(i)
+            if end is not None:
+                text_unit = doc[i:end].text.strip()
+                if text_unit and text_unit not in seen:
+                    results.append(text_unit)
+                    seen.add(text_unit)
+                i = end
+                continue
+
+            tok = doc[i]
+            if tok.dep_ in self.DEP_KEEP and any(c.isalnum() for c in tok.text):
+                if tok.text not in seen:
+                    results.append(tok.text)
+                    seen.add(tok.text)
+
+            i += 1
+
+        return results
+
+if __name__ == '__main__':
+    # Example: Using the noun phrase splitter
+    splitter = SpacyNounPhraseSplitter(
+        skip_brackets=True,
+        skip_latex=True,
+        skip_figures=True,
+        skip_arxiv_meta=True,
+        skip_pdf_noise=True
+    )
+
+    sample_texts = [
+        " O'Reilly",
+        "The population is 12,000 (approx.)."
+        "three",
+        "Two."
+    ]
+
+    results = splitter.get_referential_units_batch(sample_texts)
+
+    for i, claims in enumerate(results, 1):
+        print(f"\nText {i} claims:")
+        for claim in claims:
+            print(f"  - {claim}")