| """TensorBoard events without requiring PyTorch or TensorFlow.""" |
| import json |
| import re |
| import time |
| from pathlib import Path |
|
|
| EM = re.compile(r"EM sub_iter=(\d+) size=(\d+) obj=([\deE+.-]+) num_tokens=(\d+) (?:num_)?tokens/piece=([\deE+.-]+)") |
| STAGES = {"checking_inputs":0,"training":1,"measuring":2,"encoding":2,"merging":3,"verifying_inputs":3,"complete":4,"failed":-1} |
|
|
|
|
| class RunMonitor: |
| def __init__(self, output, log_dir, enabled): |
| self.output = output |
| self.started = time.monotonic() |
| self.sequence = self.em_step = 0 |
| self.writer = None |
| if enabled: |
| try: |
| from tensorboard.summary.writer.event_file_writer import EventFileWriter |
| except ImportError: |
| raise RuntimeError("请先安装 requirements.txt 中的 TensorBoard 依赖。") from None |
| self.writer = EventFileWriter(str(log_dir),flush_secs=5) |
| self.log_dir = log_dir |
|
|
| def scalar(self, tag, value, step): |
| if self.writer: |
| from tensorboard.compat.proto.event_pb2 import Event |
| from tensorboard.compat.proto.summary_pb2 import Summary |
| self.writer.add_event(Event(wall_time=time.time(),step=step, |
| summary=Summary(value=[Summary.Value(tag=tag,simple_value=float(value))]))) |
|
|
| def update(self, stage, **detail): |
| elapsed = time.monotonic()-self.started |
| state = {"stage":stage,"elapsed_seconds":elapsed,"tensorboard_dir":str(self.log_dir),**detail} |
| temporary = self.output/'progress.tmp' |
| temporary.write_text(json.dumps(state,ensure_ascii=False,indent=2)+'\n',encoding='utf-8') |
| temporary.replace(self.output/'progress.json') |
| self.scalar('run/stage',STAGES[stage],self.sequence) |
| self.scalar('run/elapsed_seconds',elapsed,self.sequence) |
| if 'preflight_completed_sentences' in detail: |
| self.scalar('inputs/checked_sentences',detail['preflight_completed_sentences'],self.sequence) |
| self.scalar('inputs/preflight_percent',100*detail['preflight_completed_sentences']/detail['preflight_total_sentences'],self.sequence) |
| if 'completed_sentences' in detail: |
| done,total = detail['completed_sentences'],detail['total_sentences'] |
| prefix = stage if stage in {'encoding','merging'} else ('measurement' if stage=='measuring' else 'run') |
| self.scalar(prefix+'/completed_sentences',done,self.sequence) |
| self.scalar(prefix+'/percent',100*done/total if total else 100,self.sequence) |
| if stage=='encoding': |
| self.scalar('encoding/sentences_per_second',detail.get('encoding_sentences_per_second',0),self.sequence) |
| if detail.get('encoding_eta_seconds') is not None: |
| self.scalar('encoding/eta_seconds',detail['encoding_eta_seconds'],self.sequence) |
| elif stage=='measuring': |
| seconds = detail.get('measurement_elapsed_seconds',0) |
| self.scalar('measurement/sentences_per_second',done/seconds if seconds else 0,self.sequence) |
| if self.writer: self.writer.flush() |
| self.sequence += 1 |
|
|
| def native_line(self, line): |
| match = EM.search(line) |
| if not match: return |
| sub_iter,size,obj,tokens,ratio = match.groups() |
| for name,value in [('em_objective',obj),('candidate_pieces',size),('tokens_per_piece',ratio),('segmented_tokens',tokens)]: |
| self.scalar('tokenizer/'+name,float(value),self.em_step) |
| self.update('training',em_updates=self.em_step+1,em_sub_iteration=int(sub_iter), |
| candidate_pieces=int(size),em_objective=float(obj)) |
| self.em_step += 1 |
|
|
| def final_statistics(self, statistics): |
| for split,stats in statistics.items(): |
| for name in ['characters_per_token','mean_tokens_per_sentence','byte_fallback_token_fraction','roundtrip_mismatches','unknown_tokens']: |
| value = stats[name] |
| if value is not None: self.scalar(f'corpus/{split}/{name}',value,0) |
| if self.writer: self.writer.flush() |
|
|
| def __enter__(self): return self |
|
|
| def __exit__(self, kind, error, traceback): |
| try: |
| if error is not None: self.update('failed',error_type=kind.__name__,error=str(error)) |
| finally: |
| if self.writer: self.writer.close() |
|
|
