python/rag/prepare_qa.py

"""
Download and prepare QA training data for H4 RAG.

Uses a simple extractive QA format:
- Input: [context] | [question] |
- Target: [answer]

Data sources (in order of preference):
1. SQuAD-style QA pairs generated from the sample documents
2. Downloaded SQuAD 2.0 dev set (small, freely available)

For CPU training with 2-minute budget, we need small data that
trains fast. The sample doc QA pairs are ideal for proving the
pipeline works; SQuAD provides real benchmark numbers.
"""

import json
import os
import sys
import random
from typing import List, Tuple, Dict

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))


def generate_sample_qa() -> List[Dict]:
    """
    Generate QA pairs from the sample documents.
    These are hand-crafted to match the sample_docs content.
    The model's job: learn to extract the answer from the context.
    """
    qa_pairs = [
        # golden_ratio.txt
        {"context": "The golden ratio, often denoted by the Greek letter phi, is a special number approximately equal to 1.618.",
         "question": "What is the golden ratio approximately equal to?",
         "answer": "1.618"},
        {"context": "Two quantities are in the golden ratio if their ratio is the same as the ratio of their sum to the larger of the two quantities.",
         "question": "When are two quantities in the golden ratio?",
         "answer": "if their ratio is the same as the ratio of their sum to the larger"},
        {"context": "The golden ratio is closely related to the Fibonacci sequence. As Fibonacci numbers increase, the ratio of consecutive Fibonacci numbers approaches the golden ratio.",
         "question": "How is the golden ratio related to Fibonacci numbers?",
         "answer": "the ratio of consecutive Fibonacci numbers approaches the golden ratio"},
        {"context": "The golden ratio appears in the geometry of pentagons and in the arrangement of leaves and petals in many plants.",
         "question": "Where does the golden ratio appear in nature?",
         "answer": "in the arrangement of leaves and petals in many plants"},

        # polytopes.txt
        {"context": "The 600-cell is a regular 4-polytope with 120 vertices, 720 edges, 1200 triangular faces, and 600 tetrahedral cells.",
         "question": "How many vertices does the 600-cell have?",
         "answer": "120"},
        {"context": "The 600-cell has the H4 symmetry group, which contains 14400 elements. This is the largest finite reflection group in four dimensions.",
         "question": "How many elements does the H4 symmetry group contain?",
         "answer": "14400"},
        {"context": "The 600-cell is dual to the 120-cell, which has 600 vertices.",
         "question": "What is the 600-cell dual to?",
         "answer": "the 120-cell"},
        {"context": "A polytope is a geometric object with flat sides in any number of dimensions.",
         "question": "What is a polytope?",
         "answer": "a geometric object with flat sides in any number of dimensions"},

        # e8_lattice.txt
        {"context": "The E8 lattice is the densest sphere packing in eight dimensions. This was proven by Maryna Viazovska in 2016.",
         "question": "Who proved E8 is the densest sphere packing?",
         "answer": "Maryna Viazovska"},
        {"context": "The E8 lattice has a kissing number of 240, meaning each sphere touches exactly 240 others.",
         "question": "What is the kissing number of E8?",
         "answer": "240"},
        {"context": "The Coxeter element of E8 has eigenvalues that include cosine of pi over five, which equals phi over two.",
         "question": "What eigenvalue connects E8 to the golden ratio?",
         "answer": "cosine of pi over five, which equals phi over two"},
        {"context": "When the 240 roots of E8 are projected along these eigenspaces, they map to the vertices of H4 polytopes.",
         "question": "What happens when E8 roots are projected along the eigenspaces?",
         "answer": "they map to the vertices of H4 polytopes"},
    ]

    return qa_pairs


def prepare_training_data(
    qa_pairs: List[Dict],
    val_fraction: float = 0.2,
) -> Tuple[List[Dict], List[Dict]]:
    """Split QA pairs into train and validation sets."""
    random.seed(42)
    pairs = list(qa_pairs)
    random.shuffle(pairs)
    n_val = max(1, int(len(pairs) * val_fraction))
    return pairs[n_val:], pairs[:n_val]


def format_qa_for_training(qa_pair: Dict, sep: str = " | ") -> Tuple[str, str]:
    """
    Format a QA pair for character-level training.

    Input: [context] | [question] |
    Target: [answer]

    The model learns to generate the answer given context + question.
    """
    input_text = qa_pair['context'] + sep + qa_pair['question'] + sep
    target_text = qa_pair['answer']
    return input_text, target_text


def download_squad_dev():
    """
    Download SQuAD 2.0 dev set for real benchmark evaluation.
    Returns list of QA dicts with context/question/answer.
    """
    import urllib.request

    cache_dir = os.path.join(os.path.dirname(__file__), '..', '..', 'data')
    os.makedirs(cache_dir, exist_ok=True)
    cache_path = os.path.join(cache_dir, 'squad_dev.json')

    if not os.path.exists(cache_path):
        url = "https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v2.0.json"
        print(f"Downloading SQuAD 2.0 dev set...")
        try:
            urllib.request.urlretrieve(url, cache_path)
            print(f"Saved to {cache_path}")
        except Exception as e:
            print(f"Download failed: {e}")
            return []

    with open(cache_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    qa_pairs = []
    for article in data['data']:
        for paragraph in article['paragraphs']:
            context = paragraph['context']
            for qa in paragraph['qas']:
                if qa.get('is_impossible', False):
                    continue
                if qa['answers']:
                    answer = qa['answers'][0]['text']
                    qa_pairs.append({
                        'context': context[:500],  # truncate long contexts
                        'question': qa['question'],
                        'answer': answer,
                    })

    return qa_pairs


if __name__ == '__main__':
    print("Generating sample QA pairs...")
    pairs = generate_sample_qa()
    train, val = prepare_training_data(pairs)
    print(f"Sample QA: {len(train)} train, {len(val)} val")

    for p in pairs[:3]:
        inp, tgt = format_qa_for_training(p)
        print(f"\nInput:  {inp[:80]}...")
        print(f"Target: {tgt}")

    print("\nAttempting SQuAD download...")
    squad = download_squad_dev()
    if squad:
        print(f"SQuAD 2.0 dev: {len(squad)} answerable questions")
    else:
        print("SQuAD not available (offline?). Using sample QA only.")