beichen0426/olympiads-ref

1import re2import json3 4from tqdm import tqdm5from loguru import logger6 7from pathlib import Path8from typing import Tuple, List9 10 11project_root = Path(__file__).parent.parent.parent12problem_tag = 'Problem'13solution_tag = 'Solution'14problem_pattern = re.compile(r'(?:\n|# )Problem\s+(\d+)(.)?')15solution_pattern = re.compile(r'(?:\n|# )Solution(?:\s+(\d+)|\.|\n)')16 17 18def analyze(text: str) -> Tuple[List, int]:19    """20    Analyze the text and return the tags and problem number.21 22    Args:23        text (str): The markdown text to analyze.24 25    Returns:26        Tuple[List, int]: A tuple containing the tags and problem number.27    """28    tags = []29    tags.extend([(x, problem_tag) for x in problem_pattern.finditer(text)])30    problem_num = len(tags)31 32    tags.extend([(x, solution_tag) for x in solution_pattern.finditer(text)])33    tags.sort(key=lambda x: x[0].start())34    return tags, problem_num35 36 37def segment(text: str, tags):38    starts = []39    ends = []40 41    for i in range(len(tags)):42        starts.append(tags[i][0].end())43        if i + 1 < len(tags):44            ends.append(tags[i + 1][0].start())45        else:46            ends.append(len(text))47 48    return [text[start:end].strip() for start, end in zip(starts, ends)]49 50 51def join(tags, segments):52    problem, solution = '', ''53    problem_label, problem_match, solution_match = '', '', ''54    pairs = []55 56    has_solution = any([tag[1] == solution_tag for tag in tags])57 58    for tag, segment in zip(tags, segments):59        if tag[1] == problem_tag:60            problem = segment61            problem_match = tag[0].group(0)62            problem_label = tag[0].group(1)63 64            # If there is no solution, add an empty solution65            if not has_solution:66                pairs.append((problem, '', problem_label, problem_match, ''))67        else:68            solution = segment69            solution_match = tag[0].group(0)70            pairs.append((problem, solution, problem_label, problem_match, solution_match))71 72    return pairs73 74 75def write_pairs(output_file: Path, pairs):76    year = re.search(r'(\d{4})', output_file.stem).group(1)77 78    output_jsonl_text = ""79    for problem, solution, problem_label, problem_match, solution_match in pairs:80        output_jsonl_text += json.dumps(81            {82                'year': year,83                'tier': 'T1',84                'problem_label': problem_label,85                'problem_type': None,86                'exam': 'RMM',87                'problem': problem,88                'solution': solution,89                'metadata': {90                    'resource_path': output_file.relative_to(project_root).as_posix(),91                    'problem_match': problem_match,92                    'solution_match': solution_match93                }94            },95            ensure_ascii=False96        ) + '\n'97 98    output_file.write_text(output_jsonl_text, encoding="utf-8")99 100 101def main():102    compet_base_path = Path(__file__).resolve().parent.parent103    compet_md_path = compet_base_path / "md"104    seg_output_path = compet_base_path / "segmented"105 106    total_problem_count = 0107    total_solution_count = 0108 109    for apmo_md in tqdm(list(compet_md_path.glob('**/*.md')), desc='Segmenting'):110        output_file = seg_output_path / apmo_md.relative_to(compet_md_path).with_suffix('.jsonl')111        output_file.parent.mkdir(parents=True, exist_ok=True)112 113        text = '\n' + apmo_md.read_text(encoding="utf-8")114 115        tags, problem_num = analyze(text)116 117        if problem_num != 6 and problem_num != 3:118            logger.warning(f"{apmo_md} problem number is {problem_num}")119 120        if problem_num > 0:121            segments = segment(text, tags)122            pairs = join(tags, segments)        123            write_pairs(output_file, pairs)124 125            total_problem_count += problem_num126            total_solution_count += len(pairs)127        else:128            logger.warning(f"No problem found in {apmo_md}")129    130    logger.info(f"Total problem count: {total_problem_count}")131    logger.info(f"Total solution count: {total_solution_count}")132 133if __name__ == '__main__':134    main()135