app.py

import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
from datasets import load_dataset
import numpy as np
from tqdm import tqdm
import json
from huggingface_hub import hf_hub_download, login
import os

# Retrieve the token from the environment variable
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)
    print("Successfully logged in to Hugging Face Hub!")
else:
    print("HF_TOKEN not found in environment variables. Cannot authenticate.")

def load_models(target_model, target_adapter):
    """Load base model and fine-tuned adapter"""
    base_model_name = target_model
    adapter_name = target_adapter
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
    tokenizer.pad_token = tokenizer.eos_token
    
    # Load base model
    base_model = AutoModelForCausalLM.from_pretrained(
        base_model_name,
        torch_dtype=torch.bfloat16,
        device_map="auto"
    )
    
    # Load fine-tuned model
    ft_model = PeftModel.from_pretrained(base_model, adapter_name)
    
    return tokenizer, base_model, ft_model

def evaluate_mmlu_sample(model, tokenizer, question, choices, answer_idx):
    """Evaluate a single MMLU sample"""
    
    # Format prompt with choices
    prompt = f"Question: {question}\n"
    for i, choice in enumerate(choices):
        prompt += f"{chr(65+i)}: {choice}\n"
    prompt += "Answer:"
    
    # Tokenize
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
    
    with torch.no_grad():
        # Get logits for next token
        outputs = model(**inputs)
        logits = outputs.logits[0, -1, :]  # Last token logits
        
        # Get probabilities for A, B, C, D tokens
        choice_tokens = [tokenizer.encode(chr(65+i), add_special_tokens=False)[0] for i in range(4)]
        choice_logits = [logits[token].item() for token in choice_tokens]
        
        # Predict choice with highest logit
        predicted_idx = np.argmax(choice_logits)
        
    return predicted_idx == answer_idx

def run_mmlu_evaluation(model, tokenizer, model_name, count):
    """Run MMLU evaluation on a balanced subset across all subjects"""
    
    # Load MMLU dataset
    dataset = load_dataset("cais/mmlu", "all", split="test")
    
    # Get all unique subjects
    all_subjects = list(set(dataset['subject']))
    num_subjects = len(all_subjects)
    
    # Calculate samples per subject (rounded down)
    samples_per_subject = int(np.floor(count / num_subjects))
    
    print(f"Found {num_subjects} subjects in MMLU dataset")
    print(f"Sampling {samples_per_subject} questions per subject")
    print(f"Total samples: {samples_per_subject * num_subjects}")
    
    # Create balanced subset
    balanced_indices = []
    
    for subject in all_subjects:
        # Get all indices for this subject
        subject_indices = [i for i, subj in enumerate(dataset['subject']) if subj == subject]
        
        # Sample up to samples_per_subject questions from this subject
        if len(subject_indices) >= samples_per_subject:
            # If we have enough samples, randomly select samples_per_subject
            selected_indices = np.random.choice(subject_indices, samples_per_subject, replace=False)
        else:
            # If we don't have enough samples, take all available
            selected_indices = subject_indices
            print(f"Warning: Subject '{subject}' only has {len(subject_indices)} samples (requested {samples_per_subject})")
        
        balanced_indices.extend(selected_indices.tolist())
    
    # Create balanced dataset
    dataset = dataset.select(balanced_indices)
    
    # Shuffle the final dataset
    shuffled_indices = np.random.permutation(len(dataset))
    dataset = dataset.select(shuffled_indices)
    
    correct = 0
    total = 0
    
    subject_scores = {}
    
    print(f"Evaluating {model_name} on {len(dataset)} balanced MMLU samples...")
    
    for sample in tqdm(dataset):
        subject = sample['subject']
        question = sample['question']
        choices = sample['choices']
        answer = sample['answer']  # Index of correct answer
        
        # Evaluate sample
        is_correct = evaluate_mmlu_sample(model, tokenizer, question, choices, answer)
        
        # Track overall accuracy
        if is_correct:
            correct += 1
        total += 1
        
        # Track per-subject accuracy
        if subject not in subject_scores:
            subject_scores[subject] = {'correct': 0, 'total': 0}
        
        subject_scores[subject]['total'] += 1
        if is_correct:
            subject_scores[subject]['correct'] += 1
    
    # Calculate final scores
    overall_accuracy = (correct / total) * 100
    
    subject_accuracies = {}
    for subject, scores in subject_scores.items():
        subject_accuracies[subject] = (scores['correct'] / scores['total']) * 100
    
    return overall_accuracy, subject_accuracies

def run_comparison(target_model, target_adapter, count):

    progress = gr.Progress()
    
    # count of questions per MMLU subject from form needs to be multiplied by number of subjects 
    count = int(count) * 57
    print(f"Running evaluation on {count} MMLU questions...")

    progress(0.01, desc="Loading models...")
    
    print("Loading models...")
    tokenizer, base_model, ft_model = load_models(target_model, target_adapter)

    progress(0.1, desc="Evaluating base model...")
    
    # Evaluate base model
    print("Evaluating base model...")
    base_accuracy, base_subjects = run_mmlu_evaluation(base_model, tokenizer, f"{target_model}", count)

    progress(0.5, desc="Evaluating fine-tuned adapter...")
    
    # Evaluate fine-tuned model
    print("Evaluating fine-tuned adapter...")
    ft_accuracy, ft_subjects = run_mmlu_evaluation(ft_model, tokenizer, f"{target_adapter}", count)

    progress(1, desc="Evaluation complete.")

    data = {
        "base_accuracy": base_accuracy,
        "base_subjects": base_subjects,
        "adapter_accuracy": ft_accuracy,
        "adapter_subjects": ft_subjects
    }
    
    output = ''
  
    # Print overall results
    
    output += f"## flexMMLU EVALUATION RESULTS - OVERALL PERFORMANCE\n"
    output += f"### Base Model Accuracy: {base_accuracy:.2f}%\n"
    output += f"### Adapter Accuracy: {ft_accuracy:.2f}%\n"
    output += f"### Performance Change: {ft_accuracy - base_accuracy:+.2f}%\n"
    output += f"## SUBJECT-BY-SUBJECT PERFORMANCE COMPARISON\n"
    output += f"| {'Subject':<35} | {'Base %':<10} | {'Adapter %':<15} | {'Change':<10} | {'Status'} |\n"
    output += f"|---|---|---|---|---|\n"
    
    
    # Get all subjects (union of both dictionaries in case there are differences)
    all_subjects = sorted(set(base_subjects.keys()) | set(ft_subjects.keys()))
    
    # Track statistics
    improvements = 0
    degradations = 0
    no_changes = 0
    total_change = 0
  
    subject_changes = []
    
    # Print each subject
    for subject in all_subjects:
      # Get scores (default to 0 if subject missing from either model)
      base_score = base_subjects.get(subject, 0.0)
      ft_score = ft_subjects.get(subject, 0.0)
      change = ft_score - base_score
      
      # Determine status
      if change > 0:
          status = "📈 IMPROVED"
          improvements += 1
      elif change < 0:
          status = "📉 DECLINED"  
          degradations += 1
      else:
          status = "➡️  NO CHANGE"
          no_changes += 1
      
      # Format subject name (replace underscores, capitalize)
      formatted_subject = subject.replace('_', ' ').title()
      
      output += f"| {formatted_subject:<35} | {base_score:<10.2f} | {ft_score:<15.2f} | {change:<+10.2f} | {status} |\n"
    
      subject_changes.append((formatted_subject, change))
      total_change += change
  
    # Print summary statistics
    output += f"## SUMMARY STATISTICS\n"
    output += f"#### Total Subjects Evaluated: {len(all_subjects)}\n"
    output += f"### Subjects Improved: {improvements} ({improvements/len(all_subjects)*100:.1f}%)\n"
    output += f"### Subjects Declined: {degradations} ({degradations/len(all_subjects)*100:.1f}%)\n"
    output += f"### Subjects Unchanged: {no_changes} ({no_changes/len(all_subjects)*100:.1f}%)\n"
    output += f"### Average Change per Subject: {total_change/len(all_subjects):+.2f}%\n"
    
    # Show top improvements and declines
    subject_changes.sort(key=lambda x: x[1], reverse=True)
    
    output += f"### 🏆 TOP 10 MOST IMPROVED SUBJECTS\n"
    for i, (subject, change) in enumerate(subject_changes[:10]):
      output += f"{i+1:2d}. {subject:<30} {change:+.2f}%\n"
    
    output += f"### ⚠️ TOP 10 MOST DECLINED SUBJECTS\n"
    for i, (subject, change) in enumerate(subject_changes[-10:]):
      output += f"{i+1:2d}. {subject:<30} {change:+.2f}%\n"
    
    # Categories analysis (if we can group subjects)
    output += f"### 📚 ANALYSIS BY ACADEMIC DOMAIN\n"
    
    # Group subjects by broad categories
    stem_keywords = ['math', 'physics', 'chemistry', 'biology', 'computer', 'engineering', 'statistics']
    humanities_keywords = ['history', 'philosophy', 'literature', 'art', 'religion', 'culture']
    social_keywords = ['psychology', 'sociology', 'economics', 'political', 'law', 'business']
  
    stem_changes = []
    humanities_changes = []
    social_changes = []
    other_changes = []
    
    for subject, change in subject_changes:
      subject_lower = subject.lower()
      if any(keyword in subject_lower for keyword in stem_keywords):
          stem_changes.append(change)
      elif any(keyword in subject_lower for keyword in humanities_keywords):
          humanities_changes.append(change)
      elif any(keyword in subject_lower for keyword in social_keywords):
          social_changes.append(change)
      else:
          other_changes.append(change)
    # Print category averages
    if stem_changes:
      output += f"🔬 STEM Subjects: {np.mean(stem_changes):+.2f}% avg change ({len(stem_changes)} subjects)\n\n"
    if humanities_changes:
      output += f"📖 Humanities: {np.mean(humanities_changes):+.2f}% avg change ({len(humanities_changes)} subjects)\n\n"
    if social_changes:
      output += f"👥 Social Sciences: {np.mean(social_changes):+.2f}% avg change ({len(social_changes)} subjects)\n\n"
    if other_changes:
      output += f"📋 Other: {np.mean(other_changes):+.2f}% avg change ({len(other_changes)} subjects)\n\n"

    print(data)
    print(output)

    return output

def get_base_model(adapter_repo):
    """
    Get base model for adapter via hub
    """
    try:
        # Download adapter config from hub
        config_path = hf_hub_download(repo_id=adapter_repo, filename="adapter_config.json")
        
        with open(config_path, 'r') as f:
            adapter_config = json.load(f)
        
        base_model = adapter_config.get('base_model_name_or_path', '')
        
        if base_model:
            return base_model
        else:
            print(f"❌ Base model not found")
            return "Base model not found"
        
    except Exception as e:
        print(f"Hub base model check failed: {e}")
        return "Base model not found"

# Gradio interface
def disable_button():
    return gr.Button(interactive=False)

def enable_button():
    return gr.Button(interactive=True)

def update_evaluate_button(base_model):
    if base_model != "Base model not found":
        return gr.Button(interactive=True)
    else:
        return gr.Button(interactive=False)

def process_text_to_markdown(text_input):
    """
    This function takes text from a Textbox and returns it for Markdown display.
    """
    return text_input

def hide_output_text(text_input):
  return gr.Textbox(
    visible=False
  )

def set_moe(count):
  if count == "1":
    return gr.Radio(value="8-12%")
  elif count == "2":
    return gr.Radio(value="6-9%")
  elif count == "5":
    return gr.Radio(value="4-6%")
  elif count == "10":
    return gr.Radio(value="3-4%")
  elif count == "20":
    return gr.Radio(value="2-3%")
  elif count == "50":
    return gr.Radio(value="1.5-2%")
  elif count == "100":
    return gr.Radio(value="1-1.5%")

with gr.Blocks() as interface:
    
    # Add a title and description for the app.
    gr.Markdown("# flexMMLU Fine-tune vs. Base Model Evaluator")
    gr.Markdown("### A fast, flexible, fully subject-balanced MMLU benchmark comparison tool for fine-tuned adapters and their base models")
    gr.Markdown("Compare your adapter to its base model head-to-head in just minutes using this accelerated [**MMLU (Massive Multitask Language Understanding)**](https://huggingface.co/datasets/cais/mmlu) benchmark evaluation.")
    gr.Markdown("This app uses an efficient and effective, **subject-balanced random sampling**, zero-shot **evaluation across all 57 academic and professional subjects** in the MMLU. It provides a relative performance comparison of an adapter and its base model with a side-by-side score breakdown of overall, subject-by-subject, and domain-by-domain scores.")
    gr.Markdown("Whether you’re validating a new LoRA, demonstrating gains for a project, or just curious how much your fine-tuning really moved the needle, this tool gives you a **fast, standardized report and visualization** of its results.")
    with gr.Accordion("Want it even faster?", open=False):
        gr.Markdown("This Space runs on the free “CPU basic” (2 vCPU, 16 GB RAM) hardware option. You can duplicate this Space in your own account and configure it to use a more powerful paid hardware configuration for an even faster evaluation.")
    with gr.Accordion("Using private or gated models?", open=False):
        gr.Markdown("In order to use your own private models or public gated models, obtain necessary gated model user access, duplicate this space, and set your own account HF_TOKEN environment variable/secret for the duplicated space.")
    
    with gr.Tab("Evaluation Tool"):
    # Add input components for the user to set the adapter and base model.

        with gr.Row():
          adapter = gr.Textbox(
            label="Adapter (required)",
            show_label=True,
            max_lines=1,
            placeholder="adapter account/name (e.g., highlowmedia/gpt-oss-20b-lora-enlightened)",
            container=True,
            scale=10,
          )
          submit_adapter_button = gr.Button(
              "Get Base Model",
              scale=0,
              variant="secondary"
          )

        with gr.Row():
          base_model = gr.Textbox(
            label="Base Model (auto-filled)",
            max_lines=1,
            interactive=False,
            container=True,
          )

        submit_adapter_button.click(
            fn=get_base_model,
            inputs=adapter,
            outputs=base_model
        )

        questions_per_subject_count = gr.Radio(
            ["1", "2", "5", "10", "20", "50", "100"],
            value="20",
            label="Questions Per MMLU Subject",
            interactive=True,
            info="The fewer questions per subject, the sooner the evaluation will finish; however, the trade-off is an inverse increase in the probable margin of error for the results. 20 to 50 questions per subject is recommended for a generally accepted margin of error of ~2%. (Read the 'Speed vs. Accuracy' section for more information.)"
        )
        margin_of_error_display = gr.Radio(
            ["8-12%", "6-9%", "4-6%", "3-4%", "2-3%", "1.5-2%", "1-1.5%"],
            value="2-3%",
            label="Margin of Error (auto-filled)",
            interactive=False
        )
        questions_per_subject_count.change(
            fn=set_moe,
            inputs=questions_per_subject_count,
            outputs=margin_of_error_display
        )
          
        # Add a button to trigger the evaluation.
        evaluate_button = gr.Button(
            "Evaluate",
            variant="primary",
            interactive=False
        )

        output_text = gr.Textbox(
            label="",
            show_label=False,
            lines=5, 
            interactive=False,
            show_copy_button=True,
        )
        
        output_text_markdown = gr.Markdown("""
        
        """)

        evaluate_button.click(
            fn=disable_button,
            inputs=None,
            outputs=evaluate_button,
            queue=False
        ).then(
            fn=run_comparison,
            inputs=[
                base_model,
                adapter,
                questions_per_subject_count
            ],
            outputs=output_text
        ).then(
            fn=process_text_to_markdown,
            inputs=output_text,
            outputs=output_text_markdown
        ).then(
            fn=enable_button,
            inputs=None,
            outputs=evaluate_button,
            queue=False
        ).then(
            fn=hide_output_text,
            inputs=output_text,
            outputs=output_text
        )

        # turn on evaluate button only after base model found
        # base_model textbox can only be changed via get_base_model via submit_adapter_button
        base_model.change(fn=update_evaluate_button, inputs=base_model, outputs=evaluate_button)

    with gr.Tab("Speed vs. Accuracy"):
        gr.Markdown("The balanced random approach of flexMMLU represents a practical middle ground between computational efficiency and methodological simplicity (in contrast with human-curated subset generation for instance). While you can perhaps go too far in your quest for speed and compuational efficiency, this tool gives you the option to choose the degree of trade-off between speed and accuracy that suits your needs.")
        gr.Markdown("A scaling formula for statistical error based upon sample size informs the table of estimates below which also reflects the evaluation tool configuration options. While the lowest subset size of 1 question per MMLU subject is likely not useful for your purposes (unless you want extreme computational efficiency that still well outperforms guessing), it should be noted that there are diminishing returns on performance at 100 questions or more per subject.")
        gr.Markdown("""
            | Subset Size | Questions per Subject | Estimated Error Range |
            |-------------|----------------------|----------------------|
            | 57          | 1                    | 8-12%               |
            | 114         | 2                    | 6-9%                |
            | 285         | 5                    | 4-6%                |
            | 570         | 10                   | 3-4%                |
            | 1,140       | 20                   | 2-3%                |
            | 2,850       | 50                   | 1.5-2%              |
            | 5,700       | 100                  | 1-1.5%              |
        """)
    with gr.Tab("flexMMLU"):
        gr.Markdown("The foundational LLM benchmark MMLU (Massive Multitask Language Understanding), contains 14,042 test questions on 57 different subjects across STEM, humanities, social sciences, and professional fields. The breadth of it’s subject matter makes it a still powerful tool for zero-shot and few-shot evaluations of large language models without extensive task-specific training or instructions, but it’s a large dataset, and using it can be computationally expensive and time-consuming.")
        gr.Markdown("Inspired by tinyMMLU and its conceptual precursors, flexMMLU is implemented here as an MMLU accelerator with margin of error mitigtaion that is grounded upon the utility of “stratified random sampling” - which is to say, random but balanced selection from “sub-scenarios” - MMLU subjects in our case. ([tinyBenchmarks: evaluating LLMs with fewer examples](https://arxiv.org/abs/2402.14992), Felipe Maia Polo, et al. 2024)")
        gr.Markdown("Working with a subset of MMLU, regardless of how that subset is generated, mitigates computational overhead but, of course, introduces added margin of error. Interestingly, the balanced random sampling of flexMMLU, by which there is equal representation across subjects in the subset, not only generally mitigates added margin of error relative to the use of a subset of the full MMLU dataset but may give flexMMLU a performance advantage over MMLU due to the composition of MMLU itself - specifically, with respect to the uneven distribution of questions across domains in the full MMLU test dataset.")
        gr.Markdown("All subjects in the full dataset are reprsented by at least 100 questions, but some some subjects have as many as 6 times as many questions as others. This imbalance is overcome in flexMMLU, and, accordingly, it would be fair to assume (though not scientifically verified) that flexMMLU makes inroads here wiht repsect to closing the margin of error gap with MMLU. flexMMLU prevents bias from subject representation imbalance and ultimately provides more meaningful subject-by-subject comparisons.")
        gr.Markdown("*There has been identification of errors in MMLU over the years as well - an impetus for new “massive” formulations like MMLU-Pro - and not surprisingly, it has been reported that the errors in MMLU are not evenly represented across domains either. This last matter is not mitigated by aotomated balanced subset generation from original MMLU dataset used here.*")
    
if __name__ == "__main__":
    interface.launch()