beichen0426/olympiads-ref

1import re2import json3 4from tqdm import tqdm5from loguru import logger6 7from pathlib import Path8from typing import Tuple, List9from dataclasses import dataclass10 11 12project_root = Path(__file__).parent.parent.parent13 14 15@dataclass16class Problem:17    match: re.Match18 19 20@dataclass21class Solution:22    match: re.Match23 24 25def clean_text(text: str) -> str:26    text = text.replace(27        'For a discussion, see\nW. Morris and V. Soltan. The Erdős-Szekeres Problem on Points in Convex Postion-A Survey, Bulletin of the American Math Monthly. 37 (2000), 437-458.\n\nThis article is available at\nhttp://www.ams.org/bull/2000-37-04/S0273-0979-00-00877-6/home.html.\nIf $N(7)=33$, the highest sure score on this problem would be $32-6=26$. It is not known whether there exist arbitrarily large sets of points that will fool the graders.\n\n## The unexamined life is not worth living.',28        ''29    )30    text = text.replace(31        '- Bishop: This piece can move any number of squares diagonally if there are no other pieces along its path.\n- Rook: This piece can move any number of squares either vertically or horizontally if there are no other pieces along its path\n- Knight: This piece can move either two squares along a row and one square along a column or two squares along a column and one square along a row.\n- King: This piece can move to any open adjacent square (including diagonally).',32        ''33    )34    return text35 36 37def find_problem_with_solution(38    text: str,39    problem_parttern: re.Pattern,40    solution_pattern: re.Pattern41) -> int:42    """43    Find the problem with solution start position in the text.44    Args:45        text (str): The text to search.46    Returns:47        int: The start position of the problem with solution.48    """49    matchs = list(problem_parttern.finditer(text))50 51    for index, match in enumerate(matchs):52        section_end_position = matchs[index + 1].start() if index + 1 < len(matchs) else len(text)53        if solution_pattern.search(text[match.start():section_end_position]):54            return match.start()55 56    return 057 58 59def analyze(text: str) -> Tuple[List[Problem | Solution], int]:60    """61    Analyze the text and return the tags and problem number.62    Args:63        text (str): The markdown text to analyze.64    Returns:65        Tuple[List[Problem | Solution], int]: A tuple containing the tags and problem number.66    """67    problem_pattern = re.compile(r'(?:\n|\n\#+\s+)(?:(\d{1,2})\.\s+(?:problem\:\s*|\$?\[.+?\]\$?)?|problem\s+?(\w+)\s+\[\d+(?:\spoints)?\]|\$([H|M|T]_\{\d+\})\$\.)', re.IGNORECASE)68    solution_pattern = re.compile(r'(?:\n|\n\#+\s+)(?:answer\:|solution(?:\s+\d+)?(?:\:|\.)|Proposed by:.*?\n)\s*', re.IGNORECASE)69 70    start_position = find_problem_with_solution(text, problem_pattern, solution_pattern)71 72    tags: List[Problem | Solution] = []73    tags.extend([Problem(x) for x in problem_pattern.finditer(text, start_position)])74    problem_num = len(tags)75 76    tags.extend([Solution(x) for x in solution_pattern.finditer(text, start_position)])77    tags.sort(key=lambda x: x.match.start())78    return tags, problem_num79 80 81def segment(text: str, tags: List[Problem | Solution]) -> List[str]:82    starts = []83    ends = []84 85    for i in range(len(tags)):86        starts.append(tags[i].match.end())87        if i + 1 < len(tags):88            ends.append(tags[i + 1].match.start())89        else:90            ends.append(len(text))91 92    return [text[start:end].strip() for start, end in zip(starts, ends)]93 94 95def join(tags: List[Problem | Solution], segments: List[str]) -> List[Tuple[str, str, str, str, str]]:96    problem, solution = '', ''97    problem_label, problem_match, solution_match = '', '', ''98    pairs = []99 100    for tag, segment in zip(tags, segments):101        if isinstance(tag, Problem):102            problem = segment103            problem_match = tag.match.group(0)104            problem_label = tag.match.group(1) or tag.match.group(2) or tag.match.group(3)105        elif problem.strip() != "":106            solution = segment107            solution_match = tag.match.group(0)108 109            if solution.strip() == "":110                continue111 112            pairs.append((problem, solution, problem_label, problem_match, solution_match))113 114    return pairs115 116 117def write_pairs(output_file: Path, pairs):118    year = re.search(r'(\d{4})', output_file.stem).group(1)119    problem_type_mapping = {120        "-alg-": "Algebra",121        "-comb-": "Combinatorics",122        "-geo-": "Geometry",123    }124 125    problem_type = None126    for _k, _v in problem_type_mapping.items():127        if _k in output_file.stem:128            problem_type = _v129            break130 131    output_jsonl_text = ""132    for problem, solution, problem_label, problem_match, solution_match in pairs:133        output_jsonl_text += json.dumps(134            {135                'year': year,136                'tier': "T4",137                'problem_label': problem_label,138                'problem_type': problem_type,139                "exam": "HMMT", 140                'problem': problem,141                'solution': solution,142                'metadata': {143                    'resource_path': output_file.relative_to(project_root).as_posix(),144                    'problem_match': problem_match,145                    'solution_match': solution_match146                }147            },148            ensure_ascii=False149        ) + '\n'150 151    output_file.write_text(output_jsonl_text, encoding="utf-8")152 153 154def main():155    compet_base_path = Path(__file__).resolve().parent.parent156    compet_md_path = compet_base_path / "md"157    seg_output_path = compet_base_path / "segmented"158 159    total_problem_count = 0160    total_solution_count = 0161 162    for hmmt_md in tqdm(list(compet_md_path.glob('**/*.md')), desc='Segmenting'):163        output_file = seg_output_path / hmmt_md.relative_to(compet_md_path).with_suffix('.jsonl')164        output_file.parent.mkdir(parents=True, exist_ok=True)165 166        text = '\n' + clean_text(hmmt_md.read_text(encoding="utf-8"))167 168        tags, problem_num = analyze(text)169 170        segments = segment(text, tags)171        pairs = join(tags, segments)     172        if pairs and problem_num > 0:173            write_pairs(output_file, pairs)174 175            total_problem_count += problem_num176            total_solution_count += len(pairs)177        else:178            logger.warning(f"No problem found in {hmmt_md}")179        180    logger.info(f"Total problem count: {total_problem_count}")181    logger.info(f"Total solution count: {total_solution_count}")182 183 184if __name__ == '__main__':185    main()186