paper_rewriter_OpenAPIcompatible.py

#!/usr/bin/env python3
"""
Paper Rewriter Plus V2 — Guided academic rewriter (OpenAI + OpenAI‑compatible incl. Ollama / gpt‑oss:120b)
---------------------------------------------------------------------------------------------------------

What this does
- Reads .docx, .pdf, .md/.txt, or .html and extracts clean text (heading‑aware)
- Splits text into token‑aware chunks with overlap
- Runs one task over all chunks (rewrite, guided‑rewrite, summary, extract, review, figures)
- Stitches outputs and writes both per‑chunk and full stitched files
- Works with **OpenAI** and **OpenAI‑compatible** servers (e.g., Ollama `/v1`, LM Studio, vLLM)
- Supports **streaming** and **long‑form continuation** for extensive outputs

Install
  pip install -U openai>=1.40.0 tiktoken>=0.7.0
  pip install -U python-docx>=1.1.2 pdfminer.six>=20231228 html2text>=2024.2.26 pypdf>=4.2.0

Providers
- OpenAI (cloud): set `OPENAI_API_KEY`. Optionally set `OPENAI_MODEL`.
- OpenAI‑compatible (e.g., Ollama): set `OPENAI_BASE_URL=http://localhost:11434/v1` and any `OPENAI_API_KEY=ollama`.
  *For Ollama/gpt‑oss, tuning options (num_ctx/num_predict/top_k/top_p/etc.) are sent via `extra_body={"options":...}`*

Examples
  # OpenAI
  export OPENAI_API_KEY=sk-...
  python paper_rewriter_plusV2.py --input my.pdf --task rewrite_academic_guided \
    --model gpt-4o-mini --stream --continuation_rounds 2 --max_output_tokens 4096

  # Ollama (OpenAI‑compat), model with large context (e.g., gpt‑oss:120b)
  export OPENAI_BASE_URL=http://localhost:11434/v1
  export OPENAI_API_KEY=ollama
  python paper_rewriter_plusV2.py --input my.pdf --task rewrite_academic_guided \
    --provider openai_compat --model "gpt-oss:120b" --stream \
    --num_ctx 32768 --num_predict 4096 --top_k 50 --top_p 0.9 --temperature 0.2 \
    --context_window 32768 --max_output_tokens 3072 --continuation_rounds 3

Notes
- The similarity indicator (if enabled) is a rough n‑gram overlap proxy (not a plagiarism score).
- For truly long sections, use: large `--context_window`, large `--max_output_tokens`, and `--continuation_rounds`.
- You can pass a custom guideline file via `--guidelines_file` to steer the guided rewrite.
"""

from __future__ import annotations
import argparse
import dataclasses
import logging
import os
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

# ========= Logging =========
logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s')
logger = logging.getLogger('paper_rewriter_plusV2')

# ========= Tokenization =========
try:
    import tiktoken
except Exception:  # pragma: no cover
    logger.error("tiktoken is required. Install with: pip install tiktoken")
    raise

# ========= OpenAI client (works for OpenAI + OpenAI‑compatible via base_url) =========
OPENAI_BASE_URL = os.getenv('OPENAI_BASE_URL')  # e.g. http://localhost:11434/v1 for Ollama
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY', 'sk-local')

try:
    from openai import OpenAI
except Exception:  # pragma: no cover
    OpenAI = None
    logger.warning('OpenAI SDK not installed. pip install openai')

_client = None

def get_client():
    global _client
    if _client is not None:
        return _client
    if OpenAI is None:
        raise RuntimeError('OpenAI SDK not installed. pip install openai')
    if OPENAI_BASE_URL:
        _client = OpenAI(base_url=OPENAI_BASE_URL, api_key=OPENAI_API_KEY)
    else:
        _client = OpenAI(api_key=OPENAI_API_KEY)
    return _client


def is_openai_compat() -> bool:
    url = OPENAI_BASE_URL or ''
    return bool(url and 'api.openai.com' not in url.lower())

# ========= Readers =========
try:
    from docx import Document as _Docx
except Exception:
    _Docx = None

try:
    from pdfminer.high_level import extract_text as pdf_extract_text
except Exception:
    pdf_extract_text = None

try:
    import html2text as _html2text
except Exception:
    _html2text = None

try:
    from pypdf import PdfReader
except Exception:
    PdfReader = None

# ========= Prompts / Guidelines =========

DEFAULT_GUIDELINES = (
    "Write in formal, clear academic prose (active voice, concise sentences).\n"
    "Maintain logical flow: context → method → result → implication; add transitions.\n"
    "Keep claims precise and verifiable; insert citation markers as [1], [2], ... where support is required.\n"
    "Do not fabricate facts, data, or references; if missing, indicate [Ref Needed].\n"
    "Preserve all technical details and symbols; format equations in LaTeX math mode.\n"
    "Be consistent in terminology, units, and notation; expand acronyms on first use.\n"
    "Improve structure with informative topic sentences and signposting.\n"
    "Where figures would help, insert placeholders like [FIGURE: short caption idea].\n"
    "Optimize for readability: break long sentences; prefer parallel structure; avoid redundancy.\n"
)

@dataclass
class PromptLibrary:
    SYSTEM_RESEARCH_ASSISTANT: str = (
        "You are a meticulous research assistant and scientific writing editor. "
        "You write in clear academic prose, keep claims precise, avoid hallucinations, and never invent citations."
    )
    TASKS: Dict[str, str] = dataclasses.field(default_factory=lambda: {
        'rewrite_academic': (
            "Rewrite the provided section into a coherent, engaging academic text in active voice. "
            "Keep technical accuracy; avoid assumptions not supported by the input. "
            "Preserve all core facts; improve organization and readability. "
            "Insert citation markers as [1], [2] where claims need support. Use LaTeX for equations."
        ),
        'rewrite_academic_guided': (
            "Rewrite the section for a high‑quality journal submission, following these guidelines strictly:\n\n"
            "{GUIDELINES}\n\n"
            "Return only the rewritten text (no preambles)."
        ),
        'edit_and_review': (
            "Fix grammar and style; propose a concise, informative section title. "
            "Return:\n1) Title\n2) Polished text\n3) Notes: strengths, risks, missing citations."
        ),
        'summary_cim': (
            'Summarize the section in <=650 characters for an executive memo; preserve key facts and numbers.'
        ),
        'extract_quanti': (
            "Extract all chemicals (with quantities/units), materials, devices/instruments, and statistical outcomes "
            "(e.g., p-values, confidence intervals) into a Markdown table with columns: Category | Item | Value | Unit | Context. "
            "Then list all citations from the text below the table (or 'None')."
        ),
        'outline': (
            'Create a structured outline (H1..H4) for a full paper section based on the text; keep general-purpose.'
        ),
        'figure_suggestions': (
            'Suggest where to place figures. For each: an anchor quote, rationale, and 3 caption variants (Poor/Better/Precise).'
        ),
        'peer_review': (
            'Provide a reviewer-style critique grouped into Major Issues, Minor Issues, and Typos/Edits.'
        ),
    })

# ========= Token helpers =========

def get_encoding(model: str) -> tiktoken.Encoding:
    try:
        return tiktoken.encoding_for_model(model)
    except Exception:
        return tiktoken.get_encoding('cl100k_base')


def count_tokens(text: str, model: str) -> int:
    enc = get_encoding(model)
    return len(enc.encode(text or ''))


def chunk_text(text: str, model: str, max_tokens: int, overlap_tokens: int) -> List[str]:
    enc = get_encoding(model)
    parts = re.split(r"(\n\s*\n)+", (text or '').strip())
    chunks: List[str] = []
    buf: List[str] = []
    buf_toks = 0

    def flush(force=False):
        nonlocal buf, buf_toks
        if buf and (force or buf_toks > 0):
            s = ''.join(buf).strip()
            if s:
                chunks.append(s)
            buf = []
            buf_toks = 0

    for p in parts:
        ptoks = len(enc.encode(p))
        if ptoks > max_tokens:
            toks = enc.encode(p)
            i = 0
            while i < len(toks):
                j = min(i + max_tokens, len(toks))
                piece = enc.decode(toks[i:j])
                if buf_toks:
                    flush(True)
                chunks.append(piece)
                i = max(0, j - overlap_tokens)
            continue
        if buf_toks + ptoks <= max_tokens:
            buf.append(p)
            buf_toks += ptoks
        else:
            flush(True)
            buf = [p]
            buf_toks = ptoks
    flush(True)
    return chunks

# ========= Readers (generic) =========

def read_text_from_docx(path: Path) -> str:
    if _Docx is None:
        raise RuntimeError('python-docx not installed. pip install python-docx')
    doc = _Docx(str(path))
    out: List[str] = []
    for para in doc.paragraphs:
        txt = para.text.strip('\u000b')
        if not txt:
            continue
        style = getattr(para.style, 'name', '') or ''
        if style.startswith('Heading'):
            out.append('\n\n' + txt + '\n\n')
        else:
            out.append(txt)
    return '\n'.join(out)


def read_text_from_pdf(path: Path) -> str:
    if pdf_extract_text is not None:
        try:
            return pdf_extract_text(str(path))
        except Exception:
            pass
    if PdfReader is not None:
        try:
            r = PdfReader(str(path))
            return '\n'.join(page.extract_text() or '' for page in r.pages)
        except Exception as e:
            raise
    raise RuntimeError('No PDF reader available. Install pdfminer.six or pypdf')


def read_text_from_html(path: Path) -> str:
    if _html2text is None:
        raise RuntimeError('html2text not installed. pip install html2text')
    with open(path, 'r', encoding='utf-8', errors='ignore') as f:
        return _html2text.html2text(f.read())


def read_text_generic(path: Path) -> str:
    s = path.suffix.lower()
    if s == '.docx':
        return read_text_from_docx(path)
    if s == '.pdf':
        return read_text_from_pdf(path)
    if s in {'.html', '.htm'}:
        return read_text_from_html(path)
    with open(path, 'r', encoding='utf-8', errors='ignore') as f:
        return f.read()

# ========= Provider‑aware chat call =========

def _compat_default_options(model: str) -> dict:
    """Reasonable defaults for OpenAI‑compat backends. Tweaks for gpt‑oss:120b."""
    m = (model or '').lower()
    opts = {
        'temperature': 0.2,
        'top_p': 0.9,
        'top_k': 50,
        'num_predict': 2048,
    }
    # Heuristic: large context + longer generation for gpt‑oss / 120b models
    if 'gpt-oss' in m or '120b' in m:
        opts.update({'num_ctx': 32768, 'num_predict': 4096})
    return opts


def call_chat_model(
    *,
    task_prompt: str,
    content: str,
    model: str,
    system_prompt: str,
    temperature: float,
    max_tokens: int,
    stream: bool,
    provider: str,
    # compat tuning
    top_p: Optional[float] = None,
    top_k: Optional[int] = None,
    num_ctx: Optional[int] = None,
    num_predict: Optional[int] = None,
    repetition_penalty: Optional[float] = None,
    presence_penalty: Optional[float] = None,
    frequency_penalty: Optional[float] = None,
    stop: Optional[List[str]] = None,
) -> str:
    client = get_client()

    # Provider resolution
    prov = (provider or 'auto').lower()
    if prov == 'auto':
        prov = 'openai_compat' if is_openai_compat() else 'openai'

    messages = [
        {'role': 'system', 'content': system_prompt},
        {'role': 'user', 'content': f'Task: {task_prompt}\n\nInput:\n' + content},
    ]

    kwargs = dict(
        model=model,
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
        stream=stream,
    )

    if prov == 'openai_compat':
        # Build options for compat (e.g., Ollama). Merge defaults with CLI overrides.
        options = _compat_default_options(model)
        # Apply overrides only if provided
        if top_p is not None: options['top_p'] = top_p
        if top_k is not None: options['top_k'] = top_k
        if num_ctx is not None: options['num_ctx'] = num_ctx
        if num_predict is not None: options['num_predict'] = num_predict
        if repetition_penalty is not None: options['repetition_penalty'] = repetition_penalty
        if presence_penalty is not None: options['presence_penalty'] = presence_penalty
        if frequency_penalty is not None: options['frequency_penalty'] = frequency_penalty
        if stop: options['stop'] = stop
        kwargs['extra_body'] = {'options': options}

    # Non‑streaming
    if not stream:
        resp = client.chat.completions.create(**kwargs)
        return (resp.choices[0].message.content or '').strip()

    # Streaming
    out_parts: List[str] = []
    stream_resp = client.chat.completions.create(**kwargs)
    for chunk in stream_resp:
        delta = getattr(chunk.choices[0].delta, 'content', None)
        if delta:
            out_parts.append(delta)
            print(delta, end='', flush=True)
    print()
    return ''.join(out_parts).strip()

# ========= Long‑form continuation =========

def continue_generation(prev_text: str, *, model: str, system_prompt: str, temperature: float, max_tokens: int, overlap_tokens: int) -> str:
    enc = get_encoding(model)
    toks = enc.encode(prev_text)
    anchor_toks = toks[-max(1, min(overlap_tokens, len(toks))):]
    anchor = enc.decode(anchor_toks)

    client = get_client()
    messages = [
        {'role': 'system', 'content': system_prompt},
        {'role': 'user', 'content': (
            'Continue the previous answer seamlessly. Resume right after this anchor (do not repeat any existing text):\n\n'
            + anchor + '\n\nMaintain consistent structure, numbering, and notation.'
        )},
    ]
    resp = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return (resp.choices[0].message.content or '').strip()

# ========= Core processing =========

def process_text(
    raw_text: str,
    *,
    task_key: str,
    model: str,
    prompts: PromptLibrary,
    guidelines_text: Optional[str],
    context_window: int,
    max_output_tokens: int,
    overhead_tokens: int,
    overlap_tokens: int,
    include_similarity: bool,
    continuation_rounds: int,
    continuation_target_tokens: int,
    continuation_overlap_tokens: int,
    temperature: float,
    provider: str,
    # compat tuning
    top_p: Optional[float],
    top_k: Optional[int],
    num_ctx: Optional[int],
    num_predict: Optional[int],
    repetition_penalty: Optional[float],
    presence_penalty: Optional[float],
    frequency_penalty: Optional[float],
    stop: Optional[List[str]],
    stream: bool,
) -> Tuple[str, List[str]]:
    if task_key not in prompts.TASKS:
        raise KeyError(f"Unknown task '{task_key}'. Options: {sorted(prompts.TASKS)}")

    task_prompt = prompts.TASKS[task_key]
    if '{GUIDELINES}' in task_prompt:
        guide = (guidelines_text or DEFAULT_GUIDELINES).strip()
        task_prompt = task_prompt.replace('{GUIDELINES}', guide)

    enc = get_encoding(model)
    max_input_tokens = max(512, context_window - max_output_tokens - overhead_tokens)
    chunks = chunk_text(raw_text, model, max_tokens=max_input_tokens, overlap_tokens=overlap_tokens)

    outputs: List[str] = []
    logger.info("Processing %d chunks with '%s' (provider=%s; input budget=%d, output=%d)",
                len(chunks), task_key, provider, max_input_tokens, max_output_tokens)

    for i, ch in enumerate(chunks, 1):
        logger.info(" → Chunk %d/%d (input tokens=%d)", i, len(chunks), count_tokens(ch, model))
        out = call_chat_model(
            task_prompt=task_prompt,
            content=ch,
            model=model,
            system_prompt=prompts.SYSTEM_RESEARCH_ASSISTANT,
            temperature=temperature,
            max_tokens=max_output_tokens,
            stream=stream,
            provider=provider,
            top_p=top_p,
            top_k=top_k,
            num_ctx=num_ctx,
            num_predict=num_predict,
            repetition_penalty=repetition_penalty,
            presence_penalty=presence_penalty,
            frequency_penalty=frequency_penalty,
            stop=stop,
        )
        out_text = out

        # Continuation rounds for longer output per chunk
        rounds = 0
        while rounds < continuation_rounds:
            if continuation_target_tokens and count_tokens(out_text, model) >= continuation_target_tokens:
                break
            extra = continue_generation(
                out_text,
                model=model,
                system_prompt=prompts.SYSTEM_RESEARCH_ASSISTANT,
                temperature=temperature,
                max_tokens=max_output_tokens,
                overlap_tokens=continuation_overlap_tokens,
            )
            if not extra.strip() or extra.strip() in out_text:
                break
            out_text += "\n\n" + extra.strip()
            rounds += 1

        if include_similarity:
            sim = similarity_ratio(ch, out_text)
            out_text = f"[Similarity≈{sim*100:.1f}% — n-gram overlap proxy]\n\n" + out_text
        outputs.append(out_text)

    stitched = '\n\n'.join(outputs)
    return stitched, outputs

# ========= Similarity (rough proxy) =========

def similarity_ratio(a: str, b: str, n: int = 5) -> float:
    def shingles(s: str) -> set:
        toks = re.findall(r"\w+", (s or '').lower())
        return set(tuple(toks[i:i+n]) for i in range(max(0, len(toks)-n+1)))
    A, B = shingles(a), shingles(b)
    if not A or not B:
        return 0.0
    return len(A & B) / max(1, len(A | B))

# ========= IO =========

def safe_write(path: Path, content: str):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, 'w', encoding='utf-8') as f:
        f.write(content)
    logger.info('Wrote %s', path)

# ========= Pipeline =========

def run_pipeline(
    input_path: Path,
    outdir: Path,
    task: str,
    model: str,
    include_similarity: bool,
    prompts: PromptLibrary,
    guidelines_file: Optional[Path],
    # long‑form knobs
    context_window: int,
    max_output_tokens: int,
    overhead_tokens: int,
    overlap_tokens: int,
    continuation_rounds: int,
    continuation_target_tokens: int,
    continuation_overlap_tokens: int,
    temperature: float,
    provider: str,
    # compat tuning
    top_p: Optional[float],
    top_k: Optional[int],
    num_ctx: Optional[int],
    num_predict: Optional[int],
    repetition_penalty: Optional[float],
    presence_penalty: Optional[float],
    frequency_penalty: Optional[float],
    stop: Optional[List[str]],
    stream: bool,
):
    text = read_text_generic(input_path)

    guidelines_text = None
    if guidelines_file:
        try:
            with open(guidelines_file, 'r', encoding='utf-8', errors='ignore') as f:
                guidelines_text = f.read()
                logger.info('Loaded guidelines from %s', guidelines_file)
        except Exception as e:
            logger.warning('Failed to read guidelines file %s: %s', guidelines_file, e)

    stitched, per_chunks = process_text(
        text,
        task_key=task,
        model=model,
        prompts=prompts,
        guidelines_text=guidelines_text,
        context_window=context_window,
        max_output_tokens=max_output_tokens,
        overhead_tokens=overhead_tokens,
        overlap_tokens=overlap_tokens,
        include_similarity=include_similarity,
        continuation_rounds=continuation_rounds,
        continuation_target_tokens=continuation_target_tokens,
        continuation_overlap_tokens=continuation_overlap_tokens,
        temperature=temperature,
        provider=provider,
        top_p=top_p,
        top_k=top_k,
        num_ctx=num_ctx,
        num_predict=num_predict,
        repetition_penalty=repetition_penalty,
        presence_penalty=presence_penalty,
        frequency_penalty=frequency_penalty,
        stop=stop,
        stream=stream,
    )

    base = outdir / input_path.stem
    safe_write(base.with_suffix(f'.{task}.stitched.txt'), stitched)
    for i, chunk in enumerate(per_chunks):
        safe_write(base.parent / f'{input_path.stem}.{task}.part{i+1:03d}.txt', chunk)

# ========= CLI =========

def build_argparser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description='Paper Rewriter Plus V2 — guided academic rewriting (OpenAI + OpenAI‑compat)')
    p.add_argument('--input', required=True, type=Path, help='Input file (.pdf, .docx, .txt, .md, .html)')
    p.add_argument('--outdir', default=Path('outputs'), type=Path, help='Output directory')
    p.add_argument('--task', default='rewrite_academic_guided', choices=list(PromptLibrary().TASKS.keys()), help='Task key')
    p.add_argument('--model', default=os.getenv('OPENAI_MODEL', 'gpt-4o-mini'), help='Model name (OpenAI or compat)')
    p.add_argument('--provider', default='auto', choices=['auto','openai','openai_compat'], help='Backend provider')
    p.add_argument('--similarity', action='store_true', help='Append similarity proxy vs. source per chunk')
    p.add_argument('--guidelines_file', type=Path, help='Optional path to a custom guidelines .txt file')

    # Long‑form controls
    p.add_argument('--context_window', type=int, default=8192, help='Total context window (tokens) for sizing input chunks')
    p.add_argument('--max_output_tokens', type=int, default=2048, help='Max output tokens per call')
    p.add_argument('--overhead_tokens', type=int, default=500, help='Reserved tokens for system/task/formatting')
    p.add_argument('--overlap_tokens', type=int, default=200, help='Input chunk overlap tokens')
    p.add_argument('--continuation_rounds', type=int, default=0, help='Extra calls to extend each chunk output')
    p.add_argument('--continuation_target_tokens', type=int, default=0, help='Stop early when per‑chunk output hits this size')
    p.add_argument('--continuation_overlap_tokens', type=int, default=80, help='Anchor size (tokens) between continuations')

    # Decoding controls (also forwarded to compat backends via extra_body.options)
    p.add_argument('--temperature', type=float, default=0.2)
    p.add_argument('--top_p', type=float)
    p.add_argument('--top_k', type=int)
    p.add_argument('--num_ctx', type=int, help='Compat: context window hint (e.g., 32768 for gpt‑oss:120b)')
    p.add_argument('--num_predict', type=int, help='Compat: generation length hint')
    p.add_argument('--repetition_penalty', type=float)
    p.add_argument('--presence_penalty', type=float)
    p.add_argument('--frequency_penalty', type=float)
    p.add_argument('--stop', action='append', help='Stop sequence (repeatable)')
    p.add_argument('--stream', action='store_true', help='Stream tokens to stdout while collecting the full output')

    return p


def main(argv: Optional[Sequence[str]] = None):
    args = build_argparser().parse_args(argv)

    if not args.input.exists():
        logger.error('Input not found: %s', args.input)
        sys.exit(2)

    # Friendly defaults for gpt‑oss if user forgot to pass knobs
    if args.provider in ('auto','openai_compat') and is_openai_compat():
        m = (args.model or '').lower()
        if 'gpt-oss' in m or '120b' in m:
            args.context_window = args.context_window or 32768
            if args.num_ctx is None:
                args.num_ctx = 32768
            if args.num_predict is None:
                args.num_predict = 4096
            if args.max_output_tokens < 3072:
                args.max_output_tokens = 3072
            if args.continuation_rounds == 0:
                args.continuation_rounds = 2

    prompts = PromptLibrary()

    run_pipeline(
        input_path=args.input,
        outdir=args.outdir,
        task=args.task,
        model=args.model,
        include_similarity=args.similarity,
        prompts=prompts,
        guidelines_file=args.guidelines_file,
        context_window=args.context_window,
        max_output_tokens=args.max_output_tokens,
        overhead_tokens=args.overhead_tokens,
        overlap_tokens=args.overlap_tokens,
        continuation_rounds=args.continuation_rounds,
        continuation_target_tokens=args.continuation_target_tokens,
        continuation_overlap_tokens=args.continuation_overlap_tokens,
        temperature=args.temperature,
        provider=args.provider,
        top_p=args.top_p,
        top_k=args.top_k,
        num_ctx=args.num_ctx,
        num_predict=args.num_predict,
        repetition_penalty=args.repetition_penalty,
        presence_penalty=args.presence_penalty,
        frequency_penalty=args.frequency_penalty,
        stop=args.stop,
        stream=args.stream,
    )


if __name__ == '__main__':
    main()