scripts/prepare_dataset.py

"""
ByteAstra — Dataset preparation and validation script.
Checks generated QA pairs against the official BAMS syllabus config,
ensuring strict syllabus alignment and filtering out invalid/noisy examples.
"""
from __future__ import annotations
import argparse
import json
import logging
from pathlib import Path
import sys

# Import syllabus configuration
sys.path.append(str(Path(__file__).parent.parent))
from app.domains.bams_syllabus_config import ALL_SYLLABUS_TOPICS, OUT_OF_SYLLABUS_TOPICS

logging.basicConfig(level=logging.INFO, format="%(levelname)s | %(message)s")
logger = logging.getLogger(__name__)

def is_syllabus_aligned(question: str, answer: str) -> bool:
    """
    Checks if a QA pair matches the BAMS syllabus.
    - Must NOT contain explicitly out-of-syllabus topics.
    - Must have some overlap with BAMS curriculum topics.
    """
    combined = (question + " " + answer).lower()
    
    # Strict out-of-syllabus check
    for forbidden in OUT_OF_SYLLABUS_TOPICS:
        if forbidden in combined:
            return False
            
    # Overlap with syllabus check
    matched = False
    for topic in ALL_SYLLABUS_TOPICS:
        if topic.lower() in combined:
            matched = True
            break
            
    return matched

def main():
    parser = argparse.ArgumentParser(description="Validate and filter BAMS QA dataset")
    parser.add_argument("--input", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()

    if not args.input.exists():
        logger.error("Input file %s does not exist.", args.input)
        sys.exit(1)

    logger.info("Loading dataset from %s", args.input)
    
    valid_count = 0
    total_count = 0
    
    with open(args.input, "r", encoding="utf-8") as infile, open(args.output, "w", encoding="utf-8") as outfile:
        for line in infile:
            if not line.strip():
                continue
            total_count += 1
            try:
                data = json.loads(line)
                # ShareGPT format: conversations list with role/content
                conversations = data.get("conversations", [])
                if len(conversations) < 3:
                    continue
                
                question = conversations[1]["content"]
                answer = conversations[2]["content"]
                
                if is_syllabus_aligned(question, answer):
                    outfile.write(json.dumps(data, ensure_ascii=False) + "\n")
                    valid_count += 1
            except Exception as e:
                logger.error("Error processing line: %s", e)
                
    logger.info("✓ Verification complete.")
    logger.info("Processed %d total items.", total_count)
    logger.info("Retained %d syllabus-aligned items (%d filtered out).", valid_count, total_count - valid_count)
    logger.info("Output written to: %s", args.output.resolve())

if __name__ == "__main__":
    main()