Implemented a second LLM stage for grammatical review

This commit is contained in:
2026-04-21 15:42:09 -05:00
parent ca01e46d77
commit 445329de81
13 changed files with 620 additions and 129 deletions

View File

@@ -32,9 +32,19 @@ def _build_parser() -> argparse.ArgumentParser:
process.add_argument("--model", help="OpenRouter model to use")
process.add_argument("--base-url", help="OpenAI-compatible API base URL")
process.add_argument("--max-section-tokens", type=int, help="maximum estimated tokens per transcript section")
process.add_argument("--confidence-threshold", type=float, help="minimum confidence required to apply a correction")
process.add_argument(
"--glossary-confidence-threshold",
type=float,
help="minimum confidence required to apply a glossary correction",
)
process.add_argument(
"--grammar-confidence-threshold",
type=float,
help="minimum confidence required to apply a grammar correction",
)
process.add_argument("--max-retries", type=int, help="maximum Instructor retries for structured response validation")
process.add_argument("--glossary-max-llm-passes", type=int, help="maximum total LLM passes for glossary corrections")
process.add_argument("--grammar-max-llm-passes", type=int, help="maximum total LLM passes for grammar corrections")
process.add_argument("--work-dir", type=Path, help="directory for per-run scratch diagnostics")
return parser
@@ -46,9 +56,11 @@ def _process(args: argparse.Namespace) -> int:
model=args.model,
base_url=args.base_url,
max_section_tokens=args.max_section_tokens,
confidence_threshold=args.confidence_threshold,
glossary_confidence_threshold=args.glossary_confidence_threshold,
grammar_confidence_threshold=args.grammar_confidence_threshold,
max_retries=args.max_retries,
glossary_max_llm_passes=args.glossary_max_llm_passes,
grammar_max_llm_passes=args.grammar_max_llm_passes,
work_dir=args.work_dir,
)
)

View File

@@ -9,10 +9,12 @@ from .errors import AuditaConfigError
DEFAULT_MODEL = "openrouter/mistralai/mistral-small-3.2-24b-instruct"
DEFAULT_BASE_URL = "https://openrouter.ai/api/v1"
DEFAULT_MAX_SECTION_TOKENS = 16000
DEFAULT_CONFIDENCE_THRESHOLD = 0.60
DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD = 0.60
DEFAULT_GRAMMAR_CONFIDENCE_THRESHOLD = 0.60
DEFAULT_MAX_RETRIES = 3
DEFAULT_WORK_DIR = "/tmp/audita"
DEFAULT_GLOSSARY_MAX_LLM_PASSES = 3
DEFAULT_GRAMMAR_MAX_LLM_PASSES = 3
@dataclass(frozen=True)
@@ -20,9 +22,11 @@ class ConfigOverrides:
model: Optional[str] = None
base_url: Optional[str] = None
max_section_tokens: Optional[int] = None
confidence_threshold: Optional[float] = None
glossary_confidence_threshold: Optional[float] = None
grammar_confidence_threshold: Optional[float] = None
max_retries: Optional[int] = None
glossary_max_llm_passes: Optional[int] = None
grammar_max_llm_passes: Optional[int] = None
work_dir: Optional[Path] = None
@@ -32,9 +36,11 @@ class AuditaConfig:
model: str = DEFAULT_MODEL
base_url: str = DEFAULT_BASE_URL
max_section_tokens: int = DEFAULT_MAX_SECTION_TOKENS
confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD
glossary_confidence_threshold: float = DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD
grammar_confidence_threshold: float = DEFAULT_GRAMMAR_CONFIDENCE_THRESHOLD
max_retries: int = DEFAULT_MAX_RETRIES
glossary_max_llm_passes: int = DEFAULT_GLOSSARY_MAX_LLM_PASSES
grammar_max_llm_passes: int = DEFAULT_GRAMMAR_MAX_LLM_PASSES
work_dir: Path = Path(DEFAULT_WORK_DIR)
@classmethod
@@ -55,11 +61,17 @@ class AuditaConfig:
DEFAULT_MAX_SECTION_TOKENS,
"AUDITA_MAX_SECTION_TOKENS",
)
confidence_threshold = _select_float(
selected.confidence_threshold,
source.get("AUDITA_CONFIDENCE_THRESHOLD"),
DEFAULT_CONFIDENCE_THRESHOLD,
"AUDITA_CONFIDENCE_THRESHOLD",
glossary_confidence_threshold = _select_float(
selected.glossary_confidence_threshold,
source.get("AUDITA_GLOSSARY_CONFIDENCE_THRESHOLD"),
DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD,
"AUDITA_GLOSSARY_CONFIDENCE_THRESHOLD",
)
grammar_confidence_threshold = _select_float(
selected.grammar_confidence_threshold,
source.get("AUDITA_GRAMMAR_CONFIDENCE_THRESHOLD"),
DEFAULT_GRAMMAR_CONFIDENCE_THRESHOLD,
"AUDITA_GRAMMAR_CONFIDENCE_THRESHOLD",
)
max_retries = _select_int(
selected.max_retries,
@@ -73,6 +85,12 @@ class AuditaConfig:
DEFAULT_GLOSSARY_MAX_LLM_PASSES,
"AUDITA_GLOSSARY_MAX_LLM_PASSES",
)
grammar_max_llm_passes = _select_int(
selected.grammar_max_llm_passes,
source.get("AUDITA_GRAMMAR_MAX_LLM_PASSES"),
DEFAULT_GRAMMAR_MAX_LLM_PASSES,
"AUDITA_GRAMMAR_MAX_LLM_PASSES",
)
work_dir_value = selected.work_dir or Path(source.get("AUDITA_WORK_DIR") or DEFAULT_WORK_DIR)
config = cls(
@@ -80,9 +98,11 @@ class AuditaConfig:
model=model,
base_url=base_url,
max_section_tokens=max_section_tokens,
confidence_threshold=confidence_threshold,
glossary_confidence_threshold=glossary_confidence_threshold,
grammar_confidence_threshold=grammar_confidence_threshold,
max_retries=max_retries,
glossary_max_llm_passes=glossary_max_llm_passes,
grammar_max_llm_passes=grammar_max_llm_passes,
work_dir=Path(work_dir_value),
)
config.validate()
@@ -97,12 +117,16 @@ class AuditaConfig:
raise AuditaConfigError("AUDITA_BASE_URL must not be empty.")
if self.max_section_tokens <= 0:
raise AuditaConfigError("AUDITA_MAX_SECTION_TOKENS must be greater than zero.")
if not 0.0 <= self.confidence_threshold <= 1.0:
raise AuditaConfigError("AUDITA_CONFIDENCE_THRESHOLD must be between 0.0 and 1.0.")
if not 0.0 <= self.glossary_confidence_threshold <= 1.0:
raise AuditaConfigError("AUDITA_GLOSSARY_CONFIDENCE_THRESHOLD must be between 0.0 and 1.0.")
if not 0.0 <= self.grammar_confidence_threshold <= 1.0:
raise AuditaConfigError("AUDITA_GRAMMAR_CONFIDENCE_THRESHOLD must be between 0.0 and 1.0.")
if self.max_retries < 0:
raise AuditaConfigError("AUDITA_MAX_RETRIES must be greater than or equal to zero.")
if self.glossary_max_llm_passes < 1:
raise AuditaConfigError("AUDITA_GLOSSARY_MAX_LLM_PASSES must be greater than or equal to one.")
if self.grammar_max_llm_passes < 1:
raise AuditaConfigError("AUDITA_GRAMMAR_MAX_LLM_PASSES must be greater than or equal to one.")
def _get_required_env(env: Mapping[str, str], name: str) -> str:

View File

@@ -1,9 +1,11 @@
from dataclasses import asdict, dataclass
from typing import Dict, Iterable, List, Optional, Tuple
from typing import Dict, Iterable, List, Literal, Optional, Tuple
from .errors import AuditaValidationError
from .schemas import CorrectionCandidate, TranscriptSegment
ReplacementMode = Literal["replace_all", "require_unique"]
@dataclass(frozen=True)
class SkippedCorrection:
@@ -30,9 +32,12 @@ def apply_corrections(
transcript: List[TranscriptSegment],
corrections: Iterable[CorrectionCandidate],
confidence_threshold: float,
replacement_mode: ReplacementMode = "replace_all",
) -> CorrectionApplicationResult:
if not 0.0 <= confidence_threshold <= 1.0:
raise AuditaValidationError("Confidence threshold must be between 0.0 and 1.0.")
if replacement_mode not in ("replace_all", "require_unique"):
raise AuditaValidationError("Replacement mode must be replace_all or require_unique.")
revised = list(transcript)
id_to_position = _id_to_position(revised)
@@ -44,7 +49,7 @@ def apply_corrections(
ignored_ids.append(correction.id)
continue
reason, actual_text = _target_error(revised, id_to_position, correction)
reason, actual_text = _target_error(revised, id_to_position, correction, replacement_mode)
if reason is not None:
skipped.append(_skip(correction, reason, actual_text=actual_text))
continue
@@ -67,6 +72,7 @@ def _target_error(
transcript: List[TranscriptSegment],
id_to_position: Dict[int, int],
correction: CorrectionCandidate,
replacement_mode: ReplacementMode,
) -> Tuple[Optional[str], Optional[str]]:
if correction.id not in id_to_position:
return "id does not exist in transcript", None
@@ -80,6 +86,8 @@ def _target_error(
match_count = segment.text.count(correction.original_text)
if match_count == 0:
return "original_text does not match any substring in segment text", segment.text
if replacement_mode == "require_unique" and match_count > 1:
return "original_text appears more than once in segment text", segment.text
return None, None

View File

@@ -4,7 +4,7 @@ from typing import List, Protocol
from .chunking import TranscriptSection
from .config import AuditaConfig
from .prompts import build_glossary_correction_messages
from .prompts import build_glossary_correction_messages, build_grammar_correction_messages
from .schemas import CorrectionCandidate, CorrectionSet, Glossary
@@ -45,3 +45,25 @@ class GlossaryCorrectionPass:
response_path = run_dir / f"corrections-{section.section_index:04d}.json"
response_path.write_text(response.model_dump_json(indent=2) + "\n", encoding="utf-8")
return list(response.corrections)
class GrammarCorrectionPass:
def __init__(self, llm_client: LLMClient) -> None:
self._llm_client = llm_client
def run(
self,
section: TranscriptSection,
glossary: Glossary,
config: AuditaConfig,
run_dir: Path,
retry_pass: bool = False,
) -> List[CorrectionCandidate]:
messages = build_grammar_correction_messages(section, glossary, retry_pass=retry_pass)
prompt_path = run_dir / f"prompt-{section.section_index:04d}.json"
prompt_path.write_text(json.dumps(messages, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
response = self._llm_client.create_corrections(messages, config)
response_path = run_dir / f"corrections-{section.section_index:04d}.json"
response_path.write_text(response.model_dump_json(indent=2) + "\n", encoding="utf-8")
return list(response.corrections)

View File

@@ -1,21 +1,31 @@
import json
import shutil
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Callable, Dict, List, Optional
from typing import Callable, Dict, List, Optional, Tuple
from uuid import uuid4
from .chunking import IndexedSegment, TranscriptSection, chunk_indexed_segments
from .config import AuditaConfig
from .corrections import SkippedCorrection, apply_corrections
from .corrections import ReplacementMode, SkippedCorrection, apply_corrections
from .errors import AuditaError
from .passes import GlossaryCorrectionPass, LLMClient
from .passes import CorrectionPass, GlossaryCorrectionPass, GrammarCorrectionPass, LLMClient
from .schemas import Glossary, TranscriptSegment, parse_transcript_json
ProgressCallback = Callable[[str], None]
@dataclass(frozen=True)
class StageSpec:
name: str
correction_pass: CorrectionPass
max_llm_passes: int
confidence_threshold: float
replacement_mode: ReplacementMode
def process_transcript(
transcript: List[TranscriptSegment],
glossary: Glossary,
@@ -26,8 +36,8 @@ def process_transcript(
run_dir = _create_run_dir(config.work_dir)
try:
_log(progress, f"Created work directory {run_dir}")
pass_summaries: List[dict] = []
_write_run_metadata(run_dir, config, pass_summaries)
stage_summaries: List[dict] = []
_write_run_metadata(run_dir, config, stage_summaries)
if llm_client is None:
from .llm import InstructorLLMClient
@@ -35,87 +45,56 @@ def process_transcript(
llm_client = InstructorLLMClient(config)
working = list(transcript)
correction_pass = GlossaryCorrectionPass(llm_client)
unresolved_retry_skips: Dict[int, SkippedCorrection] = {}
final_nonretry_skips: List[SkippedCorrection] = []
for pass_number in range(1, config.glossary_max_llm_passes + 1):
if pass_number == 1:
indexed_segments = _indexed_segments_for_ids(working, [segment.id for segment in working])
else:
retry_ids = sorted(unresolved_retry_skips)
if not retry_ids:
break
indexed_segments = _indexed_segments_for_ids(working, retry_ids)
if not indexed_segments:
break
pass_dir = run_dir / f"pass-{pass_number:04d}"
pass_dir.mkdir()
sections = chunk_indexed_segments(indexed_segments, config.max_section_tokens)
corrections = []
for section in sections:
_write_and_validate_section(pass_dir, section)
_log(
progress,
f"Processing glossary pass {pass_number}/{config.glossary_max_llm_passes} "
f"section {section.section_index + 1}/{len(sections)} "
f"({len(section.segments)} segments, estimated {section.token_count} tokens)",
)
corrections.extend(
correction_pass.run(
section,
glossary,
config,
pass_dir,
retry_pass=pass_number > 1,
)
)
application_result = apply_corrections(
working,
corrections,
config.confidence_threshold,
)
working = application_result.transcript
for correction_id in application_result.applied_ids:
unresolved_retry_skips.pop(correction_id, None)
for correction_id in application_result.ignored_ids:
unresolved_retry_skips.pop(correction_id, None)
for skipped in application_result.skipped:
if _is_retryable_skip(skipped, working):
unresolved_retry_skips[skipped.id] = skipped
else:
final_nonretry_skips.append(skipped)
pass_summaries.append(
{
"pass_number": pass_number,
"retry_pass": pass_number > 1,
"section_count": len(sections),
"segment_count": len(indexed_segments),
"corrections_returned": len(corrections),
"applied_count": len(application_result.applied_ids),
"ignored_below_threshold_count": len(application_result.ignored_ids),
"skipped_count": len(application_result.skipped),
"retry_segment_count": len(unresolved_retry_skips),
}
)
_write_run_metadata(run_dir, config, pass_summaries)
if not unresolved_retry_skips:
break
final_skipped = final_nonretry_skips + [
unresolved_retry_skips[correction_id] for correction_id in sorted(unresolved_retry_skips)
stages = [
StageSpec(
name="glossary",
correction_pass=GlossaryCorrectionPass(llm_client),
max_llm_passes=config.glossary_max_llm_passes,
confidence_threshold=config.glossary_confidence_threshold,
replacement_mode="replace_all",
),
StageSpec(
name="grammar",
correction_pass=GrammarCorrectionPass(llm_client),
max_llm_passes=config.grammar_max_llm_passes,
confidence_threshold=config.grammar_confidence_threshold,
replacement_mode="require_unique",
),
]
final_skipped: List[Tuple[str, SkippedCorrection]] = []
for stage in stages:
stage_dir = run_dir / stage.name
stage_dir.mkdir()
stage_summary = {
"stage": stage.name,
"max_llm_passes": stage.max_llm_passes,
"confidence_threshold": stage.confidence_threshold,
"replacement_mode": stage.replacement_mode,
"passes": [],
}
stage_summaries.append(stage_summary)
_write_run_metadata(run_dir, config, stage_summaries)
working, stage_skipped = _run_correction_stage(
working,
glossary,
config,
stage,
run_dir,
stage_dir,
stage_summaries,
stage_summary["passes"],
progress,
)
final_skipped.extend((stage.name, skipped) for skipped in stage_skipped)
_write_run_metadata(run_dir, config, stage_summaries)
_write_skipped_corrections(run_dir, final_skipped)
for skipped in final_skipped:
for stage_name, skipped in final_skipped:
_log(
progress,
f"Skipping correction for id {skipped.id}: {skipped.reason}",
f"Skipping {stage_name} correction for id {skipped.id}: {skipped.reason}",
)
revised = _sort_transcript_chronologically(working)
except Exception as exc:
@@ -132,6 +111,97 @@ def process_transcript(
return revised
def _run_correction_stage(
transcript: List[TranscriptSegment],
glossary: Glossary,
config: AuditaConfig,
stage: StageSpec,
run_dir: Path,
stage_dir: Path,
stage_summaries: List[dict],
pass_summaries: List[dict],
progress: Optional[ProgressCallback],
) -> Tuple[List[TranscriptSegment], List[SkippedCorrection]]:
working = list(transcript)
unresolved_retry_skips: Dict[int, SkippedCorrection] = {}
final_nonretry_skips: List[SkippedCorrection] = []
for pass_number in range(1, stage.max_llm_passes + 1):
if pass_number == 1:
indexed_segments = _indexed_segments_for_ids(working, [segment.id for segment in working])
else:
retry_ids = sorted(unresolved_retry_skips)
if not retry_ids:
break
indexed_segments = _indexed_segments_for_ids(working, retry_ids)
if not indexed_segments:
break
pass_dir = stage_dir / f"pass-{pass_number:04d}"
pass_dir.mkdir()
sections = chunk_indexed_segments(indexed_segments, config.max_section_tokens)
corrections = []
for section in sections:
_write_and_validate_section(pass_dir, section)
_log(
progress,
f"Processing {stage.name} pass {pass_number}/{stage.max_llm_passes} "
f"section {section.section_index + 1}/{len(sections)} "
f"({len(section.segments)} segments, estimated {section.token_count} tokens)",
)
corrections.extend(
stage.correction_pass.run(
section,
glossary,
config,
pass_dir,
retry_pass=pass_number > 1,
)
)
application_result = apply_corrections(
working,
corrections,
stage.confidence_threshold,
replacement_mode=stage.replacement_mode,
)
working = application_result.transcript
for correction_id in application_result.applied_ids:
unresolved_retry_skips.pop(correction_id, None)
for correction_id in application_result.ignored_ids:
unresolved_retry_skips.pop(correction_id, None)
for skipped in application_result.skipped:
if _is_retryable_skip(skipped, working):
unresolved_retry_skips[skipped.id] = skipped
else:
final_nonretry_skips.append(skipped)
pass_summaries.append(
{
"pass_number": pass_number,
"retry_pass": pass_number > 1,
"section_count": len(sections),
"segment_count": len(indexed_segments),
"corrections_returned": len(corrections),
"applied_count": len(application_result.applied_ids),
"ignored_below_threshold_count": len(application_result.ignored_ids),
"skipped_count": len(application_result.skipped),
"retry_segment_count": len(unresolved_retry_skips),
}
)
_write_run_metadata(run_dir, config, stage_summaries)
if not unresolved_retry_skips:
break
final_skipped = final_nonretry_skips + [
unresolved_retry_skips[correction_id] for correction_id in sorted(unresolved_retry_skips)
]
return working, final_skipped
def _create_run_dir(work_dir: Path) -> Path:
work_dir.mkdir(parents=True, exist_ok=True)
timestamp = datetime.utcnow().strftime("%Y%m%dT%H%M%SZ")
@@ -143,16 +213,18 @@ def _create_run_dir(work_dir: Path) -> Path:
def _write_run_metadata(
run_dir: Path,
config: AuditaConfig,
pass_summaries: List[dict],
stage_summaries: List[dict],
) -> None:
metadata = {
"model": config.model,
"base_url": config.base_url,
"max_section_tokens": config.max_section_tokens,
"confidence_threshold": config.confidence_threshold,
"glossary_confidence_threshold": config.glossary_confidence_threshold,
"grammar_confidence_threshold": config.grammar_confidence_threshold,
"max_retries": config.max_retries,
"glossary_max_llm_passes": config.glossary_max_llm_passes,
"passes": pass_summaries,
"grammar_max_llm_passes": config.grammar_max_llm_passes,
"stages": stage_summaries,
}
(run_dir / "metadata.json").write_text(
json.dumps(metadata, ensure_ascii=False, indent=2) + "\n",
@@ -167,11 +239,15 @@ def _write_and_validate_section(run_dir: Path, section: TranscriptSection) -> No
parse_transcript_json(section_json, require_sequential_ids=False)
def _write_skipped_corrections(run_dir: Path, skipped: List[SkippedCorrection]) -> None:
def _write_skipped_corrections(run_dir: Path, skipped: List[Tuple[str, SkippedCorrection]]) -> None:
skipped_path = run_dir / "skipped-corrections.json"
skipped_path.write_text(
json.dumps(
{"skipped_corrections": [item.to_dict() for item in skipped]},
{
"skipped_corrections": [
{"stage": stage_name, **item.to_dict()} for stage_name, item in skipped
]
},
ensure_ascii=False,
indent=2,
)

View File

@@ -57,3 +57,49 @@ def build_glossary_correction_messages(
f"Transcript section:\n{section_json}"
)
return [{"role": "system", "content": system}, {"role": "user", "content": user}]
def build_grammar_correction_messages(
section: TranscriptSection,
glossary: Glossary,
retry_pass: bool = False,
) -> List[Message]:
glossary_json = json.dumps(glossary.model_dump(mode="json"), ensure_ascii=False, indent=2)
section_json = json.dumps(section.prompt_payload(), ensure_ascii=False, indent=2)
system = (
"You are Audita, a conservative transcript readability assistant. "
"Improve readability only where the change preserves the speaker's words and meaning. "
"Allowed changes are capitalization, commas, periods, em dashes, ellipses, homophone fixes, "
"and spelling fixes. Do not paraphrase, summarize, reorder words, or change content."
)
retry_guidance = ""
if retry_pass:
retry_guidance = (
"Retry guidance:\n"
"These segments are being retried because previous correction spans did not apply cleanly. "
"Copy original_text exactly from the current segment text, using only a span that appears exactly once.\n\n"
)
user = (
"Review this transcript section and return only readability corrections that should be applied.\n\n"
f"{retry_guidance}"
"Rules:\n"
"- Allowed corrections are only capitalization changes, punctuation changes involving commas, periods, em dashes, and ellipses, homophone fixes, and spelling fixes.\n"
"- Do not add, remove, reorder, or replace words except for clear homophone or spelling corrections that preserve the spoken content.\n"
"- Do not paraphrase, summarize, clarify, smooth style, or change the speaker's intent or meaning.\n"
"- Treat the glossary as protected vocabulary and context; do not introduce new glossary substitutions during this grammar pass.\n"
"- Use the exact id from the input segment.\n"
"- For returned corrections, original_text must be only the exact text span that needs replacement, not the full segment text unless the whole segment is the replacement span.\n"
"- Choose an original_text span that appears exactly once in the current segment text.\n"
"- corrected_text must be only the replacement text for that span, not the full corrected segment text unless the whole segment is the replacement span.\n"
"- Each returned correction must contain only id, original_text, corrected_text, and confidence.\n"
"- Do not return corrections where original_text and corrected_text are identical.\n"
"- Do not return speaker, start, or end fields.\n"
"- Return only changed segments; do not return entries for unchanged segments.\n"
"- confidence must be between 0.0 and 1.0.\n"
"- If no corrections are needed, return an empty corrections list.\n\n"
f"Protected glossary/context:\n{glossary_json}\n\n"
f"Transcript section:\n{section_json}"
)
return [{"role": "system", "content": system}, {"role": "user", "content": user}]