eval_document.py

"""
Evaluate an auto-generated long-form document.

Metrics
-------
1. Average cross-section cosine similarity  (embeddings)
2. Top-k most-similar section pairs         (with titles)
3. Self-BLEU (4-gram)                       – lexical redundancy
4. Shannon entropy per section              – lexical diversity

Usage
-----
    python eval_document.py doc.json \
        --model text-embedding-ada-002 \
        --top_k 5 \
        --top_ngrams 20 \
        --lang de
"""
from __future__ import annotations

import argparse
import json
import math
import re
import sys
from collections import Counter
from pathlib import Path
from typing import List, Tuple

import numpy as np
from litellm import embedding
from nltk import word_tokenize
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
from nltk.corpus import stopwords
from rich import box
from rich.console import Console
from rich.progress import Progress
from rich.table import Table
from sklearn.feature_extraction.text import CountVectorizer

console = Console()


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Evaluate LLM document JSON with optional LaTeX cleaning and language-specific n-gram filtering"
    )
    p.add_argument("json_path", type=Path, help="Path to JSON document")
    p.add_argument("--model", default="text-embedding-ada-002",
                   help="Embedding model name supported by LiteLLM")
    p.add_argument("--top_k", type=int, default=5,
                   help="Number of most-similar section pairs to display")
    p.add_argument("--top_ngrams", type=int, default=20,
                   help="Number of top n-grams to display")
    p.add_argument("--lang", choices=["en", "de"], default="en",
                   help="Language for stopword filtering (en or de)")
    p.add_argument("--clean_latex", action="store_true",
                   help="Remove LaTeX commands and environments before analysis")
    return p.parse_args()


def load_sections(json_path: Path) -> Tuple[List[str], List[str]]:
    with json_path.open("r", encoding="utf-8") as f:
        data = json.load(f)

    titles, texts = [], []
    for idx, chap in enumerate(data.get("chapters", []), 1):
        titles.append(chap.get("title", f"Section {idx}"))
        texts.append(chap.get("content", ""))

    if not texts:
        console.print("[bold red]No sections found in JSON.[/bold red]")
        sys.exit(1)
    return titles, texts


def clean_text(text: str) -> str:
    """
    Remove common LaTeX commands, environments, and braces from text.
    """
    # Strip LaTeX commands \command{...} or \command[...]{...}
    text = re.sub(r"\\[a-zA-Z]+(?:\[[^\]]*\])*(?:\{[^}]*\})?", "", text)
    # Remove curly braces
    text = text.replace("{", "").replace("}", "")
    # Collapse multiple spaces and newlines
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def preprocess_sections(texts: List[str], clean_latex: bool) -> List[str]:
    processed = []
    for t in texts:
        t2 = clean_text(t) if clean_latex else t
        processed.append(t2)
    return processed


def embed_sections(sections: List[str], model: str) -> np.ndarray:
    try:
        resp = embedding(model=model, input=sections)
        vecs = [d["embedding"] for d in resp["data"]]
    except Exception as e:
        console.print(f"[yellow]Batch embedding failed ({e}). Falling back to per-section calls.[/yellow]")
        vecs = []
        with Progress() as progress:
            task = progress.add_task("Embedding", total=len(sections))
            for sec in sections:
                resp = embedding(model=model, input=[sec])
                vecs.append(resp["data"][0]["embedding"])
                progress.update(task, advance=1)
    return np.asarray(vecs, dtype=np.float32)


def cosine_similarity_matrix(v: np.ndarray) -> np.ndarray:
    norm = np.linalg.norm(v, axis=1, keepdims=True)
    v_norm = v / np.clip(norm, 1e-15, None)
    return v_norm @ v_norm.T


def average_off_diagonal(mat: np.ndarray) -> float:
    n = mat.shape[0]
    return (mat.sum() - n) / (n * (n - 1))


def top_k_pairs(mat: np.ndarray, k: int) -> List[Tuple[int, int, float]]:
    n = mat.shape[0]
    tri = np.triu_indices(n, k=1)
    sims = mat[tri]
    idxs = np.argsort(sims)[-k:][::-1]
    return [(tri[0][i], tri[1][i], sims[i]) for i in idxs]


def tokenize(text: str) -> List[str]:
    return word_tokenize(text.lower())


def self_bleu(texts: List[str], n_gram: int = 4) -> float:
    weights = tuple(1 / n_gram for _ in range(n_gram))
    smooth = SmoothingFunction().method1
    toks = [tokenize(t) for t in texts]
    scores = []
    for i, hyp in enumerate(toks):
        refs = toks[:i] + toks[i+1:]
        if not any(refs):
            scores.append(0.0)
            continue
        scores.append(sentence_bleu(refs, hyp, weights=weights, smoothing_function=smooth))
    return float(np.mean(scores))


def ngram_counts(tokens: List[str], n: int = 1) -> Counter[str]:
    return Counter([" ".join(tokens[i:i+n]) for i in range(len(tokens) - n + 1)])


def shannon_entropy(tokens: List[str]) -> float:
    counts = Counter(tokens)
    total = len(tokens) or 1
    return -sum((c / total) * math.log2(c / total) for c in counts.values())


def unique_token_ratio(tokens: List[str]) -> float:
    total = len(tokens) or 1
    return len(set(tokens)) / total


def section_diversity(texts: List[str]) -> float:
    ent_norms = []
    for t in texts:
        toks = tokenize(t)
        V = len(set(toks)) or 1
        ent_norms.append(shannon_entropy(toks) / math.log2(V))
    return float(np.mean(ent_norms))


def top_ngrams(sections: List[str], n: int, lang: str) -> List[Tuple[str, int]]:
    if lang == 'de':
        try:
            german = stopwords.words('german')
        except LookupError:
            import nltk; nltk.download('stopwords')
            german = stopwords.words('german')
        stop_list = german
    else:
        stop_list = 'english'

    vectorizer = CountVectorizer(ngram_range=(1, 3), stop_words=stop_list)
    X = vectorizer.fit_transform(sections)
    freqs = np.asarray(X.sum(axis=0)).flatten()
    idx = np.argsort(freqs)[::-1][:n]
    vocab = vectorizer.get_feature_names_out()
    return [(vocab[i], int(freqs[i])) for i in idx]



def main() -> None:
    args = parse_args()
    titles, texts = load_sections(args.json_path)

    # Optional LaTeX cleaning
    texts = preprocess_sections(texts, args.clean_latex)

    console.rule("[bold cyan]Embedding Sections")
    vecs = embed_sections(texts, args.model)

    console.rule("[bold cyan]Computing Similarities")
    sim_mat = cosine_similarity_matrix(vecs)
    avg_sim = average_off_diagonal(sim_mat)
    pairs = top_k_pairs(sim_mat, args.top_k)

    console.rule("[bold cyan]Lexical Redundancy Metrics")
    s_bleu = self_bleu(texts)

    ent_norm = section_diversity(texts)

    ngram_stats = top_ngrams(texts, args.top_ngrams, args.lang)

    # Calculate document length metrics
    num_sections = len(texts)
    total_tokens = sum(len(tokenize(text)) for text in texts)
    avg_section_length = total_tokens / num_sections if num_sections > 0 else 0

    summary = Table(title="Document-level Quality Metrics", box=box.ROUNDED)
    summary.add_column("Metric", style="bold cyan")
    summary.add_column("Value", style="green")
    summary.add_column("Per Section", style="yellow")
    summary.add_column("Explanation", style="dim")

    summary.add_row("Sections", str(num_sections), "—", "Total number of sections")
    summary.add_row("Total tokens", str(total_tokens), f"{avg_section_length:.1f}", "Total/average tokens per section")
    summary.add_row("Embedding model", args.model, "—", "Model used for semantic similarity")
    summary.add_row("avg. cross-section similarity", f"{avg_sim:.4f}", f"{avg_sim/num_sections:.6f}", "high → repetitive content")
    summary.add_row("Self-BLEU (4-gram)", f"{s_bleu:.4f}", f"{s_bleu/num_sections:.6f}", "High → overlapping phrasing")
    summary.add_row("Avg. entropy (norm.)", f"{ent_norm:.4f}", f"{ent_norm/num_sections:.6f}", "High → lexical diversity")
    console.print(summary)

    pair_tbl = Table(title=f"Top {args.top_k} Most Similar Section Pairs", box=box.ROUNDED)
    pair_tbl.add_column("Rank", justify="right")
    pair_tbl.add_column("Section i")
    pair_tbl.add_column("Section j")
    pair_tbl.add_column("Cosine sim.", justify="right")
    for r, (i, j, score) in enumerate(pairs, 1):
        pair_tbl.add_row(str(r), titles[i], titles[j], f"{score:.4f}")
    console.print(pair_tbl)

    ngram_tbl = Table(title=f"Top {args.top_ngrams} n-grams (1–3 tokens)", box=box.ROUNDED)
    ngram_tbl.add_column("Rank", justify="right")
    ngram_tbl.add_column("n-gram")
    ngram_tbl.add_column("Freq.", justify="right")
    for r, (gram, freq) in enumerate(ngram_stats, 1):
        ngram_tbl.add_row(str(r), gram, str(freq))
    console.print(ngram_tbl)


if __name__ == "__main__":
    main()