fount-vlm-nell-02 0.3.11__py3-none-any.whl

mlx_vlm/evals/mmmu.py

@@ -0,0 +1,528 @@
+import argparse
+import csv
+import logging
+import os
+import random
+import re
+from json import dump
+
+from datasets import load_dataset
+from tqdm import tqdm
+
+from mlx_vlm import load
+from mlx_vlm.evals.utils import inference
+
+# All 30 MMMU subjects (confirmed from dataset)
+MMMU_SUBJECTS = [
+    "Accounting",
+    "Agriculture",
+    "Architecture_and_Engineering",
+    "Art",
+    "Art_Theory",
+    "Basic_Medical_Science",
+    "Biology",
+    "Chemistry",
+    "Clinical_Medicine",
+    "Computer_Science",
+    "Design",
+    "Diagnostics_and_Laboratory_Medicine",
+    "Economics",
+    "Electronics",
+    "Energy_and_Power",
+    "Finance",
+    "Geography",
+    "History",
+    "Literature",
+    "Manage",
+    "Marketing",
+    "Materials",
+    "Math",
+    "Mechanical_Engineering",
+    "Music",
+    "Pharmacy",
+    "Physics",
+    "Psychology",
+    "Public_Health",
+    "Sociology",
+]
+
+MMMU_PRO_SUBJECTS = [
+    "vision",
+    "standard (10 options)",
+    "standard (4 options)",
+]
+
+
+def normalize_number(s):
+    """Normalize numeric strings for comparison."""
+    try:
+        return float(str(s).strip().replace(",", ""))
+    except:
+        return str(s).strip()
+
+
+def MMMU_eval(data: list, eval_file: str):
+    """
+    Evaluate MMMU results by subject.
+    Handles both multiple choice (A-F) and open-ended questions.
+    """
+
+    # Track by subject
+    subject_scores = {}
+    subject_counters = {}
+
+    total_correct = 0
+    total_questions = 0
+
+    for line in data:
+        predict = str(line["prediction"])
+        answer = str(line["answer"])
+        subject = str(line.get("subject", "Unknown"))
+
+        # Initialize subject tracking if needed
+        if subject not in subject_scores:
+            subject_scores[subject] = 0
+            subject_counters[subject] = 0
+
+        # Count this question
+        subject_counters[subject] += 1
+        total_questions += 1
+
+        # Normalize for comparison
+        predict_lower = predict.lower().strip()
+        answer_lower = answer.lower().strip()
+
+        is_correct = False
+
+        # Check if this is a multiple choice question (answer is A-F or I)
+        if answer in ["A", "B", "C", "D", "E", "F", "I"]:
+            # Multiple choice extraction with prioritized patterns
+            patterns = [
+                (r"option\s+([a-f])\b", 10),  # High priority
+                (r"answer\s+is:?\s+([a-f])\b", 10),
+                (r"choice\s+is:?\s+([a-f])\b", 10),
+                (r"correct\s+answer\s+is:?\s+([a-f])\b", 10),
+                (r"correct\s+option\s+is:?\s+\(?([a-f])\)?", 10),
+                (r"\(([a-f])\)", 8),  # Medium priority
+                (r"^([a-f])[.:\)]\s", 8),
+                (r"\b([a-f])\b", 5),  # Low priority - isolated letters
+            ]
+
+            best_match = None
+            best_priority = -1
+
+            # Try each pattern, keeping the highest priority match
+            for pattern, priority in patterns:
+                matches = re.findall(pattern, predict_lower, re.IGNORECASE)
+                if matches and priority > best_priority:
+                    best_match = matches[0].lower()
+                    best_priority = priority
+                    # Stop early if we found a high-confidence pattern
+                    if priority >= 10:
+                        break
+
+            # Check if match is correct
+            if best_match and best_match == answer_lower:
+                is_correct = True
+            # Fallback: check first character
+            elif (
+                not best_match
+                and len(predict_lower) > 0
+                and predict_lower[0] in "abcdef"
+            ):
+                if predict_lower[0] == answer_lower:
+                    is_correct = True
+
+        else:
+            # Open-ended question - check if answer appears in prediction
+            # Exact substring match (case-insensitive)
+            if answer_lower in predict_lower:
+                is_correct = True
+            # For numeric answers, try numeric comparison
+            elif answer.replace(".", "").replace("-", "").replace(",", "").isdigit():
+                numbers = re.findall(r"-?\d+\.?\d*", predict)
+                answer_num = normalize_number(answer)
+                for num_str in numbers:
+                    try:
+                        if abs(normalize_number(num_str) - answer_num) < 0.01:
+                            is_correct = True
+                            break
+                    except:
+                        pass
+            # Word-level match for text answers
+            else:
+                answer_words = set(answer_lower.split())
+                predict_words = set(predict_lower.split())
+                if answer_words and answer_words.issubset(predict_words):
+                    is_correct = True
+
+        if is_correct:
+            total_correct += 1
+            subject_scores[subject] += 1
+            line["score"] = 1
+        else:
+            line["score"] = 0
+
+    # Calculate final scores
+    results = {}
+    results["overall_accuracy"] = (
+        float(total_correct) / float(total_questions) if total_questions > 0 else 0.0
+    )
+    results["total_correct"] = total_correct
+    results["total_questions"] = total_questions
+
+    # Calculate subject scores
+    for subject in sorted(subject_scores.keys()):
+        if subject_counters[subject] > 0:
+            results[f"subject_{subject}_accuracy"] = float(
+                subject_scores[subject]
+            ) / float(subject_counters[subject])
+            results[f"subject_{subject}_correct"] = subject_scores[subject]
+            results[f"subject_{subject}_total"] = subject_counters[subject]
+
+    # Print scores
+    print("\nMMMU Evaluation Results:")
+    print("=" * 80)
+    print(f"Model: {eval_file.split('/')[-1].split('_MMMU_')[0]}")
+    print(f"Total Questions: {total_questions}")
+    print(f"Total Correct: {total_correct}")
+    print(
+        f"Overall Accuracy: {results['overall_accuracy']:.4f} ({total_correct}/{total_questions})"
+    )
+    print("=" * 80)
+    print("Subject Breakdown:")
+    for subject in sorted(subject_scores.keys()):
+        acc = results.get(f"subject_{subject}_accuracy", 0.0)
+        correct = results.get(f"subject_{subject}_correct", 0)
+        total = results.get(f"subject_{subject}_total", 0)
+        print(f"  {subject}: {acc:.4f} ({correct}/{total})")
+    print("=" * 80)
+
+    # Save results
+    score_pth = eval_file.replace(".csv", "_score.json")
+    with open(score_pth, "w") as f:
+        dump(results, f, indent=2)
+
+    with open(eval_file, "w", newline="", encoding="utf-8") as f:
+        if data:
+            writer = csv.DictWriter(f, fieldnames=data[0].keys())
+            writer.writeheader()
+            writer.writerows(data)
+
+    logging.info(
+        f"MMMU_eval successfully finished evaluating {eval_file}, results saved in {score_pth}"
+    )
+
+
+def process_question(example):
+    """
+    Process MMMU question to format it properly.
+    MMMU questions may have options and images.
+    """
+    question = example.get("question", "")
+
+    # Add options if they exist
+    options = example.get("options", None)
+    options = re.sub(r'[\[\]"\']', "", options).split(", ") if options else None
+
+    if options and isinstance(options, list):
+        question += "\n\nOptions:"
+        for i, option in enumerate(options):
+            letter = chr(65 + i)  # A, B, C, D, ...
+            question += f"\n{letter}. {option}"
+
+    # Remove <image n> tags from the question
+    question = re.sub(r"<image \d+>", "", question).strip()
+
+    return question
+
+
+def get_images(example):
+    """
+    Extract images from MMMU example.
+    MMMU can have multiple images per question.
+    """
+    images = []
+
+    # MMMU dataset may have image_1, image_2, etc.
+    if "image" in example and example["image"] is not None:
+        try:
+            img = example["image"].convert("RGB")
+            images.append(img)
+        except Exception as e:
+            print(f"Warning: Could not process image - {e}")
+    else:
+        for i in range(0, 8):  # Check up to 7 images
+            img_key = f"image_{i}"
+            if img_key in example and example[img_key] is not None:
+                try:
+                    img = example[img_key].convert("RGB")
+                    images.append(img)
+                except Exception as e:
+                    print(f"Warning: Could not process image for key {img_key} - {e}")
+                    continue
+    return images
+
+
+def list_subjects():
+    """Print all available MMMU subjects."""
+    print("\n" + "=" * 80)
+    print("MMMU Pro Subjects (3 total)")
+    print("=" * 80)
+    for i, subject in enumerate(MMMU_PRO_SUBJECTS, 1):
+        print(f"{i:2d}. {subject}")
+    print("\n" + "=" * 80)
+    print("MMMU Available Subjects (30 total)")
+    print("=" * 80)
+    for i, subject in enumerate(MMMU_SUBJECTS, 1):
+        print(f"{i:2d}. {subject}")
+    print("=" * 80 + "\n")
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description="MMMU Evaluation - Massive Multi-discipline Multimodal Understanding",
+        epilog="Use --subset to evaluate a specific subject, or omit to evaluate all 30 subjects.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="mlx-community/Qwen2-VL-2B-Instruct-bf16",
+        help="Model path",
+    )
+    parser.add_argument("--adapter_path", type=str, default=None, help="Adapter path")
+    parser.add_argument("--dataset", type=str, default="MMMU/MMMU", help="Dataset path")
+    parser.add_argument(
+        "--split", type=str, default="validation", help="Split to use for evaluation"
+    )
+    parser.add_argument(
+        "--subset",
+        type=str,
+        default=None,
+        help=f"Subset to use - one of 30 subjects: {', '.join(MMMU_SUBJECTS[:5])}... (see SUBJECTS.md for full list)",
+    )
+    parser.add_argument(
+        "--streaming", action="store_true", help="Use streaming dataset loading"
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=3000,
+        help="Maximum number of tokens to generate",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.0,
+        help="Temperature for sampling (0.0 for greedy)",
+    )
+    parser.add_argument(
+        "--top-p",
+        type=float,
+        default=0.9,
+        help="Top-p sampling parameter",
+    )
+    parser.add_argument(
+        "--repetition-penalty",
+        type=float,
+        default=1.0,
+        help="Repetition penalty parameter",
+    )
+    parser.add_argument(
+        "--resize-shape",
+        type=int,
+        nargs=2,
+        default=None,
+        help="Resize shape for the image",
+    )
+    parser.add_argument("--verbose", action="store_true", help="Verbose output")
+    parser.add_argument(
+        "--max-samples",
+        type=int,
+        default=None,
+        help="Maximum number of samples to evaluate (for testing)",
+    )
+    parser.add_argument(
+        "--list-subjects",
+        action="store_true",
+        help="List all 30 available subjects and exit",
+    )
+    parser.add_argument(
+        "--prediction-file",
+        type=str,
+        default=None,
+        help="Path to the prediction file",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="results/mmmu",
+        help="Directory to save evaluation results",
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+
+    random.seed(args.seed)
+
+    # Setup logging
+    logging.basicConfig(
+        level=logging.INFO if args.verbose else logging.WARNING,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+    )
+
+    if "pro" in args.dataset.lower():
+        subjects = MMMU_PRO_SUBJECTS
+    else:
+        subjects = MMMU_SUBJECTS
+
+    if args.prediction_file:
+        # Load predictions from file
+        logging.info(f"\033[32mLoading predictions from {args.prediction_file}\033[0m")
+        results = []
+        with open(args.prediction_file, "r", encoding="utf-8") as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                results.append(row)
+
+        # Evaluate loaded predictions
+        MMMU_eval(results, args.prediction_file)
+        logging.info(f"\033[32mEvaluation complete\033[0m")
+        return
+
+    # Handle --list-subjects flag
+    if args.list_subjects:
+        list_subjects()
+        return
+
+    logging.info("\033[32mStarting MMMU Evaluation\033[0m")
+
+    # Validate subset if provided
+    if args.subset and args.subset not in subjects:
+        logging.error(f"\033[31mError: Invalid subset '{args.subset}'\033[0m")
+        logging.error(f"\033[31mValid subjects are: {', '.join(subjects)}\033[0m")
+        logging.error(f"\033[31mSee SUBJECTS.md for more details\033[0m")
+        return
+
+    logging.info(f"\033[32mLoading dataset from {args.dataset}\033[0m")
+
+    # Load dataset
+
+    if args.subset:
+        logging.info(f"\033[32mUsing subset: {args.subset}\033[0m")
+        datasets = {
+            args.subset: load_dataset(
+                args.dataset, args.subset, split=args.split, streaming=args.streaming
+            )
+        }
+        subset_name = args.subset
+    else:
+        logging.info(f"\033[32mEvaluating all 30 subjects\033[0m")
+        datasets = {}
+
+        for subject in subjects:
+            try:
+                datasets[subject] = load_dataset(
+                    args.dataset,
+                    name=subject,
+                    split=args.split,
+                    streaming=args.streaming,
+                )
+            except Exception as e:
+                logging.error(
+                    f"\033[31mError loading dataset for {subject}: {e}\033[0m"
+                )
+                continue
+
+        subset_name = "all"
+
+    # Limit samples if specified
+    if args.max_samples:
+        datasets = {
+            k: v.select(range(min(args.max_samples, len(v))))
+            for k, v in datasets.items()
+        }
+        logging.info(f"\033[33mLimited to {len(datasets)} samples for testing\033[0m")
+
+    logging.info(f"\033[32mDataset subset size: {len(datasets.keys())}\033[0m")
+    logging.info(f"\033[32mLoading model from {args.model}\033[0m")
+
+    model, processor = load(
+        args.model, adapter_path=args.adapter_path, trust_remote_code=True
+    )
+    config = model.config
+    logging.info(f"\033[32mConfig: {config}\033[0m")
+
+    # Create results directory
+    model_name = args.model.split("/")[-1]
+    result_file = f"{args.output_dir}/{model_name}_MMMU_{subset_name}_{args.split}_predictions.csv"
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    results = []
+    for subject, dataset in tqdm(datasets.items(), desc="Processing subjects"):
+        for idx, example in enumerate(tqdm(dataset, desc=f"Processing {subject}")):
+            question = process_question(example)
+
+            images = get_images(example)
+            try:
+                # Get prediction
+                prediction = inference(
+                    model,
+                    processor,
+                    question,
+                    images,
+                    args.max_tokens,
+                    args.temperature,
+                    args.resize_shape,
+                    args.verbose,
+                )
+            except Exception as e:
+                print(f"Error during inference:", question, images, "error message:", e)
+                prediction = ""
+
+            # Store result
+            result = {
+                "id": example.get("id", idx),
+                "question": question,
+                "answer": example.get("answer", ""),
+                "subfield": example.get("subfield", "Unknown"),
+                "topic_difficulty": example.get("topic_difficulty", "Unknown"),
+                "question_type": example.get("question_type", "Unknown"),
+                "prediction": prediction,
+                "subject": example.get("subject", None) or subject,
+            }
+            results.append(result)
+
+            # Show progress
+            if (idx + 1) % 10 == 0 or idx < 5:
+                logging.info(
+                    f"Sample {idx + 1}: Answer={result['answer']}, Prediction={prediction[:50]}..."
+                )
+
+    # Print first few results
+    print("\nFirst 5 results:")
+    for i, result in enumerate(results[:5]):
+        print(
+            f"{i+1}. Question: {result['question'][:50]}... | Answer: {result['answer']} | Prediction: {result['prediction'][:50]}..."
+        )
+
+    # Save results to CSV
+    with open(result_file, "w", newline="", encoding="utf-8") as f:
+        if results:
+            writer = csv.DictWriter(f, fieldnames=results[0].keys())
+            writer.writeheader()
+            writer.writerows(results)
+
+    logging.info(f"\033[32mSaved results to {result_file}\033[0m")
+
+    # Evaluate results
+    MMMU_eval(results, result_file)
+
+    logging.info(f"\033[32mEvaluation complete\033[0m")
+
+
+if __name__ == "__main__":
+    main()