File size: 5,919 Bytes
ddbc0c8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
"""Quiet-mode utilities for notebooks and demos.

Goals:
- Suppress noisy logs (INFO/DEBUG), progress bars, and common framework warnings.
- Keep meaningful outputs (tables/figures/explicit prints) visible.
- Provide an opt-in context manager to silence extremely noisy calls.

Usage:
    from scripts.quiet import install_quiet
    install_quiet()

For a particularly noisy block:
    from scripts.quiet import suppress_output
    with suppress_output():
        do_noisy_work()
"""

from __future__ import annotations

import logging
import os
import sys
import warnings
from contextlib import contextmanager
from typing import Iterator, Optional


def install_quiet(
    *,
    python_warnings: bool = True,
    logging_level: int = logging.ERROR,
    silence_tqdm: bool = True,
    silence_hf_progress: bool = True,
    silence_tensorflow: bool = True,
    message_only_loggers: Optional[list[str]] = ['legalrag.retrieval.graph_store'],
    silence_colbert_stdout: bool = True,
) -> None:
    """Install a set of conservative defaults to reduce notebook noise.

    Parameters
    ----------
    python_warnings:
        If True, suppresses common warning categories that frequently clutter demo notebooks.
    logging_level:
        Root logging level to set (default ERROR). This suppresses INFO/DEBUG logs such as
        '2025-.. - __main__ - INFO - ...'.
    silence_tqdm:
        If True, tries to disable tqdm progress bars via environment variables.
    silence_hf_progress:
        If True, disables Hugging Face advisory warnings and some progress output.
    silence_tensorflow:
        If True, reduces TensorFlow / XLA / absl logging chatter where possible.
    """
    if python_warnings:
        # The jieba 'invalid escape sequence' lines come from SyntaxWarning during import.
        warnings.filterwarnings("ignore", category=SyntaxWarning)
        # Reduce general deprecation/runtime warnings that are rarely useful in demos.
        warnings.filterwarnings("ignore", category=DeprecationWarning)
        warnings.filterwarnings("ignore", category=FutureWarning)
        warnings.filterwarnings("ignore", category=UserWarning, message=r".*progress bar.*")
        # add more filters here 

    # Root logger: suppress INFO/DEBUG timestamped lines.
    logging.basicConfig(level=logging_level)
    logging.getLogger().setLevel(logging_level)


    for name in [
        "legalrag",
        "legalrag.retrieval",
        "legalrag.retrieval.graph_store",
        "legalrag.retrieval.hybrid_retriever",
        "legalrag.retrieval.bm25_retriever",
        "legalrag.pipeline.rag_pipeline",
        "legalrag.pipeline",
        "legalrag.api",
        "jieba",
        "transformers",
        "sentence_transformers",
        "faiss",
        "tensorflow",
        "torch",
        "urllib3",
    ]:
        logging.getLogger(name).setLevel(logging_level)


    for name in list(logging.root.manager.loggerDict.keys()):
        if name == "legalrag" or name.startswith("legalrag."):
            logging.getLogger(name).setLevel(logging_level)

    if message_only_loggers:
        for name in message_only_loggers:
            lg = logging.getLogger(name)
            lg.setLevel(logging.INFO)
            for h in lg.handlers:
                h.setFormatter(logging.Formatter("%(message)s"))
            lg.propagate = False

    if silence_colbert_stdout:
        try:
            from legalrag.retrieval.colbert_retriever import ColBERTRetriever  
            _orig_init = ColBERTRetriever._init_searcher
            _orig_search = ColBERTRetriever.search

            def _init_quiet(self):
                with suppress_output():
                    return _orig_init(self)

            def _search_quiet(self, *args, **kwargs):
                with suppress_output():
                    return _orig_search(self, *args, **kwargs)

            ColBERTRetriever._init_searcher = _init_quiet  
            ColBERTRetriever.search = _search_quiet   
        except Exception:
            pass

    if silence_tqdm:
        # Many libs respect these toggles.
        os.environ.setdefault("TQDM_DISABLE", "1")

    if silence_hf_progress:
        os.environ.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1")
        os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
        os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
        os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1")
        os.environ.setdefault("DISABLE_TORCHVISION", "1")
        os.environ.setdefault("TRANSFORMERS_NO_IMAGE_PROCESSING", "1")

    if silence_tensorflow:
        # TensorFlow C++ log level: 0 = all, 1 = INFO, 2 = WARNING, 3 = ERROR
        os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
        os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
        os.environ.setdefault("USE_TF", "0")

    # If transformers is imported, we can try to disable its verbosity.
    try:
        from transformers.utils import logging as hf_logging   
        hf_logging.set_verbosity_error()
        hf_logging.disable_progress_bar()
    except Exception:
        pass


@contextmanager
def suppress_output(*, stdout: bool = True, stderr: bool = True) -> Iterator[None]:
    """Temporarily redirect stdout/stderr to os.devnull.

    Use this sparingly—only for extremely noisy calls (e.g., one-time index builds).
    """
    devnull = open(os.devnull, "w")
    old_stdout: Optional[object] = None
    old_stderr: Optional[object] = None
    try:
        if stdout:
            old_stdout = sys.stdout
            sys.stdout = devnull  
        if stderr:
            old_stderr = sys.stderr
            sys.stderr = devnull  
        yield
    finally:
        try:
            if stdout and old_stdout is not None:
                sys.stdout = old_stdout  
            if stderr and old_stderr is not None:
                sys.stderr = old_stderr  
        finally:
            devnull.close()