Coverage for langbrainscore/utils/logging.py: 59%

1"""Utility functions for logging to stdout and tracking parameters using Weights and Biases""" 

2 

3################################################################ 

4# stuff for type annotations 

5################################################################ 

6import typing 

7 

8# from langbrainscore.interface.cacheable import _Cacheable 

9 

10################################################################ 

11# stuff for logging to W&B (https://wandb.ai) 

12################################################################ 

13 

14 

15def init_wandb(project: str, group: str, use_wandb=True): 

16 """Initializes an instance of Weights & Biases for logging 

17 

18 Args: 

19 project (str): a broad grouping of what project this wandb instance 

20 should be under, e.g., alpha testing, or human-predictive-processing-1 

21 group (str): a more immediate grouping, recommended: by benchmark 

22 use_wandb (bool, optional): if False, we will not initialize W&B. 

23 Mainly for compatibility with CLI. Defaults to True. 

24 """ 

25 if use_wandb: 

26 import wandb 

27 

28 wandb.init( 

29 project=project, 

30 entity="langbrainscore", 

31 # group=group, 

32 ) 

33 

34 

35def log_to_wandb( 

36 *obs: typing.Collection[typing.Union["_Cacheable", dict]], 

37 use_wandb=True, 

38 commit=False, 

39): 

40 """Given a `_Cacheable` object, logs its salient parameters to W&B by 

41 prefixing them with class name. 

42 

43 Args: 

44 ob (typing.Union[_Cacheable, typing.Mapping]): Object to log 

45 use_wandb (bool, optional): if False, we will not log anything. For compatibility with CLI. 

46 Defaults to True. 

47 """ 

48 if use_wandb: 

49 import wandb 

50 

51 # if a singleton object is passed, we convert it to a collection 

52 try: 

53 iter(obs) 

54 except TypeError: 

55 obs = [obs] 

56 

57 for ob in obs: 

58 if hasattr(ob, "params"): 

59 d = {f"{ob.__class__.__name__}/{k}": v for k, v in ob.params.items()} 

60 else: 

61 d = ob 

62 wandb.log(d, commit=False) 

63 

64 wandb.log({}, commit=commit) 

65 

66 

67################################################################ 

68# stuff for logging to the terminal 

69################################################################ 

70import textwrap 

71from datetime import date 

72from sys import stderr, stdout 

73from time import time 

74import shutil 

75import os 

76 

77from colorama import Back, Fore, Style, init 

78from tqdm import tqdm 

79 

80 

81def verbose() -> bool: 

82 """returns True if env variable "VERBOSE" is set to 1""" 

83 return os.environ.get("VERBOSE", None) == "1" 

84 

85 

86init(autoreset=True) 

87_START_TIME = time() 

88 

89 

90def START_TIME(): 

91 return _START_TIME 

92 

93 

94def log(message, cmap="INFO", type=None, verbosity_check=False, **kwargs): 

95 """Utility function to log a `message` to stdout 

96 

97 Args: 

98 message (typing.Any): an object that supports `__str__()` 

99 cmap (str, optional): what colormap to use. "INFO" corresponds to blue, 

100 "WARN" Defaults to "INFO". 

101 type (str, optional): Type of message, for user knowledge. If provided, will be used as the 

102 tag for this output (e.g. "info"). If no value is provided, the same string as `cmap` is 

103 used as the tag. Defaults to None. 

104 verbosity_check (bool, optional): Whether to check for a "VERBOSE" environment flag before 

105 outputting. If false, always output text regardless of verbosity setting. Defaults to False. 

106 """ 

107 if verbosity_check and not verbose(): 

108 return 

109 

110 _message = str(message) 

111 

112 class T: 

113 HEADER = "\033[95m" 

114 OKBLUE = "\033[94m" 

115 OKCYAN = "\033[96m" 

116 OKGREEN = "\033[92m" 

117 WARNING = "\033[93m" 

118 FAIL = "\033[91m" 

119 ENDC = "\033[0m" 

120 BOLD = "\033[1m" 

121 UNDERLINE = "\033[4m" 

122 

123 if cmap == "INFO": 

124 c = T.OKBLUE 

125 elif cmap == "WARN": 

126 c = T.BOLD + T.WARNING 

127 elif cmap == "ANNOUNCE": 

128 c = T.BOLD + T.OKGREEN 

129 elif cmap == "ERR": 

130 c = "\n" + T.BOLD + T.FAIL 

131 else: 

132 c = T.OKBLUE 

133 

134 timestamp = f"{time() - START_TIME():.2f}s" 

135 lines = textwrap.wrap( 

136 _message + T.ENDC, 

137 width=shutil.get_terminal_size((120, 24))[0] - 1, 

138 initial_indent=c + "%" * 3 + f" [{type or cmap.lower()} @ {timestamp}] ", 

139 subsequent_indent=". " * 6 + "", 

140 ) 

141 tqdm.write("\n".join(lines), file=stderr) 

142 # print(*lines, sep='\n', file=stderr)