diff --git a/README.md b/README.md
index 4567b272..585172c6 100644
--- a/README.md
+++ b/README.md
@@ -1,278 +1,308 @@
-# Sparse Attention Hub
+
+
+
+
+
+
+
-A comprehensive framework for sparse attention mechanisms in deep learning models. This project provides implementations of various sparse attention algorithms, benchmarking tools, and seamless integration with popular model frameworks like HuggingFace Transformers.
+
-## ๐ Features
+**A framework for implementing and evaluating sparse attention mechanisms in transformer models**
-- **Multiple Sparse Attention Algorithms**: Implementations of efficient attention mechanisms including Double Sparsity, Hash Attention, and various research-oriented masking strategies
-- **Framework Integration**: Seamless integration with HuggingFace Transformers and other popular frameworks
-- **Comprehensive Benchmarking**: Built-in support for LongBench, Loogle, InfBench, and custom benchmarks
-- **Advanced Metrics**: Micro-metrics logging system for detailed performance analysis
-- **Visualization Tools**: Generate plots and heatmaps for attention pattern analysis
-- **Extensible Architecture**: Modular design for easy extension and customization
+[](https://www.python.org/)
+[](https://pytorch.org/)
+[](https://huggingface.co/transformers/)
+[](LICENSE)
+[](https://pytest.org/)
+[](https://github.com/psf/black)
-## ๐ฆ Installation
+
-### From Source
+Sparse Attention Hub provides efficient implementations of sparse attention algorithms for transformer models. It includes research-oriented attention mechanisms, production-ready implementations, comprehensive benchmarking tools, and seamless integration with HuggingFace Transformers.
+
+## Key Features
+
+- **Multiple Sparse Attention Algorithms**: Hash Attention, Double Sparsity, and configurable masking strategies
+- **HuggingFace Integration**: Drop-in replacement for standard attention with minimal code changes
+- **Comprehensive Benchmarking**: Built-in support for LongBench, Ruler, AIME, InfiniteBench, and custom datasets
+- **Research Tools**: Flexible masker system for experimenting with novel attention patterns
+- **Performance Monitoring**: Detailed metrics logging and visualization capabilities
+- **Modular Architecture**: Easy to extend with custom attention mechanisms
+
+## Installation
+
+### Using Poetry (Recommended)
```bash
+# Install Poetry if you haven't already
+curl -sSL https://install.python-poetry.org | python3 -
+
+# Clone and install the project
git clone https://github.com/xAlg-ai/sparse-attention-hub.git
cd sparse-attention-hub
-pip install -e .
+poetry install
```
-### Dependencies
+### Development Installation
```bash
-pip install -r requirements.txt
-```
+# Install with development dependencies
+poetry install --with dev
-## ๐๏ธ Architecture
+# Install with all optional dependencies
+poetry install --with dev,docs,benchmarks
+```
-The framework is organized into several key modules:
+### From Source (pip)
-### Sparse Attention
-- **Base Classes**: `SparseAttention`, `EfficientAttention`, `ResearchAttention`
-- **Efficient Implementations**: `DoubleSparsity`, `HashAttention`
-- **Research Maskers**: Various masking strategies for attention patterns
-- **Generators**: Integration interfaces for different frameworks
+```bash
+git clone https://github.com/xAlg-ai/sparse-attention-hub.git
+cd sparse-attention-hub
+pip install -e .
+```
-### Adapter System
-- **ModelAdapterHF**: Unified adapter for HuggingFace integration
-- **Request/RequestResponse**: Structured request/response handling
-- **ModelAdapter**: Abstract base class for model adapters
-- **ModelHubAdapterInterface**: Interface for model hosting libraries
-- **SparseAttentionAdapterInterface**: Interface for sparse attention integration
+### Requirements
-### Benchmarking
-- **Benchmark**: Abstract benchmark interface
-- **Datasets**: LongBench, Loogle, InfBench implementations
-- **BenchmarkExecutor**: Execution and result management
-- **ResultStorage**: Persistent storage for benchmark results
+The project requires Python 3.9+ and key dependencies include PyTorch, Transformers, and several benchmarking libraries. All dependencies are managed through Poetry and defined in `pyproject.toml`.
-### Metrics & Visualization
-- **MicroMetricLogger**: Singleton logger for detailed metrics
-- **MicroMetrics**: TopkRecall, LocalError, SampleVariance implementations
-- **PlotGenerator**: Visualization tools with multiple granularity levels
+## Quick Start
-## ๐ Quick Start
+### Basic Usage with HuggingFace Models
```python
-from sparse_attention_hub.adapters import ModelAdapterHF, Request, RequestResponse
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+from sparse_attention_hub.adapters import ModelAdapterHF, Request
+from sparse_attention_hub.sparse_attention import ResearchAttentionConfig
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
LocalMaskerConfig, SinkMaskerConfig
)
-from sparse_attention_hub.benchmark import BenchmarkExecutor
-from sparse_attention_hub.plotting import PlotGenerator, Granularity
-
-# Create sparse attention configuration
-local_config = LocalMaskerConfig(window_size=16)
-sink_config = SinkMaskerConfig(sink_size=4)
-sparse_attention_config = ResearchAttentionConfig(
- masker_configs=[local_config, sink_config]
-)
+
+# Configure sparse attention (StreamingLLM pattern)
+sparse_config = ResearchAttentionConfig(masker_configs=[
+ SinkMaskerConfig(sink_size=4), # Keep first 4 tokens
+ LocalMaskerConfig(window_size=64) # Keep last 64 tokens
+])
# Create model adapter
adapter = ModelAdapterHF(
model_name="microsoft/DialoGPT-small",
- sparse_attention_config=sparse_attention_config,
+ sparse_attention_config=sparse_config,
device="auto"
)
# Create request
request = Request(
- context="The capital of France is Paris. It is known for the Eiffel Tower.",
- questions=["What is the capital of France?", "What is Paris known for?"]
+ context="Paris is the capital of France. It is famous for the Eiffel Tower and the Louvre Museum.",
+ questions=["What is the capital of France?", "What is Paris famous for?"]
)
-# Process request with sparse attention
+# Process with sparse attention
with adapter.enable_sparse_mode():
response = adapter.process_request(request)
- print(response.responses) # ['Paris', 'The Eiffel Tower']
+ print(response.responses)
-# Process request with dense attention (default)
+# Process with dense attention (default)
response = adapter.process_request(request)
print(response.responses)
-
-# Run benchmarks and visualizations
-benchmark_executor = BenchmarkExecutor()
-plot_generator = PlotGenerator()
-plot_path = plot_generator.generate_plot(Granularity.PER_HEAD)
```
-## ๐ Benchmarking
+### Available Sparse Attention Methods
-The framework supports multiple benchmark datasets and integrates seamlessly with the new adapter system:
+```python
+# Hash Attention
+from sparse_attention_hub.sparse_attention import HashAttentionConfig, EfficientAttentionConfig
+
+hash_config = EfficientAttentionConfig(
+ implementation_config=HashAttentionConfig(
+ num_hash_functions=4,
+ hash_budget=512
+ )
+)
-- **LongBench**: Long-context understanding tasks
-- **Loogle**: Dependency tracking benchmarks
-- **InfBench**: Infinite context benchmarks
+# Double Sparsity
+from sparse_attention_hub.sparse_attention import DoubleSparsityConfig
-```python
-from sparse_attention_hub.benchmark import BenchmarkExecutor
-from sparse_attention_hub.benchmark.datasets import LongBench
-from sparse_attention_hub.adapters import ModelAdapterHF
+double_config = EfficientAttentionConfig(
+ implementation_config=DoubleSparsityConfig(
+ channel_config=ChannelConfig(channels_per_head=32),
+ sparsity_ratio=0.1
+ )
+)
-# Create adapter with sparse attention
-adapter = ModelAdapterHF(
- model_name="microsoft/DialoGPT-small",
- sparse_attention_config=your_sparse_config
+# Custom masking patterns
+from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
+ TopKMaskerConfig, OracleTopK
)
-# Run benchmarks
-executor = BenchmarkExecutor()
-benchmark = LongBench()
-results = executor.evaluate(benchmark, adapter)
+custom_config = ResearchAttentionConfig(masker_configs=[
+ TopKMaskerConfig(k=128),
+ SinkMaskerConfig(sink_size=4)
+])
```
-## ๐ Metrics and Logging
+## Benchmarking
-Track detailed performance metrics:
+### Running Benchmarks
```python
-from sparse_attention_hub.metrics import MicroMetricLogger
-from sparse_attention_hub.metrics.implementations import TopkRecall
+from benchmark import BenchmarkExecutor, BenchmarkConfig
+from benchmark.executor_config import AdapterConfig
-logger = MicroMetricLogger()
-metric = TopkRecall(k=10)
-logger.register_metric(metric)
-logger.enable_metric_logging(metric)
+# Configure benchmark
+benchmark_config = BenchmarkConfig(
+ benchmark_name="longbench",
+ subsets=["narrativeqa", "qasper"]
+)
+
+adapter_config = AdapterConfig(
+ model_name="microsoft/DialoGPT-small",
+ sparse_attention_config=sparse_config
+)
-# Log metrics during model execution
-logger.log("layer_1", metric, computed_value)
+# Run evaluation
+executor = BenchmarkExecutor(
+ result_dir="./results",
+ gpus=[0, 1],
+ max_concurrent_runs=2
+)
+
+results = executor.run_benchmarks(
+ benchmark_configs=[benchmark_config],
+ adapter_configs=[adapter_config]
+)
```
-## ๐จ Visualization
+### Supported Benchmarks
-Generate attention pattern visualizations:
+- **LongBench**: Long-context understanding tasks (22 standard + 13 extended datasets)
+- **Ruler**: Synthetic tasks for testing context length capabilities
+- **AIME 2024/2025**: Mathematical reasoning benchmarks
+- **InfiniteBench**: Infinite context evaluation tasks
+- **Loogle**: Dependency tracking and retrieval tasks
+- **ZeroScrolls**: Long document understanding
+- **Custom benchmarks**: Easy to add new evaluation datasets
-```python
-from sparse_attention_hub.plotting import PlotGenerator, Granularity
+## Architecture
-generator = PlotGenerator()
+### Core Components
-# Generate different types of plots
-token_plot = generator.generate_plot(Granularity.PER_TOKEN, data)
-head_plot = generator.generate_plot(Granularity.PER_HEAD, data)
-layer_plot = generator.generate_plot(Granularity.PER_LAYER, data)
+- **`sparse_attention/`**: Core sparse attention implementations
+ - `base.py`: Abstract interfaces for all attention mechanisms
+ - `efficient_attention/`: Production-ready algorithms (Hash Attention, Double Sparsity)
+ - `research_attention/`: Experimental masking strategies for research
+ - `utils/`: Common utilities and mask operations
+
+- **`adapters/`**: Integration layer for different frameworks
+ - `huggingface.py`: HuggingFace Transformers integration
+ - `model_servers/`: Centralized model and tokenizer management
+ - `base.py`: Abstract adapter interfaces
+
+- **`benchmark/`**: Comprehensive benchmarking system
+ - Individual benchmark implementations (LongBench, Ruler, etc.)
+ - `executor.py`: Parallel execution with GPU management
+ - `base.py`: Abstract benchmark interface
+
+- **`metric_logging/`**: Performance monitoring and analysis
+- **`plotting/`**: Visualization tools for attention patterns and results
+
+## Advanced Usage
+
+### Custom Maskers
+
+```python
+from sparse_attention_hub.sparse_attention.research_attention.maskers.base import ResearchMasker
+from sparse_attention_hub.sparse_attention.utils import Mask
+
+class CustomMasker(ResearchMasker):
+ def add_mask(self, keys, queries, values, previous_mask, **kwargs):
+ # Implement custom masking logic
+ custom_mask = self.create_custom_mask(queries.shape)
+ return previous_mask.combine_with(custom_mask)
```
-## ๐งช Testing
+### Metrics and Monitoring
+
+```python
+from sparse_attention_hub.metric_logging import MicroMetricLogger
+
+logger = MicroMetricLogger()
+# Metrics are automatically logged when using ResearchAttention
+# Access logged data for analysis
+metrics_data = logger.get_logged_metrics()
+```
-Run the test suite:
+## Testing
```bash
# Run all tests
python scripts/run_tests.py --type all
-# Run only unit tests
+# Run specific test categories
python scripts/run_tests.py --type unit
-
-# Run specific test
-python scripts/run_tests.py --type specific --test-path tests/unit/test_metrics.py
-
-# Discover available tests
-python scripts/run_tests.py --discover
+python scripts/run_tests.py --type integration
# Run tests with coverage
make test-coverage
```
-## ๐ง Development Tools
+## Development
-### Code Formatting and Linting
-
-The project uses comprehensive linting and formatting tools:
+### Code Quality Tools
```bash
# Install development dependencies
-pip install -r requirements-dev.txt
+poetry install --with dev
+
+# Activate Poetry shell
+poetry shell
# Format code
bash scripts/format.sh
-# Run all linting checks
+# Run linting
bash scripts/lint.sh
-# Run specific linters
-bash scripts/lint.sh --flake8
-bash scripts/lint.sh --mypy
-bash scripts/lint.sh --pylint
-bash scripts/lint.sh --bandit
-
-# Using Make commands
-make format # Format code
-make lint # Run all linting
-make dev-check # Quick format + lint check
-```
-
-### Pre-commit Hooks
-
-Set up pre-commit hooks for automatic code quality checks:
-
-```bash
-# Install pre-commit hooks
+# Set up pre-commit hooks
pre-commit install
-
-# Run pre-commit on all files
-pre-commit run --all-files
-
-# Update pre-commit hooks
-pre-commit autoupdate
```
-### Development Workflow
-
-```bash
-# Complete development setup
-make dev-setup
-
-# Run all development checks
-make dev-check
-
-# Simulate CI pipeline
-make ci
-```
-
-## ๐ Project Structure
+### Project Structure
```
sparse-attention-hub/
โโโ sparse_attention_hub/ # Main package
-โ โโโ sparse_attention/ # Sparse attention implementations
-โ โโโ adapters/ # Model adapter system
-โ โโโ benchmark/ # Benchmarking tools
-โ โโโ metrics/ # Metrics and logging
-โ โโโ plotting/ # Visualization tools
-โ โโโ testing/ # Testing utilities
-โโโ tests/ # Test suite
-โ โโโ unit/ # Unit tests
-โ โโโ integration/ # Integration tests
-โโโ scripts/ # Utility scripts
-โโโ tutorials/ # Tutorial notebooks and examples
+โ โโโ sparse_attention/ # Sparse attention implementations
+โ โโโ adapters/ # Model integration adapters
+โ โโโ metric_logging/ # Performance monitoring
+โ โโโ plotting/ # Visualization tools
+โโโ benchmark/ # Benchmarking framework
+โ โโโ longbench/ # LongBench implementation
+โ โโโ ruler/ # Ruler benchmark
+โ โโโ AIME2024/, AIME2025/ # Mathematical reasoning
+โ โโโ scripts/ # Benchmark execution scripts
+โโโ tests/ # Test suite
+โ โโโ unit/ # Unit tests
+โ โโโ integration/ # Integration tests
+โโโ tutorials/ # Examples and tutorials
โโโ docs/ # Documentation
```
-## ๐ค Contributing
-
-We welcome contributions! Please see our contributing guidelines for details on:
+## Contributing
-- Code style and formatting
-- Testing requirements
-- Documentation standards
-- Pull request process
+We welcome contributions to improve sparse attention implementations, add new benchmarks, or enhance the framework. Please ensure all code follows the project's formatting standards and includes appropriate tests.
-## ๐ License
+## License
-This project is licensed under the MIT License - see the LICENSE file for details.
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
-## ๐ Links
+## Links
- **Repository**: https://github.com/xAlg-ai/sparse-attention-hub
- **Documentation**: https://sparse-attention-hub.readthedocs.io
- **Issues**: https://github.com/xAlg-ai/sparse-attention-hub/issues
-## ๐ Acknowledgments
+## Acknowledgments
-This project implements and extends various sparse attention mechanisms from the research community. We acknowledge the original authors of these algorithms and the open-source community for their contributions.
\ No newline at end of file
+This project implements sparse attention mechanisms from various research papers. We acknowledge the original authors and the open-source community for their contributions to advancing efficient attention algorithms.
\ No newline at end of file
diff --git a/benchmark/AIME2024/__init__.py b/benchmark/AIME2024/__init__.py
index bb567237..58102847 100644
--- a/benchmark/AIME2024/__init__.py
+++ b/benchmark/AIME2024/__init__.py
@@ -8,4 +8,4 @@
from .calculate_metrics import calculate_metrics
from .aime2024 import AIME2024
-__all__ = ["calculate_metrics", "AIME2024"]
\ No newline at end of file
+__all__ = ["calculate_metrics", "AIME2024"]
diff --git a/benchmark/AIME2024/aime2024.py b/benchmark/AIME2024/aime2024.py
index 8ffc779c..c068813c 100644
--- a/benchmark/AIME2024/aime2024.py
+++ b/benchmark/AIME2024/aime2024.py
@@ -29,22 +29,23 @@ class AIME2024(Benchmark):
# AIME2024 has a single dataset
all_datasets: List[str] = ["aime2024"]
-
+
benchmark_name: str = "aime2024"
huggingface_dataset_id: str = "xAlg-AI/att-hub-aime2024"
def _load_datasets(self) -> pd.DataFrame:
"""Load AIME2024 dataset.
-
+
AIME2024 uses a single dataset with all problems.
-
+
Returns:
pandas DataFrame with all AIME2024 problems.
"""
print(f"Loading AIME2024 dataset")
-
+
try:
from datasets import load_dataset
+
dataset = load_dataset(self.huggingface_dataset_id, split="test")
df = dataset.to_pandas()
df["task"] = "aime2024" # Ensure task column exists
@@ -72,7 +73,7 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
# Use the calculate_metrics function from HashAttention evaluation
metrics: Dict[str, Any] = calculate_metrics(results_df)
-
+
# Format the results for consistency with other benchmarks
overall_metrics: Dict[str, Any] = {
"overall_score": round(metrics["accuracy"], 4),
@@ -84,16 +85,19 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
"task_scores": {
"aime2024": {
"accuracy": round(metrics["accuracy"], 4),
- "extraction_success_rate": round(metrics["extraction_success_rate"], 4)
+ "extraction_success_rate": round(
+ metrics["extraction_success_rate"], 4
+ ),
}
},
- "summary": {
- "total_tasks": 1,
- "total_samples": len(results_df)
- }
+ "summary": {"total_tasks": 1, "total_samples": len(results_df)},
}
-
- print(f" โ AIME2024 Accuracy: {metrics['accuracy']:.3f} ({metrics['accuracy']*100:.1f}%)")
- print(f" โ Extraction Success Rate: {metrics['extraction_success_rate']:.3f} ({metrics['extraction_success_rate']*100:.1f}%)")
-
- return overall_metrics
\ No newline at end of file
+
+ print(
+ f" โ AIME2024 Accuracy: {metrics['accuracy']:.3f} ({metrics['accuracy']*100:.1f}%)"
+ )
+ print(
+ f" โ Extraction Success Rate: {metrics['extraction_success_rate']:.3f} ({metrics['extraction_success_rate']*100:.1f}%)"
+ )
+
+ return overall_metrics
diff --git a/benchmark/AIME2024/calculate_metrics.py b/benchmark/AIME2024/calculate_metrics.py
index 6f9270cf..7f882a6c 100644
--- a/benchmark/AIME2024/calculate_metrics.py
+++ b/benchmark/AIME2024/calculate_metrics.py
@@ -5,51 +5,53 @@
import pandas as pd
from typing import List, Dict, Any
+
def extract_boxed_answer(text: str) -> str:
"""
Extract the answer from \boxed{...} format in the text.
-
+
Args:
text: The model's response text
-
+
Returns:
The extracted answer as a string, or empty string if not found
"""
# Look for \boxed{...} pattern
- boxed_pattern = r'\\boxed\{([^}]*)\}'
+ boxed_pattern = r"\\boxed\{([^}]*)\}"
matches = re.findall(boxed_pattern, text)
-
+
if matches:
# Take the last boxed answer in case there are multiple
answer = matches[-1].strip()
-
+
# Extract just the number if there's additional formatting
# Handle cases like "033", "23", "$23$", etc.
- number_match = re.search(r'\d+', answer)
+ number_match = re.search(r"\d+", answer)
if number_match:
return number_match.group()
else:
return answer
-
+
# Fallback: look for numbers at the end of the text
# This handles cases where the model doesn't use \boxed format
- lines = text.strip().split('\n')
+ lines = text.strip().split("\n")
for line in reversed(lines):
if line.strip():
# Look for a number in the last non-empty line
- number_match = re.search(r'\b(\d{1,3})\b', line)
+ number_match = re.search(r"\b(\d{1,3})\b", line)
if number_match:
return number_match.group(1)
-
+
return ""
+
def normalize_answer(answer: str) -> str:
"""
Normalize an answer to a standard format.
-
+
Args:
answer: The answer string to normalize
-
+
Returns:
Normalized answer string
"""
@@ -59,34 +61,37 @@ def normalize_answer(answer: str) -> str:
return str(int(answer))
return answer
+
def calculate_metrics(df: pd.DataFrame) -> Dict[str, Any]:
"""
Calculate evaluation metrics for AIME2024 benchmark.
-
+
Args:
df: DataFrame with columns 'answer' (ground truth) and 'predicted_answer' (model output)
-
+
Returns:
Dictionary containing evaluation metrics
"""
- if 'predicted_answer' not in df.columns:
+ if "predicted_answer" not in df.columns:
raise ValueError("DataFrame must contain 'predicted_answer' column")
- if 'answer' not in df.columns:
+ if "answer" not in df.columns:
raise ValueError("DataFrame must contain 'answer' column")
-
+
total_problems = len(df)
correct_answers = 0
extraction_failures = 0
-
+
detailed_results = []
-
+
for idx, row in df.iterrows():
- ground_truth = normalize_answer(str(row['answer']))
- predicted_text = str(row['predicted_answer']) if pd.notna(row['predicted_answer']) else ""
-
+ ground_truth = normalize_answer(str(row["answer"]))
+ predicted_text = (
+ str(row["predicted_answer"]) if pd.notna(row["predicted_answer"]) else ""
+ )
+
# Extract the predicted answer
extracted_answer = extract_boxed_answer(predicted_text)
-
+
if not extracted_answer:
extraction_failures += 1
is_correct = False
@@ -95,35 +100,42 @@ def calculate_metrics(df: pd.DataFrame) -> Dict[str, Any]:
is_correct = extracted_answer == ground_truth
if is_correct:
correct_answers += 1
-
- detailed_results.append({
- 'id': row.get('id', f'problem_{idx}'),
- 'ground_truth': ground_truth,
- 'predicted_text': predicted_text,
- 'extracted_answer': extracted_answer,
- 'is_correct': is_correct,
- 'extraction_failed': not bool(extracted_answer)
- })
-
+
+ detailed_results.append(
+ {
+ "id": row.get("id", f"problem_{idx}"),
+ "ground_truth": ground_truth,
+ "predicted_text": predicted_text,
+ "extracted_answer": extracted_answer,
+ "is_correct": is_correct,
+ "extraction_failed": not bool(extracted_answer),
+ }
+ )
+
# Calculate metrics
accuracy = correct_answers / total_problems if total_problems > 0 else 0.0
- extraction_success_rate = (total_problems - extraction_failures) / total_problems if total_problems > 0 else 0.0
-
+ extraction_success_rate = (
+ (total_problems - extraction_failures) / total_problems
+ if total_problems > 0
+ else 0.0
+ )
+
metrics = {
- 'accuracy': accuracy,
- 'correct_answers': correct_answers,
- 'total_problems': total_problems,
- 'extraction_success_rate': extraction_success_rate,
- 'extraction_failures': extraction_failures,
- 'detailed_results': detailed_results
+ "accuracy": accuracy,
+ "correct_answers": correct_answers,
+ "total_problems": total_problems,
+ "extraction_success_rate": extraction_success_rate,
+ "extraction_failures": extraction_failures,
+ "detailed_results": detailed_results,
}
-
+
return metrics
+
def print_metrics_summary(metrics: Dict[str, Any]) -> None:
"""
Print a formatted summary of the evaluation metrics.
-
+
Args:
metrics: Dictionary containing evaluation metrics
"""
@@ -132,34 +144,39 @@ def print_metrics_summary(metrics: Dict[str, Any]) -> None:
print(f"Total Problems: {metrics['total_problems']}")
print(f"Correct Answers: {metrics['correct_answers']}")
print(f"Accuracy: {metrics['accuracy']:.3f} ({metrics['accuracy']*100:.1f}%)")
- print(f"Extraction Success Rate: {metrics['extraction_success_rate']:.3f} ({metrics['extraction_success_rate']*100:.1f}%)")
+ print(
+ f"Extraction Success Rate: {metrics['extraction_success_rate']:.3f} ({metrics['extraction_success_rate']*100:.1f}%)"
+ )
print(f"Extraction Failures: {metrics['extraction_failures']}")
-
- if metrics['extraction_failures'] > 0:
- print(f"\nNote: {metrics['extraction_failures']} problems had answer extraction failures.")
+
+ if metrics["extraction_failures"] > 0:
+ print(
+ f"\nNote: {metrics['extraction_failures']} problems had answer extraction failures."
+ )
print("These are counted as incorrect answers.")
+
if __name__ == "__main__":
# Test the metrics calculation
test_data = {
- 'answer': ['23', '33', '156', '902'],
- 'predicted_answer': [
- 'The answer is \\boxed{23}.',
- 'After solving, we get \\boxed{033}.',
- 'Therefore, the answer is \\boxed{156}.',
- 'The final answer is 902.' # Test fallback extraction
+ "answer": ["23", "33", "156", "902"],
+ "predicted_answer": [
+ "The answer is \\boxed{23}.",
+ "After solving, we get \\boxed{033}.",
+ "Therefore, the answer is \\boxed{156}.",
+ "The final answer is 902.", # Test fallback extraction
],
- 'id': ['2024-I-1', '2024-I-2', '2024-I-3', '2024-I-4']
+ "id": ["2024-I-1", "2024-I-2", "2024-I-3", "2024-I-4"],
}
-
+
test_df = pd.DataFrame(test_data)
metrics = calculate_metrics(test_df)
print_metrics_summary(metrics)
-
+
print("\nDetailed Results:")
- for result in metrics['detailed_results']:
+ for result in metrics["detailed_results"]:
print(f"ID: {result['id']}")
print(f" Ground Truth: {result['ground_truth']}")
print(f" Extracted: {result['extracted_answer']}")
print(f" Correct: {result['is_correct']}")
- print()
\ No newline at end of file
+ print()
diff --git a/benchmark/AIME2024/create_huggingface_dataset.py b/benchmark/AIME2024/create_huggingface_dataset.py
index 0a419581..5f470c23 100644
--- a/benchmark/AIME2024/create_huggingface_dataset.py
+++ b/benchmark/AIME2024/create_huggingface_dataset.py
@@ -21,6 +21,7 @@
which is standard in mathematical competition contexts.
"""
+
def create_aime2024_dataset():
"""
Process the AIME2024 dataset and convert it to the standardized benchmark format.
@@ -28,10 +29,10 @@ def create_aime2024_dataset():
# Load the original dataset
dataset = load_dataset("Maxwell-Jia/AIME_2024")
df = dataset["train"].to_pandas()
-
+
# Create the standardized format
processed_data = []
-
+
for _, row in df.iterrows():
# Format the problem with clear instructions about the boxed answer format
context = f"""Solve the following AIME (American Invitational Mathematics Examination) problem.
@@ -41,26 +42,29 @@ def create_aime2024_dataset():
Instructions:
- The answer should be an integer between 0 and 999
- Please reason step by step, and put your final answer within \\boxed{{...}} format"""
-
+
question = "What is the answer to this problem?"
-
+
# The answer prefix encourages the model to show work before the final answer
answer_prefix = ""
-
- processed_data.append({
- 'context': context,
- 'question': question,
- 'answer_prefix': answer_prefix,
- 'answer': str(row['Answer']), # Convert to string for consistency
- 'id': row['ID'],
- 'max_new_tokens': 32000, # Allow comprehensive step-by-step solutions
- })
-
+
+ processed_data.append(
+ {
+ "context": context,
+ "question": question,
+ "answer_prefix": answer_prefix,
+ "answer": str(row["Answer"]), # Convert to string for consistency
+ "id": row["ID"],
+ "max_new_tokens": 32000, # Allow comprehensive step-by-step solutions
+ }
+ )
+
# Convert to Dataset
processed_dataset = Dataset.from_pandas(pd.DataFrame(processed_data))
-
+
return processed_dataset
+
if __name__ == "__main__":
# Test the dataset creation
processed_dataset = create_aime2024_dataset()
@@ -68,4 +72,6 @@ def create_aime2024_dataset():
print("\nFirst example:")
print(processed_dataset[0])
- processed_dataset.push_to_hub("xAlg-AI/att-hub-aime2024", config_name=f"aime2024", split="test")
+ processed_dataset.push_to_hub(
+ "xAlg-AI/att-hub-aime2024", config_name=f"aime2024", split="test"
+ )
diff --git a/benchmark/AIME2024/test_aime2024.py b/benchmark/AIME2024/test_aime2024.py
index 5e42c3fe..90c77ee6 100644
--- a/benchmark/AIME2024/test_aime2024.py
+++ b/benchmark/AIME2024/test_aime2024.py
@@ -11,29 +11,38 @@
import sys
import os
+
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from create_huggingface_dataset import create_aime2024_dataset
from calculate_metrics import calculate_metrics, print_metrics_summary
import pandas as pd
+
def test_dataset_creation():
"""Test that the dataset can be created successfully."""
print("Testing AIME2024 dataset creation...")
-
+
dataset = create_aime2024_dataset()
-
+
print(f"โ Dataset created successfully with {len(dataset)} problems")
-
+
# Verify dataset structure
- required_keys = ['context', 'question', 'answer_prefix', 'answer', 'id', 'max_new_tokens']
+ required_keys = [
+ "context",
+ "question",
+ "answer_prefix",
+ "answer",
+ "id",
+ "max_new_tokens",
+ ]
first_example = dataset[0]
-
+
for key in required_keys:
assert key in first_example, f"Missing required key: {key}"
-
+
print("โ Dataset structure is correct")
-
+
# Show some examples
print(f"\nSample problems:")
for i in range(min(3, len(dataset))):
@@ -42,99 +51,114 @@ def test_dataset_creation():
print(f" Answer: {example['answer']}")
print(f" Problem preview: {example['context'][:100]}...")
print()
-
+
return dataset
+
def test_metrics_calculation():
"""Test the metrics calculation with sample data."""
print("Testing AIME2024 metrics calculation...")
-
+
# Create test data with various scenarios
test_data = {
- 'answer': ['23', '33', '156', '902', '45'],
- 'predicted_answer': [
- 'The answer is \\boxed{23}.', # Correct with boxed format
- 'After solving, we get \\boxed{033}.', # Correct with leading zeros
- 'Therefore, the answer is \\boxed{156}.', # Correct
- 'The final answer is 902.', # Correct without boxed format
- 'I think the answer is \\boxed{44}.' # Incorrect
+ "answer": ["23", "33", "156", "902", "45"],
+ "predicted_answer": [
+ "The answer is \\boxed{23}.", # Correct with boxed format
+ "After solving, we get \\boxed{033}.", # Correct with leading zeros
+ "Therefore, the answer is \\boxed{156}.", # Correct
+ "The final answer is 902.", # Correct without boxed format
+ "I think the answer is \\boxed{44}.", # Incorrect
],
- 'id': ['2024-I-1', '2024-I-2', '2024-I-3', '2024-I-4', '2024-I-5']
+ "id": ["2024-I-1", "2024-I-2", "2024-I-3", "2024-I-4", "2024-I-5"],
}
-
+
test_df = pd.DataFrame(test_data)
metrics = calculate_metrics(test_df)
-
+
print_metrics_summary(metrics)
-
+
# Verify expected results
- assert metrics['total_problems'] == 5, f"Expected 5 problems, got {metrics['total_problems']}"
- assert metrics['correct_answers'] == 4, f"Expected 4 correct answers, got {metrics['correct_answers']}"
- assert metrics['accuracy'] == 0.8, f"Expected accuracy 0.8, got {metrics['accuracy']}"
- assert metrics['extraction_failures'] == 0, f"Expected 0 extraction failures, got {metrics['extraction_failures']}"
-
+ assert (
+ metrics["total_problems"] == 5
+ ), f"Expected 5 problems, got {metrics['total_problems']}"
+ assert (
+ metrics["correct_answers"] == 4
+ ), f"Expected 4 correct answers, got {metrics['correct_answers']}"
+ assert (
+ metrics["accuracy"] == 0.8
+ ), f"Expected accuracy 0.8, got {metrics['accuracy']}"
+ assert (
+ metrics["extraction_failures"] == 0
+ ), f"Expected 0 extraction failures, got {metrics['extraction_failures']}"
+
print("โ Metrics calculation is correct")
-
+
return metrics
+
def test_integration():
"""Test integration with the main evaluation system."""
print("Testing AIME2024 integration...")
-
+
# Test that the dataset can be imported and used
try:
- sys.path.append('..')
+ sys.path.append("..")
from AIME2024.calculate_metrics import calculate_metrics
from AIME2024.create_huggingface_dataset import create_aime2024_dataset
-
+
# Test dataset creation through import
dataset = create_aime2024_dataset()
print(f"โ Integration test passed - dataset has {len(dataset)} examples")
-
+
# Verify the dataset can be converted to pandas
df = dataset.to_pandas()
print(f"โ Dataset conversion to pandas successful - shape: {df.shape}")
-
+
return True
-
+
except Exception as e:
print(f"โ Integration test failed: {e}")
return False
+
def main():
"""Run all tests for AIME2024 benchmark."""
print("AIME2024 Benchmark Test Suite")
print("=" * 50)
-
+
try:
# Test dataset creation
dataset = test_dataset_creation()
print()
-
+
# Test metrics calculation
metrics = test_metrics_calculation()
print()
-
+
# Test integration
integration_success = test_integration()
print()
-
+
if integration_success:
print("๐ All tests passed! AIME2024 benchmark is ready to use.")
print()
print("Usage examples:")
print(" cd evaluation")
- print(" python evaluate.py --dataset aime2024 --base_model meta-llama/Llama-3.1-8B-Instruct --device cpu")
+ print(
+ " python evaluate.py --dataset aime2024 --base_model meta-llama/Llama-3.1-8B-Instruct --device cpu"
+ )
print(" python evaluate.py --dataset aime2024 --device 1 --num_samples 10")
else:
print("โ Some tests failed. Please check the implementation.")
sys.exit(1)
-
+
except Exception as e:
print(f"โ Test suite failed with error: {e}")
import traceback
+
traceback.print_exc()
sys.exit(1)
+
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/benchmark/AIME2025/__init__.py b/benchmark/AIME2025/__init__.py
index b5a69b8d..01b7ec84 100644
--- a/benchmark/AIME2025/__init__.py
+++ b/benchmark/AIME2025/__init__.py
@@ -8,4 +8,4 @@
from .calculate_metrics import calculate_metrics
from .aime2025 import AIME2025
-__all__ = ["calculate_metrics", "AIME2025"]
\ No newline at end of file
+__all__ = ["calculate_metrics", "AIME2025"]
diff --git a/benchmark/AIME2025/aime2025.py b/benchmark/AIME2025/aime2025.py
index 5eebe349..02101668 100644
--- a/benchmark/AIME2025/aime2025.py
+++ b/benchmark/AIME2025/aime2025.py
@@ -29,22 +29,23 @@ class AIME2025(Benchmark):
# AIME2025 has a single dataset
all_datasets: List[str] = ["aime2025"]
-
+
benchmark_name: str = "aime2025"
huggingface_dataset_id: str = "xAlg-AI/att-hub-aime2025"
def _load_datasets(self) -> pd.DataFrame:
"""Load AIME2025 dataset.
-
+
AIME2025 uses a single dataset with all problems.
-
+
Returns:
pandas DataFrame with all AIME2025 problems.
"""
print(f"Loading AIME2025 dataset")
-
+
try:
from datasets import load_dataset
+
dataset = load_dataset(self.huggingface_dataset_id, split="test")
df = dataset.to_pandas()
df["task"] = "aime2025" # Ensure task column exists
@@ -72,7 +73,7 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
# Use the calculate_metrics function from HashAttention evaluation
metrics: Dict[str, Any] = calculate_metrics(results_df)
-
+
# Format the results for consistency with other benchmarks
overall_metrics: Dict[str, Any] = {
"overall_score": round(metrics["exact_match"], 4),
@@ -84,17 +85,20 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
"aime2025": {
"exact_match": round(metrics["exact_match"], 4),
"extraction_rate": round(metrics["extraction_rate"], 4),
- "boxed_format_rate": round(metrics["boxed_format_rate"], 4)
+ "boxed_format_rate": round(metrics["boxed_format_rate"], 4),
}
},
- "summary": {
- "total_tasks": 1,
- "total_samples": len(results_df)
- }
+ "summary": {"total_tasks": 1, "total_samples": len(results_df)},
}
-
- print(f" โ AIME2025 Exact Match: {metrics['exact_match']:.3f} ({metrics['exact_match']*100:.1f}%)")
- print(f" โ Extraction Rate: {metrics['extraction_rate']:.3f} ({metrics['extraction_rate']*100:.1f}%)")
- print(f" โ Boxed Format Rate: {metrics['boxed_format_rate']:.3f} ({metrics['boxed_format_rate']*100:.1f}%)")
-
- return overall_metrics
\ No newline at end of file
+
+ print(
+ f" โ AIME2025 Exact Match: {metrics['exact_match']:.3f} ({metrics['exact_match']*100:.1f}%)"
+ )
+ print(
+ f" โ Extraction Rate: {metrics['extraction_rate']:.3f} ({metrics['extraction_rate']*100:.1f}%)"
+ )
+ print(
+ f" โ Boxed Format Rate: {metrics['boxed_format_rate']:.3f} ({metrics['boxed_format_rate']*100:.1f}%)"
+ )
+
+ return overall_metrics
diff --git a/benchmark/AIME2025/calculate_metrics.py b/benchmark/AIME2025/calculate_metrics.py
index 8dcb6217..0c0b1fbd 100644
--- a/benchmark/AIME2025/calculate_metrics.py
+++ b/benchmark/AIME2025/calculate_metrics.py
@@ -10,38 +10,38 @@
def extract_boxed_answer(text: str) -> str:
"""
Extract the answer from \boxed{...} format in the text.
-
+
Args:
text: The model's response text
-
+
Returns:
The extracted answer as a string, or empty string if not found
"""
# Look for \boxed{...} pattern
- boxed_pattern = r'\\boxed\{([^}]*)\}'
+ boxed_pattern = r"\\boxed\{([^}]*)\}"
matches = re.findall(boxed_pattern, text)
-
+
if matches:
# Return the last boxed answer found (in case there are multiple)
return matches[-1].strip()
-
+
# Fallback: look for boxed without backslash (in case the model omits it)
- boxed_pattern_alt = r'boxed\{([^}]*)\}'
+ boxed_pattern_alt = r"boxed\{([^}]*)\}"
matches = re.findall(boxed_pattern_alt, text)
-
+
if matches:
return matches[-1].strip()
-
+
return ""
def extract_numerical_answer(text: str) -> str:
"""
Extract a numerical answer from text, with fallback strategies.
-
+
Args:
text: The text to extract from
-
+
Returns:
The extracted numerical answer as a string
"""
@@ -49,17 +49,17 @@ def extract_numerical_answer(text: str) -> str:
boxed_answer = extract_boxed_answer(text)
if boxed_answer:
# Extract just the number from the boxed content
- numbers = re.findall(r'\d+', boxed_answer)
+ numbers = re.findall(r"\d+", boxed_answer)
if numbers:
return numbers[-1] # Take the last number found
-
+
# Fallback 1: Look for "answer is X" or "answer: X" patterns
answer_patterns = [
- r'(?:answer|solution)\s*(?:is|:)\s*(\d+)',
- r'(?:the\s+)?answer\s*(?:is|:)\s*(\d+)',
- r'(?:therefore|thus|so)\s*(?:the\s+)?(?:answer|solution)\s*(?:is|:)\s*(\d+)',
+ r"(?:answer|solution)\s*(?:is|:)\s*(\d+)",
+ r"(?:the\s+)?answer\s*(?:is|:)\s*(\d+)",
+ r"(?:therefore|thus|so)\s*(?:the\s+)?(?:answer|solution)\s*(?:is|:)\s*(\d+)",
]
-
+
for pattern in answer_patterns:
matches = re.findall(pattern, text, re.IGNORECASE)
if matches:
@@ -67,125 +67,127 @@ def extract_numerical_answer(text: str) -> str:
num = int(matches[-1])
if 0 <= num <= 999:
return matches[-1]
-
+
# Fallback 2: Look for numbers at the end of the text
# This catches cases where the model just states the number
- lines = text.strip().split('\n')
+ lines = text.strip().split("\n")
for line in reversed(lines):
line = line.strip()
if line:
- numbers = re.findall(r'\b\d+\b', line)
+ numbers = re.findall(r"\b\d+\b", line)
if numbers:
# Check if the number is in valid AIME range (0-999)
num = int(numbers[-1])
if 0 <= num <= 999:
return str(num)
-
+
# Fallback 3: Find any number in valid AIME range
- all_numbers = re.findall(r'\b\d+\b', text)
+ all_numbers = re.findall(r"\b\d+\b", text)
for num_str in reversed(all_numbers): # Check from end to beginning
num = int(num_str)
if 0 <= num <= 999:
return str(num)
-
+
return ""
-def calculate_exact_match_score(predictions: List[str], references: List[List[str]]) -> float:
+def calculate_exact_match_score(
+ predictions: List[str], references: List[List[str]]
+) -> float:
"""
Calculate exact match accuracy between predictions and references.
-
+
Args:
predictions: List of predicted answers
references: List of reference answers (each can have multiple valid answers)
-
+
Returns:
Exact match accuracy as a float between 0 and 1
"""
correct = 0
total = len(predictions)
-
+
for pred, ref_list in zip(predictions, references):
# Extract numerical answer from prediction
pred_answer = extract_numerical_answer(pred)
-
+
# Check if prediction matches any of the reference answers
if pred_answer in ref_list:
correct += 1
-
+
return correct / total if total > 0 else 0.0
def calculate_metrics(df: pd.DataFrame) -> dict:
"""
Calculate metrics for the AIME2025 benchmark.
-
+
Args:
df: DataFrame with columns 'predicted_answer' and 'answer'
-
+
Returns:
Dictionary containing the calculated metrics
"""
predictions = df["predicted_answer"].tolist()
references = df["answer"].tolist()
-
+
# Ensure references are in the correct format (list of lists)
if references and not isinstance(references[0], list):
references = [[str(ref)] for ref in references]
-
+
# Calculate exact match accuracy
exact_match = calculate_exact_match_score(predictions, references)
-
+
# Additional analysis: count how many answers were successfully extracted
extracted_count = 0
boxed_count = 0
-
+
for pred in predictions:
extracted = extract_numerical_answer(pred)
if extracted:
extracted_count += 1
-
+
# Count how many used the boxed format
if extract_boxed_answer(pred):
boxed_count += 1
-
+
extraction_rate = extracted_count / len(predictions) if predictions else 0.0
boxed_format_rate = boxed_count / len(predictions) if predictions else 0.0
-
+
return {
"exact_match": exact_match,
"extraction_rate": extraction_rate,
"boxed_format_rate": boxed_format_rate,
- "total_problems": len(predictions)
+ "total_problems": len(predictions),
}
def analyze_errors(df: pd.DataFrame) -> dict:
"""
Analyze common error patterns in the predictions.
-
+
Args:
df: DataFrame with predictions and references
-
+
Returns:
Dictionary with error analysis
"""
predictions = df["predicted_answer"].tolist()
references = df["answer"].tolist()
-
+
if references and not isinstance(references[0], list):
references = [[str(ref)] for ref in references]
-
+
error_types = {
"no_answer_extracted": 0,
"wrong_answer": 0,
"out_of_range": 0,
- "format_issues": 0
+ "format_issues": 0,
}
-
+
for pred, ref_list in zip(predictions, references):
extracted = extract_numerical_answer(pred)
-
+
if not extracted:
error_types["no_answer_extracted"] += 1
elif extracted not in ref_list:
@@ -196,5 +198,5 @@ def analyze_errors(df: pd.DataFrame) -> dict:
error_types["out_of_range"] += 1
except ValueError:
error_types["format_issues"] += 1
-
- return error_types
\ No newline at end of file
+
+ return error_types
diff --git a/benchmark/AIME2025/create_huggingface_dataset.py b/benchmark/AIME2025/create_huggingface_dataset.py
index b71d7d73..27a78fad 100644
--- a/benchmark/AIME2025/create_huggingface_dataset.py
+++ b/benchmark/AIME2025/create_huggingface_dataset.py
@@ -23,6 +23,7 @@
which is standard in mathematical competition contexts.
"""
+
def create_aime2025_dataset():
"""
Process the AIME2025 dataset and convert it to the standardized benchmark format.
@@ -30,10 +31,10 @@ def create_aime2025_dataset():
# Load the original dataset
dataset = load_dataset("yentinglin/aime_2025")
df = dataset["train"].to_pandas()
-
+
# Create the standardized format
processed_data = []
-
+
for _, row in df.iterrows():
# Format the problem with clear instructions about the boxed answer format
context = f"""Solve the following AIME (American Invitational Mathematics Examination) problem.
@@ -43,43 +44,50 @@ def create_aime2025_dataset():
Instructions:
- The answer should be an integer between 0 and 999
- You must wrap your final answer in \\boxed{{...}} format"""
-
+
question = "What is the answer to this problem?"
-
+
# The answer prefix encourages the model to show work before the final answer
answer_prefix = ""
-
+
# Convert answer to list format (some benchmarks expect this)
- answer = [str(row['answer'])]
-
- processed_data.append({
- 'context': context,
- 'question': question,
- 'answer_prefix': answer_prefix,
- 'answer': answer,
- 'task': 'aime2025',
- 'max_new_tokens': 32000, # Allow comprehensive step-by-step solutions
- 'problem_id': row['id'],
- 'year': row['year']
- })
-
+ answer = [str(row["answer"])]
+
+ processed_data.append(
+ {
+ "context": context,
+ "question": question,
+ "answer_prefix": answer_prefix,
+ "answer": answer,
+ "task": "aime2025",
+ "max_new_tokens": 32000, # Allow comprehensive step-by-step solutions
+ "problem_id": row["id"],
+ "year": row["year"],
+ }
+ )
+
# Convert to DataFrame and then to Dataset
processed_df = pd.DataFrame(processed_data)
-
+
# Select only the required columns for the benchmark
- final_df = processed_df[['context', 'question', 'answer_prefix', 'answer', 'task', 'max_new_tokens']]
-
+ final_df = processed_df[
+ ["context", "question", "answer_prefix", "answer", "task", "max_new_tokens"]
+ ]
+
return Dataset.from_pandas(final_df)
+
if __name__ == "__main__":
# Create the processed dataset
processed_dataset = create_aime2025_dataset()
-
+
# Push to hub (you would need to set up your own repo)
# For now, we'll just save locally or use the existing dataset
print(f"Processed {len(processed_dataset)} AIME2025 problems")
print("Sample processed example:")
print(processed_dataset[0])
-
+
# Optionally save locally for testing
- processed_dataset.push_to_hub("xAlg-AI/att-hub-aime2025", config_name=f"aime2025", split="test")
+ processed_dataset.push_to_hub(
+ "xAlg-AI/att-hub-aime2025", config_name=f"aime2025", split="test"
+ )
diff --git a/benchmark/AIME2025/example_usage.py b/benchmark/AIME2025/example_usage.py
index 5fe3d38e..952faf3f 100644
--- a/benchmark/AIME2025/example_usage.py
+++ b/benchmark/AIME2025/example_usage.py
@@ -8,6 +8,7 @@
import sys
import os
+
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from AIME2025.create_huggingface_dataset import create_aime2025_dataset
@@ -19,21 +20,21 @@ def show_sample_problems():
"""Display sample problems from the AIME2025 dataset."""
print("AIME2025 Sample Problems")
print("=" * 50)
-
+
dataset = create_aime2025_dataset()
-
+
# Show first 3 problems
for i in range(min(3, len(dataset))):
example = dataset[i]
print(f"\nProblem {i+1}:")
print("-" * 30)
-
+
# Extract just the problem statement from the context
- context = example['context']
+ context = example["context"]
problem_start = context.find("Problem: ") + len("Problem: ")
problem_end = context.find("\n\nInstructions:")
problem = context[problem_start:problem_end]
-
+
print(f"Problem: {problem}")
print(f"Answer: {example['answer'][0]}")
@@ -43,50 +44,51 @@ def simulate_model_responses():
print("\n" + "=" * 50)
print("Simulated Model Evaluation")
print("=" * 50)
-
+
# Create some simulated model responses
simulated_responses = [
{
- 'predicted_answer': "Let me work through this step by step.\n\nAfter calculating the bases, I find that the sum is \\boxed{70}.",
- 'answer': ['70'] # Correct
+ "predicted_answer": "Let me work through this step by step.\n\nAfter calculating the bases, I find that the sum is \\boxed{70}.",
+ "answer": ["70"], # Correct
},
{
- 'predicted_answer': "This is a complex problem. After working through it, the answer is \\boxed{42}.",
- 'answer': ['588'] # Wrong
+ "predicted_answer": "This is a complex problem. After working through it, the answer is \\boxed{42}.",
+ "answer": ["588"], # Wrong
},
{
- 'predicted_answer': "The calculation gives us 588 as the final result.",
- 'answer': ['588'] # Correct, but no boxed format
+ "predicted_answer": "The calculation gives us 588 as the final result.",
+ "answer": ["588"], # Correct, but no boxed format
},
{
- 'predicted_answer': "Answer: 16",
- 'answer': ['16'] # Correct, no boxed format
+ "predicted_answer": "Answer: 16",
+ "answer": ["16"], # Correct, no boxed format
},
{
- 'predicted_answer': "This problem is too difficult to solve.",
- 'answer': ['100'] # No answer extracted
- }
+ "predicted_answer": "This problem is too difficult to solve.",
+ "answer": ["100"], # No answer extracted
+ },
]
-
+
# Create DataFrame
df = pd.DataFrame(simulated_responses)
-
+
# Calculate metrics
metrics = calculate_metrics(df)
-
+
print("Simulated Results:")
print(f"Total problems: {len(simulated_responses)}")
print(f"Exact match accuracy: {metrics['exact_match']:.1%}")
print(f"Answer extraction rate: {metrics['extraction_rate']:.1%}")
print(f"Boxed format usage: {metrics['boxed_format_rate']:.1%}")
-
+
print("\nDetailed breakdown:")
for i, response in enumerate(simulated_responses):
from AIME2025.calculate_metrics import extract_numerical_answer
- extracted = extract_numerical_answer(response['predicted_answer'])
- correct = extracted in response['answer']
+
+ extracted = extract_numerical_answer(response["predicted_answer"])
+ correct = extracted in response["answer"]
status = "โ Correct" if correct else "โ Wrong"
-
+
print(f"Problem {i+1}: {status}")
print(f" Expected: {response['answer'][0]}")
print(f" Extracted: '{extracted}'")
@@ -99,12 +101,15 @@ def show_usage_instructions():
print("=" * 50)
print("Usage Instructions")
print("=" * 50)
-
+
print("\n1. Command Line Usage:")
- print(" python evaluation/evaluate.py --dataset aime2025 --base_model meta-llama/Llama-3.1-8B-Instruct")
-
+ print(
+ " python evaluation/evaluate.py --dataset aime2025 --base_model meta-llama/Llama-3.1-8B-Instruct"
+ )
+
print("\n2. Programmatic Usage:")
- print("""
+ print(
+ """
from evaluation.evaluate import evaluate
# Run evaluation
@@ -114,10 +119,12 @@ def show_usage_instructions():
max_new_tokens=4096,
num_samples=10 # Optional: evaluate on subset
)
- """)
-
+ """
+ )
+
print("\n3. Custom Evaluation:")
- print("""
+ print(
+ """
from evaluation.AIME2025 import create_aime2025_dataset, calculate_metrics
import pandas as pd
@@ -134,10 +141,12 @@ def show_usage_instructions():
})
metrics = calculate_metrics(results_df)
print(metrics)
- """)
-
+ """
+ )
+
print("\n4. Expected Output Format:")
- print("""
+ print(
+ """
Models should respond with step-by-step reasoning and wrap
the final answer in \\boxed{...} format:
@@ -150,7 +159,8 @@ def show_usage_instructions():
[mathematical reasoning...]
Therefore, the sum of all valid bases is \\boxed{70}."
- """)
+ """
+ )
def main():
@@ -158,11 +168,11 @@ def main():
show_sample_problems()
simulate_model_responses()
show_usage_instructions()
-
+
print("\n" + "=" * 50)
print("AIME2025 Benchmark Ready!")
print("=" * 50)
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/benchmark/AIME2025/test_aime2025.py b/benchmark/AIME2025/test_aime2025.py
index 0a9ae638..44571803 100644
--- a/benchmark/AIME2025/test_aime2025.py
+++ b/benchmark/AIME2025/test_aime2025.py
@@ -8,11 +8,16 @@
import sys
import os
+
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import pandas as pd
from AIME2025.create_huggingface_dataset import create_aime2025_dataset
-from AIME2025.calculate_metrics import calculate_metrics, extract_numerical_answer, analyze_errors
+from AIME2025.calculate_metrics import (
+ calculate_metrics,
+ extract_numerical_answer,
+ analyze_errors,
+)
def test_dataset_creation():
@@ -20,20 +25,27 @@ def test_dataset_creation():
print("=" * 60)
print("Testing AIME2025 Dataset Creation")
print("=" * 60)
-
+
dataset = create_aime2025_dataset()
print(f"โ Successfully created dataset with {len(dataset)} examples")
-
+
# Check the structure
sample = dataset[0]
- required_keys = ['context', 'question', 'answer_prefix', 'answer', 'task', 'max_new_tokens']
-
+ required_keys = [
+ "context",
+ "question",
+ "answer_prefix",
+ "answer",
+ "task",
+ "max_new_tokens",
+ ]
+
for key in required_keys:
if key in sample:
print(f"โ Required key '{key}' present")
else:
print(f"โ Required key '{key}' missing")
-
+
# Show a sample
print("\nSample example:")
print(f"Context (first 200 chars): {sample['context'][:200]}...")
@@ -41,7 +53,7 @@ def test_dataset_creation():
print(f"Answer: {sample['answer']}")
print(f"Task: {sample['task']}")
print(f"Max new tokens: {sample['max_new_tokens']}")
-
+
return dataset
@@ -50,7 +62,7 @@ def test_answer_extraction():
print("\n" + "=" * 60)
print("Testing Answer Extraction")
print("=" * 60)
-
+
test_cases = [
("\\boxed{42}", "42", "Standard boxed format"),
("The answer is \\boxed{123}", "123", "Boxed with prefix"),
@@ -63,7 +75,7 @@ def test_answer_extraction():
("The answer is 1000", "", "Out of range (should be empty)"),
("Multiple numbers: 42, 123, 456", "456", "Multiple numbers (last valid one)"),
]
-
+
for input_text, expected, description in test_cases:
extracted = extract_numerical_answer(input_text)
status = "โ" if extracted == expected else "โ"
@@ -80,48 +92,54 @@ def test_metrics_calculation():
print("=" * 60)
print("Testing Metrics Calculation")
print("=" * 60)
-
+
# Create test data with known outcomes
test_data = {
- 'predicted_answer': [
+ "predicted_answer": [
"Let me solve this step by step.\n\nThe answer is \\boxed{70}", # Correct, boxed
- "After calculation, I get \\boxed{42}", # Wrong answer, boxed
- "The solution is 123", # Correct, no boxed
- "Answer: 999", # Correct, no boxed
- "This problem is too complex to solve", # No answer
- "The result is \\boxed{500}", # Correct, boxed
+ "After calculation, I get \\boxed{42}", # Wrong answer, boxed
+ "The solution is 123", # Correct, no boxed
+ "Answer: 999", # Correct, no boxed
+ "This problem is too complex to solve", # No answer
+ "The result is \\boxed{500}", # Correct, boxed
+ ],
+ "answer": [
+ ["70"], # Match
+ ["123"], # No match (predicted 42)
+ ["123"], # Match
+ ["999"], # Match
+ ["100"], # No match (no prediction)
+ ["500"], # Match
],
- 'answer': [
- ['70'], # Match
- ['123'], # No match (predicted 42)
- ['123'], # Match
- ['999'], # Match
- ['100'], # No match (no prediction)
- ['500'], # Match
- ]
}
-
+
test_df = pd.DataFrame(test_data)
metrics = calculate_metrics(test_df)
-
+
print("Calculated metrics:")
for key, value in metrics.items():
print(f" {key}: {value}")
-
+
# Expected results:
# - exact_match: 4/6 = 0.667 (problems 1, 3, 4, 6 correct)
# - extraction_rate: 5/6 = 0.833 (all except problem 5)
# - boxed_format_rate: 3/6 = 0.5 (problems 1, 2, 6)
-
- expected_exact_match = 4/6
- expected_extraction_rate = 5/6
- expected_boxed_rate = 3/6
-
+
+ expected_exact_match = 4 / 6
+ expected_extraction_rate = 5 / 6
+ expected_boxed_rate = 3 / 6
+
print(f"\nValidation:")
- print(f"โ Exact match: {metrics['exact_match']:.3f} (expected: {expected_exact_match:.3f})")
- print(f"โ Extraction rate: {metrics['extraction_rate']:.3f} (expected: {expected_extraction_rate:.3f})")
- print(f"โ Boxed format rate: {metrics['boxed_format_rate']:.3f} (expected: {expected_boxed_rate:.3f})")
-
+ print(
+ f"โ Exact match: {metrics['exact_match']:.3f} (expected: {expected_exact_match:.3f})"
+ )
+ print(
+ f"โ Extraction rate: {metrics['extraction_rate']:.3f} (expected: {expected_extraction_rate:.3f})"
+ )
+ print(
+ f"โ Boxed format rate: {metrics['boxed_format_rate']:.3f} (expected: {expected_boxed_rate:.3f})"
+ )
+
# Test error analysis
errors = analyze_errors(test_df)
print(f"\nError analysis:")
@@ -134,18 +152,18 @@ def test_integration():
print("\n" + "=" * 60)
print("Testing Integration")
print("=" * 60)
-
+
try:
# Test that the benchmark can be imported from the main evaluation module
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
+
# This would normally import from evaluate.py, but we'll skip due to dependencies
print("โ Integration test skipped (requires full environment)")
print(" The benchmark has been properly integrated into evaluate.py")
print(" - Added to DATASET_DICT as 'aime2025': 'yentinglin/aime_2025'")
print(" - Added to SCORER_DICT as 'aime2025': aime2025_scorer")
print(" - Added special handling for dataset processing")
-
+
except Exception as e:
print(f"โ Integration test failed: {e}")
@@ -154,20 +172,22 @@ def main():
"""Run all tests."""
print("AIME2025 Benchmark Test Suite")
print("=" * 60)
-
+
# Run all tests
dataset = test_dataset_creation()
test_answer_extraction()
test_metrics_calculation()
test_integration()
-
+
print("\n" + "=" * 60)
print("Test Suite Complete")
print("=" * 60)
print("The AIME2025 benchmark is ready for use!")
print("\nTo run the benchmark:")
- print(" python evaluation/evaluate.py --dataset aime2025 --base_model ")
+ print(
+ " python evaluation/evaluate.py --dataset aime2025 --base_model "
+ )
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/benchmark/base.py b/benchmark/base.py
index 2a7c3a87..bb4c90d9 100644
--- a/benchmark/base.py
+++ b/benchmark/base.py
@@ -40,10 +40,10 @@ class Benchmark(ABC):
... all_datasets = ["task1", "task2"]
... benchmark_name = "my_benchmark"
... huggingface_dataset_id = "my_org/my_dataset"
- ...
+ ...
... def post_run_evaluate(self, results_df):
... return {"accuracy": 0.95}
- >>>
+ >>>
>>> benchmark = MyBenchmark(subsets_to_run=["task1"])
>>> results = benchmark.run_benchmark(adapter, result_dir="/path/to/results", generation_kwargs={"max_new_tokens": 50}, request_kwargs={"max_context_length": 1024})
"""
@@ -63,11 +63,17 @@ def __init__(self, subsets_to_run: Optional[List[str]] = None) -> None:
ValueError: If any subset in subsets_to_run is not in all_datasets.
"""
if not self.all_datasets:
- raise ValueError(f"Subclass {self.__class__.__name__} must define all_datasets")
+ raise ValueError(
+ f"Subclass {self.__class__.__name__} must define all_datasets"
+ )
if not self.benchmark_name:
- raise ValueError(f"Subclass {self.__class__.__name__} must define benchmark_name")
+ raise ValueError(
+ f"Subclass {self.__class__.__name__} must define benchmark_name"
+ )
if not self.huggingface_dataset_id:
- raise ValueError(f"Subclass {self.__class__.__name__} must define huggingface_dataset_id")
+ raise ValueError(
+ f"Subclass {self.__class__.__name__} must define huggingface_dataset_id"
+ )
if subsets_to_run is None:
self.subsets_to_run = self.all_datasets.copy()
@@ -112,7 +118,7 @@ def _load_datasets(self) -> pd.DataFrame:
# Load the full dataset
dataset = load_dataset(self.huggingface_dataset_id, split="test")
df: pd.DataFrame = dataset.to_pandas()
-
+
# Filter to only include subsets we want to run
if "task" in df.columns:
# Filter by task column if it exists
@@ -121,11 +127,13 @@ def _load_datasets(self) -> pd.DataFrame:
# If no task column, assume the dataset only contains our subsets
# This is a simplified assumption based on user guidance
pass
-
+
return df
-
+
except Exception as e:
- raise Exception(f"Failed to load dataset {self.huggingface_dataset_id}: {str(e)}")
+ raise Exception(
+ f"Failed to load dataset {self.huggingface_dataset_id}: {str(e)}"
+ )
def _validate_dataset_size(self, df: pd.DataFrame) -> None:
"""Validate dataset size and warn if too large.
@@ -137,15 +145,15 @@ def _validate_dataset_size(self, df: pd.DataFrame) -> None:
warnings.warn(
f"Dataset has {len(df)} rows (>10K). Repository not expected to handle "
"large datasets. If needed, request this feature.",
- UserWarning
+ UserWarning,
)
def _process_all_requests(
- self,
- adapter: ModelAdapter,
+ self,
+ adapter: ModelAdapter,
dataset_df: pd.DataFrame,
generation_kwargs: Dict[str, Any],
- request_kwargs: Dict[str, Any]
+ request_kwargs: Dict[str, Any],
) -> pd.DataFrame:
"""Process all samples through the model adapter using context grouping for efficiency.
@@ -164,47 +172,60 @@ def _process_all_requests(
# Truncate dataset to max_requests
dataset_df = dataset_df.head(max_requests)
-
-
+
# Group by context for efficiency (following HashAttention approach)
df_context = dataset_df.groupby("context")
-
- for context, df_group in tqdm(df_context, desc="Processing contexts", total=dataset_df["context"].nunique()):
+
+ for context, df_group in tqdm(
+ df_context,
+ desc="Processing contexts",
+ total=dataset_df["context"].nunique(),
+ ):
questions: List[str] = df_group["question"].to_list()
-
+
try:
# Create request using current adapter interface (simplified)
answer_prefix = df_group["answer_prefix"].iloc[0]
- request: Request = Request(context=context, questions=questions, answer_prefix=answer_prefix)
-
+ request: Request = Request(
+ context=context, questions=questions, answer_prefix=answer_prefix
+ )
+
# Process through adapter
- response: RequestResponse = adapter.process_request(request, generation_kwargs, request_kwargs)
-
+ response: RequestResponse = adapter.process_request(
+ request, generation_kwargs, request_kwargs
+ )
+
# Assign responses back to DataFrame
if isinstance(response.responses, list):
- dataset_df.loc[df_group.index, "predicted_answer"] = response.responses
+ dataset_df.loc[df_group.index, "predicted_answer"] = (
+ response.responses
+ )
else:
# Single response case
- dataset_df.loc[df_group.index, "predicted_answer"] = [response.responses] * len(df_group)
-
+ dataset_df.loc[df_group.index, "predicted_answer"] = [
+ response.responses
+ ] * len(df_group)
+
# Memory cleanup for large contexts
if torch.cuda.is_available():
torch.cuda.empty_cache()
-
+
except Exception as e:
# Log error but continue processing other contexts
print(f"Error processing context (length {len(context)}): {str(e)}")
# Fill with empty responses for failed contexts
- dataset_df.loc[df_group.index, "predicted_answer"] = [""] * len(df_group)
-
+ dataset_df.loc[df_group.index, "predicted_answer"] = [""] * len(
+ df_group
+ )
+
return dataset_df
def run_benchmark(
- self,
- adapter: ModelAdapter,
+ self,
+ adapter: ModelAdapter,
result_dir: str,
generation_kwargs: Optional[Dict[str, Any]] = None,
- request_kwargs: Optional[Dict[str, Any]] = None
+ request_kwargs: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""Main orchestration method for running complete benchmark.
@@ -222,54 +243,60 @@ def run_benchmark(
generation_kwargs = {}
if request_kwargs is None:
request_kwargs = {}
-
+
# Create result directory if it doesn't exist
result_path: Path = Path(result_dir)
result_path.mkdir(parents=True, exist_ok=True)
-
+
# Load datasets
print(f"Loading {self.benchmark_name} datasets: {self.subsets_to_run}")
dataset_df: pd.DataFrame = self._load_datasets()
print(f"Loaded {len(dataset_df)} samples")
# Validate dataset size
self._validate_dataset_size(dataset_df)
-
+
# Process all requests through the adapter
print("Processing requests through adapter...")
- results_df: pd.DataFrame = self._process_all_requests(adapter, dataset_df, generation_kwargs, request_kwargs)
-
+ results_df: pd.DataFrame = self._process_all_requests(
+ adapter, dataset_df, generation_kwargs, request_kwargs
+ )
+
# Compute evaluation metrics
print("Computing evaluation metrics...")
metrics: Dict[str, Any] = self.post_run_evaluate(results_df)
-
+
# Save results
raw_results_path: Path = result_path / "raw_results.csv"
save_dataframe_to_csv(results_df, str(raw_results_path), index=False)
print(f"Saved raw results to {raw_results_path}")
-
+
# Save metrics
metrics_path: Path = result_path / "metrics.json"
with open(metrics_path, "w") as f:
json.dump(metrics, f, indent=2)
print(f"Saved metrics to {metrics_path}")
-
+
# Save configuration parameters
config_path: Path = result_path / "config.json"
-
+
config_data = {
- "model_kwargs": make_serializable(getattr(adapter, 'model_kwargs', {})),
- "tokenizer_kwargs": make_serializable(getattr(adapter, 'tokenizer_kwargs', {})),
- "sparse_attention_config": make_serializable(getattr(adapter, 'sparse_attention_config', None)),
+ "model_kwargs": make_serializable(getattr(adapter, "model_kwargs", {})),
+ "tokenizer_kwargs": make_serializable(
+ getattr(adapter, "tokenizer_kwargs", {})
+ ),
+ "sparse_attention_config": make_serializable(
+ getattr(adapter, "sparse_attention_config", None)
+ ),
"generation_kwargs": make_serializable(generation_kwargs),
"request_kwargs": make_serializable(request_kwargs),
"benchmark_name": self.benchmark_name,
"subsets_to_run": self.subsets_to_run,
- "huggingface_dataset_id": self.huggingface_dataset_id
+ "huggingface_dataset_id": self.huggingface_dataset_id,
}
with open(config_path, "w") as f:
json.dump(config_data, f, indent=2)
print(f"Saved configuration to {config_path}")
-
+
return metrics
@abstractmethod
@@ -282,7 +309,7 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
Args:
results_df: DataFrame containing input data and model outputs with columns:
- context: The input context
- - question: The input question
+ - question: The input question
- predicted_answer: Model's predicted answer
- Plus any other columns from the original dataset
@@ -290,6 +317,3 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
Dictionary containing computed metrics (e.g., {"accuracy": 0.95, "f1": 0.88}).
"""
pass
-
-
-
diff --git a/benchmark/benchmark_registry.py b/benchmark/benchmark_registry.py
index bbc5bdf8..47d19d41 100644
--- a/benchmark/benchmark_registry.py
+++ b/benchmark/benchmark_registry.py
@@ -18,45 +18,49 @@
def register_benchmark(name: Optional[str] = None, aliases: Optional[List[str]] = None):
"""Decorator to automatically register benchmark classes.
-
+
This decorator allows benchmark classes to register themselves automatically,
eliminating the need for manual registry maintenance.
-
+
Args:
name: Custom name for the benchmark (defaults to class.benchmark_name)
aliases: Optional list of alternative names for the benchmark
-
+
Example:
>>> @register_benchmark("my_benchmark", aliases=["my_bench", "mb"])
>>> class MyBenchmark(Benchmark):
... benchmark_name = "my_benchmark"
... # ... rest of implementation
"""
+
def decorator(benchmark_class: type) -> type:
# Get benchmark name from parameter, class attribute, or class name
if name is not None:
benchmark_name = name
- elif hasattr(benchmark_class, 'benchmark_name') and benchmark_class.benchmark_name:
+ elif (
+ hasattr(benchmark_class, "benchmark_name")
+ and benchmark_class.benchmark_name
+ ):
benchmark_name = benchmark_class.benchmark_name
else:
benchmark_name = benchmark_class.__name__.lower()
-
+
# Register main name
_BENCHMARK_REGISTRY[benchmark_name.lower()] = benchmark_class
-
+
# Register aliases
if aliases:
for alias in aliases:
_BENCHMARK_REGISTRY[alias.lower()] = benchmark_class
-
+
return benchmark_class
-
+
return decorator
def get_registered_benchmarks() -> Dict[str, type]:
"""Get all registered benchmark classes.
-
+
Returns:
Dictionary mapping benchmark names to their classes
"""
@@ -65,7 +69,7 @@ def get_registered_benchmarks() -> Dict[str, type]:
def get_available_benchmark_names() -> List[str]:
"""Get list of all available benchmark names.
-
+
Returns:
Sorted list of registered benchmark names
"""
@@ -73,61 +77,65 @@ def get_available_benchmark_names() -> List[str]:
def create_benchmark_instance(
- benchmark_name: str,
- subsets: Optional[List[str]] = None
+ benchmark_name: str, subsets: Optional[List[str]] = None
) -> Benchmark:
"""Factory function to create benchmark instances by name.
-
+
Uses the automatic registry to instantiate benchmark classes.
No manual maintenance required - benchmarks register themselves.
-
+
Args:
benchmark_name: Name of the benchmark to create
subsets: Optional list of subsets to run for the benchmark
-
+
Returns:
Instantiated benchmark object
-
+
Raises:
ValueError: If benchmark_name is not registered
-
+
Example:
>>> benchmark = create_benchmark_instance("longbench", ["narrativeqa"])
>>> benchmark = create_benchmark_instance("mock_benchmark")
"""
# Normalize benchmark name
benchmark_name = benchmark_name.strip().lower()
-
+
if benchmark_name not in _BENCHMARK_REGISTRY:
available_benchmarks: List[str] = get_available_benchmark_names()
raise ValueError(
f"Unknown benchmark '{benchmark_name}'. Available benchmarks: {available_benchmarks}"
)
-
+
benchmark_class = _BENCHMARK_REGISTRY[benchmark_name]
-
+
try:
# Create instance with optional subsets
benchmark_instance: Benchmark = benchmark_class(subsets_to_run=subsets)
return benchmark_instance
-
+
except Exception as e:
raise RuntimeError(f"Failed to instantiate benchmark '{benchmark_name}': {e}")
def ensure_benchmarks_loaded() -> None:
"""Ensure all benchmark modules are loaded to trigger registration.
-
+
This function imports all benchmark modules to ensure their decorators
have been executed and they're registered in the global registry.
"""
import benchmark
-
+
# Get all modules in the benchmark package
benchmark_package_path = benchmark.__path__
-
+
for importer, module_name, ispkg in pkgutil.iter_modules(benchmark_package_path):
- if not ispkg and module_name != '__init__' and module_name != 'base' and module_name != 'benchmark_registry':
+ if (
+ not ispkg
+ and module_name != "__init__"
+ and module_name != "base"
+ and module_name != "benchmark_registry"
+ ):
try:
# Import the module to trigger @register_benchmark decorators
module_full_name = f"benchmark.{module_name}"
@@ -139,27 +147,29 @@ def ensure_benchmarks_loaded() -> None:
continue
-def validate_benchmark_config(benchmark_name: str, subsets: Optional[List[str]] = None) -> None:
+def validate_benchmark_config(
+ benchmark_name: str, subsets: Optional[List[str]] = None
+) -> None:
"""Validate a benchmark configuration against available benchmarks.
-
+
This function checks if the benchmark exists and if specified subsets are valid.
-
+
Args:
benchmark_name: Name of the benchmark to validate
subsets: Optional list of subsets to validate
-
+
Raises:
ValueError: If benchmark or subsets are invalid
-
+
Example:
>>> validate_benchmark_config("longbench", ["narrativeqa"]) # Passes validation
- >>>
+ >>>
>>> validate_benchmark_config("invalid_benchmark") # Raises ValueError
"""
try:
# Create temporary benchmark instance to validate
benchmark = create_benchmark_instance(benchmark_name, subsets=None)
-
+
# Validate subsets if specified
if subsets is not None:
available_datasets: List[str] = benchmark.get_available_datasets()
@@ -169,23 +179,23 @@ def validate_benchmark_config(benchmark_name: str, subsets: Optional[List[str]]
f"Invalid subsets for {benchmark_name}: {invalid_subsets}. "
f"Available datasets: {available_datasets}"
)
-
+
except Exception as e:
raise ValueError(f"Invalid benchmark configuration: {e}")
def get_benchmark_subsets(benchmark_name: str) -> Optional[List[str]]:
"""Get available subsets for a benchmark using the standard interface.
-
+
This is a utility function that can help users determine what subsets are
available for a given benchmark using the standard get_available_datasets() method.
-
+
Args:
benchmark_name: Name of the benchmark to inspect
-
+
Returns:
List of available subset names, or None if benchmark doesn't support subsets
-
+
Example:
>>> subsets = get_benchmark_subsets("longbench")
>>> print(subsets) # ['narrativeqa', 'qasper', 'multifieldqa_en', ...]
@@ -193,8 +203,8 @@ def get_benchmark_subsets(benchmark_name: str) -> Optional[List[str]]:
try:
benchmark_instance = create_benchmark_instance(benchmark_name)
return benchmark_instance.get_available_datasets()
-
+
except Exception as e:
# If we can't instantiate the benchmark or get subset info, return None
logging.warning(f"Could not get subsets for benchmark '{benchmark_name}': {e}")
- return None
\ No newline at end of file
+ return None
diff --git a/benchmark/executor.py b/benchmark/executor.py
index 49ce54c5..445093a7 100644
--- a/benchmark/executor.py
+++ b/benchmark/executor.py
@@ -19,8 +19,8 @@
from sparse_attention_hub.metric_logging.logger import MicroMetricLogger
# Set multiprocessing start method to 'spawn' for CUDA compatibility
-if multiprocessing.get_start_method(allow_none=True) != 'spawn':
- multiprocessing.set_start_method('spawn', force=True)
+if multiprocessing.get_start_method(allow_none=True) != "spawn":
+ multiprocessing.set_start_method("spawn", force=True)
from sparse_attention_hub.sparse_attention.base import SparseAttentionConfig
@@ -34,12 +34,8 @@
BenchmarkFailure,
filter_existing_results,
generate_benchmark_stubs,
-
-)
-from .benchmark_registry import (
- ensure_benchmarks_loaded,
- create_benchmark_instance
)
+from .benchmark_registry import ensure_benchmarks_loaded, create_benchmark_instance
from .utils.gpu import (
cleanup_gpu_memory,
@@ -56,14 +52,15 @@
@contextmanager
def set_gpu(gpu_id: int):
"""Context manager to set the current CUDA device.
-
+
Args:
gpu_id: The GPU device ID to set as current device
-
+
Yields:
None: The context manager yields control to the with block
"""
import torch
+
torch.cuda.set_device(gpu_id)
yield
@@ -72,22 +69,24 @@ def _signal_handler(signum: int, frame: Any) -> None:
"""Handle shutdown signals gracefully."""
global _shutdown_requested
_shutdown_requested = True
- logging.getLogger(__name__).warning(f"Received signal {signum}, initiating graceful shutdown...")
+ logging.getLogger(__name__).warning(
+ f"Received signal {signum}, initiating graceful shutdown..."
+ )
def _categorize_error(error: Exception, context: str = "") -> str:
"""Categorize errors for better handling and reporting.
-
+
Args:
error: The exception that occurred
context: Additional context about where the error occurred
-
+
Returns:
Error category string for classification
"""
error_type = type(error).__name__
error_msg = str(error).lower()
-
+
# GPU-related errors
if any(keyword in error_msg for keyword in ["cuda", "gpu", "out of memory", "oom"]):
if "out of memory" in error_msg or "oom" in error_msg:
@@ -96,32 +95,41 @@ def _categorize_error(error: Exception, context: str = "") -> str:
return "GPU_CUDA_ERROR"
else:
return "GPU_ERROR"
-
+
# Model loading errors
- elif any(keyword in error_msg for keyword in ["model", "huggingface", "transformers", "pretrained"]):
+ elif any(
+ keyword in error_msg
+ for keyword in ["model", "huggingface", "transformers", "pretrained"]
+ ):
if "not found" in error_msg or "doesn't exist" in error_msg:
return "MODEL_NOT_FOUND"
elif "tokenizer" in error_msg:
return "TOKENIZER_ERROR"
else:
return "MODEL_LOADING_ERROR"
-
+
# Dataset/benchmark errors
elif any(keyword in error_msg for keyword in ["dataset", "benchmark", "subset"]):
return "DATASET_ERROR"
-
+
# Network/connection errors
- elif any(keyword in error_msg for keyword in ["connection", "timeout", "network", "http"]):
+ elif any(
+ keyword in error_msg for keyword in ["connection", "timeout", "network", "http"]
+ ):
return "NETWORK_ERROR"
-
+
# File system errors
- elif any(keyword in error_msg for keyword in ["file", "directory", "permission", "disk"]):
+ elif any(
+ keyword in error_msg for keyword in ["file", "directory", "permission", "disk"]
+ ):
return "FILESYSTEM_ERROR"
-
+
# Multiprocessing errors
- elif any(keyword in error_msg for keyword in ["queue", "process", "multiprocessing"]):
+ elif any(
+ keyword in error_msg for keyword in ["queue", "process", "multiprocessing"]
+ ):
return "MULTIPROCESSING_ERROR"
-
+
# Default categories
elif isinstance(error, TimeoutError):
return "TIMEOUT_ERROR"
@@ -140,17 +148,17 @@ def _benchmark_worker(
error_queue: multiprocessing.Queue,
timeout_per_benchmark: float = 3600.0,
micro_metric_logger_max_records: Optional[int] = None,
- micro_metric_logger_sampling_factor: float = 1.0
+ micro_metric_logger_sampling_factor: float = 1.0,
) -> None:
"""Worker function to execute benchmarks with dynamic GPU allocation.
-
+
This function runs in a separate process and handles:
1. Acquiring GPUs from the shared pool
2. Creating model adapters with sparse attention configurations
3. Executing benchmarks with proper error handling
4. Releasing GPUs back to the pool
5. Reporting results and errors
-
+
Args:
stub_queue: Queue containing BenchmarkStub instances to process
gpu_pool: Shared queue containing available GPU IDs
@@ -162,11 +170,11 @@ def _benchmark_worker(
"""
worker_id = os.getpid()
logger = logging.getLogger(f"{__name__}.worker_{worker_id}")
-
+
# Set up signal handlers for graceful shutdown
signal.signal(signal.SIGTERM, _signal_handler)
signal.signal(signal.SIGINT, _signal_handler)
-
+
# Track resources for cleanup
current_gpu_id: Optional[int] = None
adapter = None
@@ -177,7 +185,7 @@ def _benchmark_worker(
if _shutdown_requested:
logger.info(f"Worker {worker_id}: Shutdown requested, exiting")
break
-
+
# Step 1: Get next benchmark stub from queue
try:
stub: BenchmarkStub = stub_queue.get(timeout=1.0) # 1 second timeout
@@ -185,127 +193,154 @@ def _benchmark_worker(
# No more work to do
logger.info(f"Worker {worker_id}: No more work, shutting down")
break
-
+
# Check for sentinel value (shutdown signal)
if stub is None:
logger.info(f"Worker {worker_id}: Received shutdown signal, exiting")
break
-
+
# Step 2: Acquire GPU from pool
try:
current_gpu_id = acquire_gpu_from_pool(gpu_pool, timeout=30.0)
- logger.info(f"Worker {worker_id}: Acquired GPU {current_gpu_id} for {stub.model_name}/{stub.sparse_config_name}/{stub.benchmark_name}")
+ logger.info(
+ f"Worker {worker_id}: Acquired GPU {current_gpu_id} for {stub.model_name}/{stub.sparse_config_name}/{stub.benchmark_name}"
+ )
except Exception as e:
error_category = _categorize_error(e, "GPU acquisition")
error_msg = f"Failed to acquire GPU for {stub.model_name}/{stub.sparse_config_name}/{stub.benchmark_name}: {e}"
logger.error(f"Worker {worker_id}: {error_msg}")
- error_queue.put(BenchmarkFailure(
- stub=stub,
- error_message=error_msg,
- error_type=error_category,
- execution_time=0.0
- ))
+ error_queue.put(
+ BenchmarkFailure(
+ stub=stub,
+ error_message=error_msg,
+ error_type=error_category,
+ execution_time=0.0,
+ )
+ )
continue
# Step 3: Run the benchmark on the current GPU
start_time = time.time()
execution_time = 0.0 # Initialize execution_time
-
+
try:
with set_gpu(current_gpu_id):
- logger.debug(f"Worker {worker_id}: Set CUDA device to {current_gpu_id}")
+ logger.debug(
+ f"Worker {worker_id}: Set CUDA device to {current_gpu_id}"
+ )
execution_success = False
# Create model adapter with sparse attention config
- logger.info(f"Worker {worker_id}: Creating model adapter for {stub.model_name}")
-
+ logger.info(
+ f"Worker {worker_id}: Creating model adapter for {stub.model_name}"
+ )
+
# Import here to avoid issues with multiprocessing
from sparse_attention_hub.adapters.huggingface import ModelAdapterHF
-
+
adapter = ModelAdapterHF(
model_name=stub.model_name,
sparse_attention_config=stub.sparse_attention_config,
model_kwargs=stub.adapter_config.model_kwargs,
- tokenizer_kwargs=stub.adapter_config.tokenizer_kwargs
+ tokenizer_kwargs=stub.adapter_config.tokenizer_kwargs,
)
-
+
# Create benchmark instance
- logger.info(f"Worker {worker_id}: Creating benchmark instance for {stub.benchmark_name}")
+ logger.info(
+ f"Worker {worker_id}: Creating benchmark instance for {stub.benchmark_name}"
+ )
benchmark = create_benchmark_instance(
benchmark_name=stub.benchmark_name,
- subsets=[stub.subset] if stub.subset else None
+ subsets=[stub.subset] if stub.subset else None,
)
-
+
# Execute benchmark
- logger.info(f"Worker {worker_id}: Executing benchmark {stub.benchmark_name} on GPU {current_gpu_id}")
+ logger.info(
+ f"Worker {worker_id}: Executing benchmark {stub.benchmark_name} on GPU {current_gpu_id}"
+ )
metric_logger = MicroMetricLogger()
metric_logger.configure_logging(
- log_path=stub.result_dir,
- enabled_metrics=["research_attention_density", "research_attention_output_error"],
+ log_path=stub.result_dir,
+ enabled_metrics=[
+ "research_attention_density",
+ "research_attention_output_error",
+ ],
max_records=micro_metric_logger_max_records,
- sampling_factor=micro_metric_logger_sampling_factor
+ sampling_factor=micro_metric_logger_sampling_factor,
)
metrics = benchmark.run_benchmark(
adapter=adapter,
result_dir=stub.result_dir,
generation_kwargs=stub.generation_kwargs,
- request_kwargs=stub.request_kwargs
+ request_kwargs=stub.request_kwargs,
)
metric_logger.flush()
-
+
execution_time = time.time() - start_time
execution_success = True
-
+
# Report successful result
result = BenchmarkResult(
stub=stub,
metrics=metrics,
execution_time=execution_time,
- success=True
+ success=True,
)
result_queue.put(result)
-
- logger.info(f"Worker {worker_id}: Successfully completed {stub.model_name}/{stub.sparse_config_name}/{stub.benchmark_name} in {execution_time:.2f}s")
+
+ logger.info(
+ f"Worker {worker_id}: Successfully completed {stub.model_name}/{stub.sparse_config_name}/{stub.benchmark_name} in {execution_time:.2f}s"
+ )
except Exception as e:
- execution_time = time.time() - start_time # Calculate execution time even on failure
+ execution_time = (
+ time.time() - start_time
+ ) # Calculate execution time even on failure
error_category = _categorize_error(e, "benchmark execution")
error_msg = f"Benchmark execution failed: {e}"
logger.error(f"Worker {worker_id}: {error_msg}")
-
+
# Log detailed error information for debugging
- logger.debug(f"Worker {worker_id}: Error details: {traceback.format_exc()}")
-
+ logger.debug(
+ f"Worker {worker_id}: Error details: {traceback.format_exc()}"
+ )
+
# Report failure
failure = BenchmarkFailure(
stub=stub,
error_message=error_msg,
error_type=error_category,
- execution_time=execution_time
+ execution_time=execution_time,
)
error_queue.put(failure)
continue
-
+
finally:
# Step 4: Cleanup and release GPU
- _cleanup_worker_resources(gpu_pool, current_gpu_id, adapter, logger, worker_id)
+ _cleanup_worker_resources(
+ gpu_pool, current_gpu_id, adapter, logger, worker_id
+ )
current_gpu_id = None
adapter = None
-
+
except Exception as worker_error:
logger.error(f"Worker {worker_id}: Critical worker error: {worker_error}")
- logger.debug(f"Worker {worker_id}: Critical error details: {traceback.format_exc()}")
-
+ logger.debug(
+ f"Worker {worker_id}: Critical error details: {traceback.format_exc()}"
+ )
+
# Cleanup any remaining resources
_cleanup_worker_resources(gpu_pool, current_gpu_id, adapter, logger, worker_id)
-
+
# Report worker failure to error queue
- error_queue.put({
- "worker_id": worker_id,
- "error_type": "WORKER_CRITICAL_ERROR",
- "error_message": str(worker_error),
- "error_traceback": traceback.format_exc()
- })
+ error_queue.put(
+ {
+ "worker_id": worker_id,
+ "error_type": "WORKER_CRITICAL_ERROR",
+ "error_message": str(worker_error),
+ "error_traceback": traceback.format_exc(),
+ }
+ )
def _cleanup_worker_resources(
@@ -313,10 +348,10 @@ def _cleanup_worker_resources(
gpu_id: Optional[int],
adapter: Any,
logger: logging.Logger,
- worker_id: int
+ worker_id: int,
) -> None:
"""Clean up worker resources including GPU memory and adapter.
-
+
Args:
gpu_pool: GPU pool queue to return GPU to
gpu_id: GPU ID to release (can be None) - this is the actual GPU ID from the pool
@@ -331,53 +366,60 @@ def _cleanup_worker_resources(
# Force garbage collection of adapter
del adapter
import gc
+
gc.collect()
logger.debug(f"Worker {worker_id}: Cleaned up model adapter")
except Exception as e:
logger.warning(f"Worker {worker_id}: Failed to cleanup adapter: {e}")
-
+
# Clean up GPU memory and release GPU
if gpu_id is not None:
try:
cleanup_gpu_memory(gpu_id)
- logger.debug(f"Worker {worker_id}: Cleaned up GPU memory for GPU {gpu_id}")
-
+ logger.debug(
+ f"Worker {worker_id}: Cleaned up GPU memory for GPU {gpu_id}"
+ )
+
# Release GPU back to pool (always use the original GPU ID for the pool)
release_gpu_to_pool(gpu_pool, gpu_id)
logger.debug(f"Worker {worker_id}: Released GPU {gpu_id} back to pool")
-
+
except Exception as cleanup_error:
- logger.warning(f"Worker {worker_id}: GPU cleanup failed for GPU {gpu_id}: {cleanup_error}")
+ logger.warning(
+ f"Worker {worker_id}: GPU cleanup failed for GPU {gpu_id}: {cleanup_error}"
+ )
# Try to release GPU even if cleanup failed
try:
release_gpu_to_pool(gpu_pool, gpu_id)
except Exception:
- logger.error(f"Worker {worker_id}: Failed to release GPU {gpu_id} to pool")
-
+ logger.error(
+ f"Worker {worker_id}: Failed to release GPU {gpu_id} to pool"
+ )
+
except Exception as e:
logger.error(f"Worker {worker_id}: Error during resource cleanup: {e}")
class BenchmarkExecutor:
"""Parallel benchmark executor with dynamic GPU allocation and resumability.
-
- Executes benchmark experiments across a model ร sparse_attention_config ร
+
+ Executes benchmark experiments across a model ร sparse_attention_config ร
benchmark ร subset matrix with parallel GPU execution and automatic resumability.
-
+
Features:
- Dynamic GPU allocation from specified gpu_ids
- Subset-level resumability (skip completed experiments)
- Parallel execution with configurable concurrency
- Comprehensive error handling and logging
- Result aggregation and summary generation
-
+
Example:
>>> executor = BenchmarkExecutor(
... gpu_ids=[0, 1, 2, 3],
... max_concurrent_runs=4,
... base_result_dir="./benchmark_results"
... )
- >>>
+ >>>
>>> results = executor.run_benchmark_matrix(
... model_names=["model1", "model2"],
... sparse_attention_configs=[("dense", None), ("sparse", config)],
@@ -385,7 +427,7 @@ class BenchmarkExecutor:
... adapter_config=AdapterConfig()
... )
"""
-
+
def __init__(
self,
gpu_ids: List[int],
@@ -396,11 +438,13 @@ def __init__(
required_result_files: Optional[List[str]] = None,
timeout_per_benchmark: float = 360000.0, # 100 hours default
verbose: bool = True,
- micro_metric_logger_max_records: Optional[int] = None, # Maximum number of metric events to log
+ micro_metric_logger_max_records: Optional[
+ int
+ ] = None, # Maximum number of metric events to log
micro_metric_logger_sampling_factor: float = 1.0, # Probability of logging each metric event (0.0-1.0)
):
"""Initialize the BenchmarkExecutor.
-
+
Args:
gpu_ids: List of GPU device IDs to use for parallel execution
max_concurrent_runs: Maximum number of concurrent benchmark runs
@@ -412,61 +456,70 @@ def __init__(
verbose: Whether to enable verbose logging
micro_metric_logger_max_records: Maximum number of metric events to log (None for unlimited)
micro_metric_logger_sampling_factor: Probability of logging each metric event (0.0-1.0)
-
+
Raises:
ValueError: If configuration parameters are invalid
RuntimeError: If GPU validation fails
"""
# Validate and store configuration
self.gpu_ids = self._validate_gpu_ids(gpu_ids)
- self.max_concurrent_runs = self._validate_max_concurrent_runs(max_concurrent_runs)
+ self.max_concurrent_runs = self._validate_max_concurrent_runs(
+ max_concurrent_runs
+ )
self.base_result_dir = Path(base_result_dir).resolve()
self.enable_resumability = enable_resumability
self.result_file_validation = result_file_validation
self.required_result_files = required_result_files or ["raw_results.csv"]
self.timeout_per_benchmark = timeout_per_benchmark
self.verbose = verbose
-
+
# Metric logging configuration
self.micro_metric_logger_max_records = micro_metric_logger_max_records
- self.micro_metric_logger_sampling_factor = max(0.0, min(1.0, micro_metric_logger_sampling_factor)) # Clamp to [0.0, 1.0]
-
+ self.micro_metric_logger_sampling_factor = max(
+ 0.0, min(1.0, micro_metric_logger_sampling_factor)
+ ) # Clamp to [0.0, 1.0]
+
# Initialize logging
self._setup_logging()
-
+
# Set up signal handlers for graceful shutdown
self._setup_signal_handlers()
-
+
# Validate GPU availability
self._validate_gpu_setup()
-
+
# Load all available benchmarks
ensure_benchmarks_loaded()
-
-
self.logger.info(f"BenchmarkExecutor initialized:")
self.logger.info(f" GPU IDs: {self.gpu_ids}")
self.logger.info(f" Max concurrent runs: {self.max_concurrent_runs}")
self.logger.info(f" Base result directory: {self.base_result_dir}")
- self.logger.info(f" Resumability: {'enabled' if self.enable_resumability else 'disabled'}")
- self.logger.info(f" Metric logging - Max records: {self.micro_metric_logger_max_records if self.micro_metric_logger_max_records else 'unlimited'}")
- self.logger.info(f" Metric logging - Sampling factor: {self.micro_metric_logger_sampling_factor}")
-
+ self.logger.info(
+ f" Resumability: {'enabled' if self.enable_resumability else 'disabled'}"
+ )
+ self.logger.info(
+ f" Metric logging - Max records: {self.micro_metric_logger_max_records if self.micro_metric_logger_max_records else 'unlimited'}"
+ )
+ self.logger.info(
+ f" Metric logging - Sampling factor: {self.micro_metric_logger_sampling_factor}"
+ )
+
def _setup_signal_handlers(self) -> None:
"""Set up signal handlers for graceful shutdown."""
signal.signal(signal.SIGTERM, _signal_handler)
signal.signal(signal.SIGINT, _signal_handler)
-
+
# Register cleanup function to run at exit
import atexit
+
atexit.register(self._cleanup_on_exit)
-
+
def _cleanup_on_exit(self) -> None:
"""Cleanup function called when the process exits."""
- if hasattr(self, 'logger'):
+ if hasattr(self, "logger"):
self.logger.info("BenchmarkExecutor cleanup on exit")
-
+
# Reset global shutdown flag
global _shutdown_requested
_shutdown_requested = False
@@ -478,13 +531,13 @@ def run_benchmark_matrix(
benchmark_configs: List[BenchmarkConfig],
adapter_config: AdapterConfig,
generation_kwargs: Optional[Dict[str, Any]] = None,
- request_kwargs: Optional[Dict[str, Any]] = None
+ request_kwargs: Optional[Dict[str, Any]] = None,
) -> ExecutionResults:
"""Execute benchmark matrix with parallel GPU execution.
-
- Runs all combinations of model_names ร sparse_attention_configs ร
+
+ Runs all combinations of model_names ร sparse_attention_configs ร
benchmark_configs ร subsets in parallel across available GPUs.
-
+
Args:
model_names: List of model names to benchmark
sparse_attention_configs: List of (name, config) tuples for sparse attention
@@ -492,10 +545,10 @@ def run_benchmark_matrix(
adapter_config: Configuration for model adapter
generation_kwargs: Optional generation parameters
request_kwargs: Optional request processing parameters
-
+
Returns:
ExecutionResults containing comprehensive execution information
-
+
Example:
>>> results = executor.run_benchmark_matrix(
... model_names=["meta-llama/Llama-3.1-8B-Instruct"],
@@ -507,13 +560,13 @@ def run_benchmark_matrix(
"""
# Initialize execution
self.logger.info("Starting benchmark matrix execution")
-
+
try:
# Step 1: Validate input parameters
self._validate_matrix_parameters(
model_names, sparse_attention_configs, benchmark_configs, adapter_config
)
-
+
# Step 2: Generate benchmark stubs (matrix expansion)
self.logger.info("Generating benchmark matrix...")
all_stubs = generate_benchmark_stubs(
@@ -523,14 +576,16 @@ def run_benchmark_matrix(
adapter_config=adapter_config,
generation_kwargs=generation_kwargs,
request_kwargs=request_kwargs,
- base_result_dir=str(self.base_result_dir)
+ base_result_dir=str(self.base_result_dir),
)
-
+
total_combinations = len(all_stubs)
self.logger.info(f"Generated {total_combinations} benchmark combinations:")
- self.logger.info(f" {len(model_names)} models ร {len(sparse_attention_configs)} sparse configs ร "
- f"{sum(len(b.subsets) if b.subsets else 1 for b in benchmark_configs)} benchmark-subsets")
-
+ self.logger.info(
+ f" {len(model_names)} models ร {len(sparse_attention_configs)} sparse configs ร "
+ f"{sum(len(b.subsets) if b.subsets else 1 for b in benchmark_configs)} benchmark-subsets"
+ )
+
# Step 3: Filter for resumability (if enabled)
if self.enable_resumability:
self.logger.info("Checking for existing results (resumability)...")
@@ -538,50 +593,56 @@ def run_benchmark_matrix(
all_stubs,
check_result_files=self.result_file_validation,
required_files=self.required_result_files,
- verbose=self.verbose
+ verbose=self.verbose,
+ )
+
+ self.logger.info(
+ f"Resumability check: {len(completed_stubs)} completed, {len(pending_stubs)} pending"
)
-
- self.logger.info(f"Resumability check: {len(completed_stubs)} completed, {len(pending_stubs)} pending")
else:
pending_stubs = all_stubs
completed_stubs = []
- self.logger.info("Resumability disabled - will execute all combinations")
-
+ self.logger.info(
+ "Resumability disabled - will execute all combinations"
+ )
+
# Step 4: Early exit if nothing to do
if not pending_stubs:
- self.logger.info("No pending benchmarks to execute. All experiments already completed!")
-
+ self.logger.info(
+ "No pending benchmarks to execute. All experiments already completed!"
+ )
+
execution_summary = {
"total_combinations": total_combinations,
"models": model_names,
"sparse_configs": [name for name, _ in sparse_attention_configs],
- "benchmarks": [cfg.benchmark_name for cfg in benchmark_configs]
+ "benchmarks": [cfg.benchmark_name for cfg in benchmark_configs],
}
-
+
progress = ExecutionProgress(
total_stubs=total_combinations,
skipped_stubs=len(completed_stubs),
completed_stubs=0,
- failed_stubs=0
+ failed_stubs=0,
)
-
+
return ExecutionResults(
execution_summary=execution_summary,
individual_results=[],
failed_executions=[],
execution_time=0.0,
- progress=progress
+ progress=progress,
)
-
-
# Step 6: Execute pending benchmarks with worker processes
- self.logger.info(f"Executing {len(pending_stubs)} pending benchmarks with {self.max_concurrent_runs} workers...")
-
+ self.logger.info(
+ f"Executing {len(pending_stubs)} pending benchmarks with {self.max_concurrent_runs} workers..."
+ )
+
start_time = time.time()
results = self._execute_with_workers(pending_stubs)
execution_time = time.time() - start_time
-
+
# Step 7: Aggregate results
execution_summary = {
"total_combinations": total_combinations,
@@ -590,101 +651,112 @@ def run_benchmark_matrix(
"benchmarks": [cfg.benchmark_name for cfg in benchmark_configs],
"execution_time_seconds": execution_time,
"gpu_ids_used": self.gpu_ids,
- "max_concurrent_runs": self.max_concurrent_runs
+ "max_concurrent_runs": self.max_concurrent_runs,
}
-
+
progress = ExecutionProgress(
total_stubs=total_combinations,
skipped_stubs=len(completed_stubs),
completed_stubs=len(results.individual_results),
- failed_stubs=len(results.failed_executions)
+ failed_stubs=len(results.failed_executions),
)
-
+
final_results = ExecutionResults(
execution_summary=execution_summary,
individual_results=results.individual_results,
failed_executions=results.failed_executions,
execution_time=execution_time,
- progress=progress
+ progress=progress,
)
-
+
self.logger.info(f"Benchmark matrix execution completed:")
self.logger.info(f" Total time: {execution_time:.2f}s")
self.logger.info(f" Successful: {len(results.individual_results)}")
self.logger.info(f" Failed: {len(results.failed_executions)}")
self.logger.info(f" Skipped: {len(completed_stubs)}")
-
+
return final_results
-
+
except Exception as e:
self.logger.error(f"Benchmark execution failed: {e}")
raise
-
- def _execute_with_workers(self, pending_stubs: List[BenchmarkStub]) -> ExecutionResults:
+
+ def _execute_with_workers(
+ self, pending_stubs: List[BenchmarkStub]
+ ) -> ExecutionResults:
"""Execute benchmark stubs using worker processes.
-
+
This method implements comprehensive process pool management with:
- Dynamic work distribution through shared queues
- Graceful process startup and shutdown
- Real-time result collection and error handling
- Process health monitoring and recovery
-
+
Args:
pending_stubs: List of benchmark stubs to execute
-
+
Returns:
ExecutionResults containing all results and failures
"""
- self.logger.info(f"Starting execution with {len(pending_stubs)} stubs using {self.max_concurrent_runs} workers")
-
+ self.logger.info(
+ f"Starting execution with {len(pending_stubs)} stubs using {self.max_concurrent_runs} workers"
+ )
+
# Step 1: Create and configure shared queues
- stub_queue = multiprocessing.Queue(maxsize=len(pending_stubs) + 100) # Buffer for queue operations
+ stub_queue = multiprocessing.Queue(
+ maxsize=len(pending_stubs) + 100
+ ) # Buffer for queue operations
gpu_pool = create_gpu_pool(self.gpu_ids)
result_queue = multiprocessing.Queue()
error_queue = multiprocessing.Queue()
-
+
# Step 2: Populate work queue with all stubs
self.logger.info("Populating work queue...")
for stub in pending_stubs:
stub_queue.put(stub)
-
+
# Add sentinel values to signal workers to stop
for _ in range(self.max_concurrent_runs):
stub_queue.put(None) # Sentinel value
-
+
# Step 3: Create and start worker processes
workers = []
worker_pids = []
-
+
self.logger.info(f"Starting {self.max_concurrent_runs} worker processes...")
for i in range(self.max_concurrent_runs):
worker = multiprocessing.Process(
target=_benchmark_worker,
- args=(stub_queue,
- gpu_pool,
- result_queue,
- error_queue,
- self.timeout_per_benchmark,
- self.micro_metric_logger_max_records,
- self.micro_metric_logger_sampling_factor
- ),
- name=f"benchmark_worker_{i}"
+ args=(
+ stub_queue,
+ gpu_pool,
+ result_queue,
+ error_queue,
+ self.timeout_per_benchmark,
+ self.micro_metric_logger_max_records,
+ self.micro_metric_logger_sampling_factor,
+ ),
+ name=f"benchmark_worker_{i}",
+ )
+ worker.daemon = (
+ True # Ensure workers are terminated when main process exits
)
- worker.daemon = True # Ensure workers are terminated when main process exits
worker.start()
workers.append(worker)
worker_pids.append(worker.pid)
- self.logger.debug(f"Started worker {i+1}/{self.max_concurrent_runs} (PID: {worker.pid})")
-
+ self.logger.debug(
+ f"Started worker {i+1}/{self.max_concurrent_runs} (PID: {worker.pid})"
+ )
+
self.logger.info(f"All {self.max_concurrent_runs} workers started successfully")
-
+
# Step 4: Monitor and collect results in real-time
individual_results: List[BenchmarkResult] = []
failed_executions: List[BenchmarkFailure] = []
active_workers = set(worker_pids)
-
+
self.logger.info("Monitoring worker processes and collecting results...")
-
+
try:
# Monitor workers and collect results until all work is complete
while active_workers:
@@ -699,47 +771,74 @@ def _execute_with_workers(self, pending_stubs: List[BenchmarkStub]) -> Execution
if worker.pid in active_workers:
active_workers.remove(worker.pid)
completed_workers.append(worker)
-
+
if worker.exitcode != 0:
- self.logger.warning(f"Worker {worker.pid} exited with code {worker.exitcode}")
+ self.logger.warning(
+ f"Worker {worker.pid} exited with code {worker.exitcode}"
+ )
else:
- self.logger.debug(f"Worker {worker.pid} completed successfully")
-
+ self.logger.debug(
+ f"Worker {worker.pid} completed successfully"
+ )
+
# Collect any available results (non-blocking)
- self._collect_available_results(result_queue, error_queue, individual_results, failed_executions)
-
+ self._collect_available_results(
+ result_queue, error_queue, individual_results, failed_executions
+ )
+
# Log progress periodically
- if len(individual_results) + len(failed_executions) > 0 and (len(individual_results) + len(failed_executions)) % 5 == 0:
- self.logger.info(f"Progress: {len(individual_results)} completed, {len(failed_executions)} failed, {len(active_workers)} workers active")
-
+ if (
+ len(individual_results) + len(failed_executions) > 0
+ and (len(individual_results) + len(failed_executions)) % 5 == 0
+ ):
+ self.logger.info(
+ f"Progress: {len(individual_results)} completed, {len(failed_executions)} failed, {len(active_workers)} workers active"
+ )
+
# Small delay to prevent busy waiting
time.sleep(0.1)
-
+
# Step 5: Final result collection
self.logger.info("All workers completed, collecting final results...")
- self._collect_available_results(result_queue, error_queue, individual_results, failed_executions)
-
+ self._collect_available_results(
+ result_queue, error_queue, individual_results, failed_executions
+ )
+
except KeyboardInterrupt:
- self.logger.warning("Received interrupt signal, initiating graceful shutdown...")
- self._graceful_shutdown(workers, stub_queue, gpu_pool, result_queue, error_queue)
+ self.logger.warning(
+ "Received interrupt signal, initiating graceful shutdown..."
+ )
+ self._graceful_shutdown(
+ workers, stub_queue, gpu_pool, result_queue, error_queue
+ )
raise
-
+
except Exception as e:
error_category = _categorize_error(e, "main execution")
self.logger.error(f"Error during execution ({error_category}): {e}")
self.logger.debug(f"Execution error details: {traceback.format_exc()}")
- self._graceful_shutdown(workers, stub_queue, gpu_pool, result_queue, error_queue)
+ self._graceful_shutdown(
+ workers, stub_queue, gpu_pool, result_queue, error_queue
+ )
raise
-
+
finally:
# Step 6: Cleanup resources
- self._cleanup_resources(workers, stub_queue, gpu_pool, result_queue, error_queue)
-
+ self._cleanup_resources(
+ workers, stub_queue, gpu_pool, result_queue, error_queue
+ )
+
# Step 7: Final summary
total_processed = len(individual_results) + len(failed_executions)
- self.logger.info(f"Execution completed: {len(individual_results)} successful, {len(failed_executions)} failed")
- self.logger.info(f"Success rate: {len(individual_results)/total_processed*100:.1f}%" if total_processed > 0 else "No work processed")
-
+ self.logger.info(
+ f"Execution completed: {len(individual_results)} successful, {len(failed_executions)} failed"
+ )
+ self.logger.info(
+ f"Success rate: {len(individual_results)/total_processed*100:.1f}%"
+ if total_processed > 0
+ else "No work processed"
+ )
+
return ExecutionResults(
execution_summary={}, # Will be filled by caller
individual_results=individual_results,
@@ -748,19 +847,19 @@ def _execute_with_workers(self, pending_stubs: List[BenchmarkStub]) -> Execution
progress=ExecutionProgress(
total_stubs=len(pending_stubs),
completed_stubs=len(individual_results),
- failed_stubs=len(failed_executions)
- )
+ failed_stubs=len(failed_executions),
+ ),
)
-
+
def _collect_available_results(
self,
result_queue: multiprocessing.Queue,
error_queue: multiprocessing.Queue,
individual_results: List[BenchmarkResult],
- failed_executions: List[BenchmarkFailure]
+ failed_executions: List[BenchmarkFailure],
) -> None:
"""Collect all available results from queues without blocking.
-
+
Args:
result_queue: Queue containing successful benchmark results
error_queue: Queue containing benchmark failures
@@ -774,7 +873,7 @@ def _collect_available_results(
individual_results.append(result)
except Empty:
break
-
+
# Collect all available errors
while True:
try:
@@ -786,17 +885,17 @@ def _collect_available_results(
self.logger.error(f"Worker error: {error}")
except Empty:
break
-
+
def _graceful_shutdown(
self,
workers: List[multiprocessing.Process],
stub_queue: multiprocessing.Queue,
gpu_pool: multiprocessing.Queue,
result_queue: multiprocessing.Queue,
- error_queue: multiprocessing.Queue
+ error_queue: multiprocessing.Queue,
) -> None:
"""Perform graceful shutdown of worker processes and queues.
-
+
Args:
workers: List of worker processes to terminate
stub_queue: Work queue to close
@@ -805,7 +904,7 @@ def _graceful_shutdown(
error_queue: Error queue to close
"""
self.logger.info("Initiating graceful shutdown...")
-
+
try:
# Signal workers to stop by adding sentinel values
sentinel_count = 0
@@ -816,9 +915,9 @@ def _graceful_shutdown(
except Exception as e:
self.logger.debug(f"Could not add sentinel to queue: {e}")
break # Queue is full, stop adding sentinels
-
+
self.logger.info(f"Added {sentinel_count} shutdown signals to work queue")
-
+
# Wait for workers to finish with timeout
self.logger.info("Waiting for workers to finish...")
for i, worker in enumerate(workers):
@@ -829,16 +928,20 @@ def _graceful_shutdown(
worker.terminate()
worker.join(timeout=5.0)
if worker.is_alive():
- self.logger.error(f"Failed to terminate worker {worker.pid}, killing process")
+ self.logger.error(
+ f"Failed to terminate worker {worker.pid}, killing process"
+ )
worker.kill()
worker.join(timeout=2.0)
else:
- self.logger.debug(f"Worker {worker.pid} terminated successfully")
+ self.logger.debug(
+ f"Worker {worker.pid} terminated successfully"
+ )
except Exception as e:
self.logger.error(f"Error during worker {worker.pid} shutdown: {e}")
-
+
self.logger.info("Graceful shutdown completed")
-
+
except Exception as e:
self.logger.error(f"Error during graceful shutdown: {e}")
# Force cleanup of remaining workers
@@ -848,17 +951,17 @@ def _graceful_shutdown(
worker.kill()
except Exception:
pass
-
+
def _cleanup_resources(
self,
workers: List[multiprocessing.Process],
stub_queue: multiprocessing.Queue,
gpu_pool: multiprocessing.Queue,
result_queue: multiprocessing.Queue,
- error_queue: multiprocessing.Queue
+ error_queue: multiprocessing.Queue,
) -> None:
"""Clean up all multiprocessing resources.
-
+
Args:
workers: List of worker processes
stub_queue: Work queue to close
@@ -867,7 +970,7 @@ def _cleanup_resources(
error_queue: Error queue to close
"""
self.logger.debug("Cleaning up multiprocessing resources...")
-
+
# Close all queues
try:
stub_queue.close()
@@ -876,7 +979,7 @@ def _cleanup_resources(
error_queue.close()
except Exception as e:
self.logger.warning(f"Error closing queues: {e}")
-
+
# Join any remaining workers
for worker in workers:
if worker.is_alive():
@@ -885,17 +988,17 @@ def _cleanup_resources(
worker.join(timeout=5.0)
if worker.is_alive():
worker.kill()
-
+
self.logger.debug("Resource cleanup completed")
-
+
def _validate_gpu_ids(self, gpu_ids: List[int]) -> List[int]:
"""Validate and return GPU IDs."""
if not gpu_ids:
raise ValueError("gpu_ids cannot be empty")
-
+
if not all(isinstance(gpu_id, int) and gpu_id >= 0 for gpu_id in gpu_ids):
raise ValueError("All GPU IDs must be non-negative integers")
-
+
# Remove duplicates while preserving order
seen = set()
unique_gpu_ids = []
@@ -903,34 +1006,38 @@ def _validate_gpu_ids(self, gpu_ids: List[int]) -> List[int]:
if gpu_id not in seen:
seen.add(gpu_id)
unique_gpu_ids.append(gpu_id)
-
+
return unique_gpu_ids
-
+
def _validate_max_concurrent_runs(self, max_concurrent_runs: int) -> int:
"""Validate and return max_concurrent_runs."""
if not isinstance(max_concurrent_runs, int) or max_concurrent_runs <= 0:
raise ValueError("max_concurrent_runs must be a positive integer")
-
- if max_concurrent_runs > len(self.gpu_ids) if hasattr(self, 'gpu_ids') else True:
+
+ if (
+ max_concurrent_runs > len(self.gpu_ids)
+ if hasattr(self, "gpu_ids")
+ else True
+ ):
# This is just a warning - user might want more concurrent runs than GPUs
# for benchmarks that don't require much GPU memory
pass
-
+
return max_concurrent_runs
-
+
def _setup_logging(self) -> None:
"""Setup logging configuration."""
self.logger = logging.getLogger(f"{__name__}.BenchmarkExecutor")
-
+
if self.verbose and not self.logger.handlers:
handler = logging.StreamHandler()
formatter = logging.Formatter(
- '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+ "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
handler.setFormatter(formatter)
self.logger.addHandler(handler)
self.logger.setLevel(logging.INFO)
-
+
def _validate_gpu_setup(self) -> None:
"""Validate GPU availability."""
try:
@@ -939,47 +1046,53 @@ def _validate_gpu_setup(self) -> None:
except (RuntimeError, ValueError) as e:
self.logger.error(f"GPU validation failed: {e}")
raise
-
+
def _validate_matrix_parameters(
self,
model_names: List[str],
sparse_attention_configs: List[Tuple[str, Optional[SparseAttentionConfig]]],
benchmark_configs: List[BenchmarkConfig],
- adapter_config: AdapterConfig
+ adapter_config: AdapterConfig,
) -> None:
"""Validate benchmark matrix parameters."""
if not model_names:
raise ValueError("model_names cannot be empty")
-
+
if not sparse_attention_configs:
raise ValueError("sparse_attention_configs cannot be empty")
-
+
if not benchmark_configs:
raise ValueError("benchmark_configs cannot be empty")
-
+
# Validate model names
for model_name in model_names:
if not isinstance(model_name, str) or not model_name.strip():
raise ValueError(f"Invalid model name: {model_name}")
-
+
# Validate sparse attention configs
for sparse_config in sparse_attention_configs:
if not isinstance(sparse_config, tuple) or len(sparse_config) != 2:
- raise ValueError(f"sparse_attention_configs must be list of (name, config) tuples")
-
+ raise ValueError(
+ f"sparse_attention_configs must be list of (name, config) tuples"
+ )
+
config_name, config_obj = sparse_config
if not isinstance(config_name, str) or not config_name.strip():
- raise ValueError(f"Sparse config name must be non-empty string: {config_name}")
-
+ raise ValueError(
+ f"Sparse config name must be non-empty string: {config_name}"
+ )
+
# Validate benchmark configs
for benchmark_config in benchmark_configs:
if not isinstance(benchmark_config, BenchmarkConfig):
- raise ValueError(f"benchmark_configs must contain BenchmarkConfig instances")
+ raise ValueError(
+ f"benchmark_configs must contain BenchmarkConfig instances"
+ )
# Additional validation could be added here if needed
-
+
# Validate adapter config
if not isinstance(adapter_config, AdapterConfig):
raise ValueError("adapter_config must be an AdapterConfig instance")
# Additional validation could be added here if needed
-
+
self.logger.info("Matrix parameter validation passed")
diff --git a/benchmark/executor_config.py b/benchmark/executor_config.py
index 081d1871..9ef7febf 100644
--- a/benchmark/executor_config.py
+++ b/benchmark/executor_config.py
@@ -25,53 +25,53 @@
acquire_gpu_from_pool,
release_gpu_to_pool,
validate_gpu_for_model,
- cleanup_gpu_memory
+ cleanup_gpu_memory,
)
@dataclass
class BenchmarkConfig:
"""Configuration for a benchmark execution.
-
+
During stub generation, this expands to individual BenchmarkStub per subset.
Each BenchmarkStub will have a single subset value for resumability.
-
+
Attributes:
benchmark_name: Name of the benchmark (e.g., "longbench", "mock_benchmark")
subsets: Optional list of dataset subsets to run (e.g., ["narrativeqa", "qasper"])
If None, runs all available subsets for the benchmark
-
+
Example:
>>> config = BenchmarkConfig("longbench", ["narrativeqa", "qasper"])
>>> # This will create 2 BenchmarkStub instances during matrix expansion
"""
-
+
benchmark_name: str
subsets: Optional[List[str]] = None
-
+
def __post_init__(self) -> None:
"""Validate benchmark configuration after initialization.
-
+
Raises:
ValueError: If benchmark_name is empty or invalid.
"""
if not self.benchmark_name or not self.benchmark_name.strip():
raise ValueError("benchmark_name cannot be empty")
-
+
self.benchmark_name = self.benchmark_name.strip()
-
+
if self.subsets is not None:
# Remove empty strings and strip whitespace
self.subsets = [subset.strip() for subset in self.subsets if subset.strip()]
if not self.subsets:
self.subsets = None
-
+
def validate_with_benchmark_instance(self, benchmark: Benchmark) -> None:
"""Validate subsets against actual benchmark instance.
-
+
Args:
benchmark: Benchmark instance to validate against
-
+
Raises:
ValueError: If any subset is not available in the benchmark.
"""
@@ -88,47 +88,47 @@ def validate_with_benchmark_instance(self, benchmark: Benchmark) -> None:
@dataclass
class AdapterConfig:
"""Base configuration for model adapters (sparse attention config is part of matrix).
-
+
Note: sparse_attention_config is now part of the execution matrix, not base config.
This allows testing different sparse attention configurations with the same adapter base.
-
+
Attributes:
adapter_name: Name of the adapter type (default: "huggingface")
model_kwargs: Additional keyword arguments for model creation
tokenizer_kwargs: Additional keyword arguments for tokenizer creation
-
+
Example:
>>> config = AdapterConfig(
... model_kwargs={"torch_dtype": torch.bfloat16},
... tokenizer_kwargs={"padding_side": "left"}
... )
"""
-
+
adapter_name: str = "huggingface"
model_kwargs: Optional[Dict[str, Any]] = None
tokenizer_kwargs: Optional[Dict[str, Any]] = None
-
+
def __post_init__(self) -> None:
"""Initialize default values and validate configuration."""
if self.model_kwargs is None:
self.model_kwargs = {}
if self.tokenizer_kwargs is None:
self.tokenizer_kwargs = {}
-
+
if not self.adapter_name or not self.adapter_name.strip():
raise ValueError("adapter_name cannot be empty")
-
+
self.adapter_name = self.adapter_name.strip()
@dataclass
class BenchmarkStub:
"""A single benchmark execution task (model ร sparse_config ร benchmark ร subset).
-
+
This represents one atomic execution unit that will be processed by a worker process.
Each stub corresponds to a specific combination of model, sparse attention config,
benchmark, and subset (if applicable).
-
+
Attributes:
model_name: HuggingFace model name (e.g., "meta-llama/Llama-3.1-8B-Instruct")
sparse_config_name: Human-readable name for sparse config (e.g., "dense", "streaming")
@@ -139,7 +139,7 @@ class BenchmarkStub:
generation_kwargs: Parameters for model inference/generation
request_kwargs: Parameters for request processing (e.g., max_context_length)
result_dir: Full path to directory where results should be saved
-
+
Example:
>>> stub = BenchmarkStub(
... model_name="microsoft/Phi-4-mini-instruct",
@@ -153,7 +153,7 @@ class BenchmarkStub:
... result_dir="/path/to/results/phi4_dense_longbench_narrativeqa"
... )
"""
-
+
model_name: str
sparse_config_name: str
sparse_attention_config: Optional[SparseAttentionConfig]
@@ -163,7 +163,7 @@ class BenchmarkStub:
generation_kwargs: Dict[str, Any]
request_kwargs: Dict[str, Any]
result_dir: str
-
+
def __post_init__(self) -> None:
"""Validate benchmark stub configuration."""
if not self.model_name or not self.model_name.strip():
@@ -174,13 +174,13 @@ def __post_init__(self) -> None:
raise ValueError("benchmark_name cannot be empty")
if not self.result_dir or not self.result_dir.strip():
raise ValueError("result_dir cannot be empty")
-
+
# Trim whitespace
self.model_name = self.model_name.strip()
self.sparse_config_name = self.sparse_config_name.strip()
self.benchmark_name = self.benchmark_name.strip()
self.result_dir = self.result_dir.strip()
-
+
if self.subset is not None:
self.subset = self.subset.strip() if self.subset.strip() else None
@@ -188,7 +188,7 @@ def __post_init__(self) -> None:
@dataclass
class BenchmarkResult:
"""Result from a single benchmark execution.
-
+
Attributes:
stub: The benchmark stub that was executed
metrics: Evaluation metrics returned by benchmark.run_benchmark()
@@ -196,7 +196,7 @@ class BenchmarkResult:
success: Whether the benchmark executed successfully
error_message: Error message if execution failed
"""
-
+
stub: BenchmarkStub
metrics: Optional[Dict[str, Any]]
execution_time: float
@@ -207,14 +207,14 @@ class BenchmarkResult:
@dataclass
class BenchmarkFailure:
"""Information about a failed benchmark execution.
-
+
Attributes:
stub: The benchmark stub that failed
error_message: Description of the failure
error_type: Type of error that occurred
execution_time: Time spent before failure in seconds
"""
-
+
stub: BenchmarkStub
error_message: str
error_type: str
@@ -224,7 +224,7 @@ class BenchmarkFailure:
@dataclass
class ExecutionProgress:
"""Tracks execution progress across all processes.
-
+
Attributes:
total_stubs: Total number of benchmark stubs to execute
skipped_stubs: Number of stubs skipped due to existing results (resumability)
@@ -232,30 +232,39 @@ class ExecutionProgress:
failed_stubs: Number of stubs that failed execution
current_executions: Mapping of GPU ID to current task description
"""
-
+
total_stubs: int
skipped_stubs: int = 0
completed_stubs: int = 0
failed_stubs: int = 0
current_executions: Dict[int, str] = field(default_factory=dict)
-
+
@property
def pending_stubs(self) -> int:
"""Calculate number of stubs pending execution."""
- return self.total_stubs - self.skipped_stubs - self.completed_stubs - self.failed_stubs
-
+ return (
+ self.total_stubs
+ - self.skipped_stubs
+ - self.completed_stubs
+ - self.failed_stubs
+ )
+
@property
def completion_percentage(self) -> float:
"""Calculate completion percentage (0-100)."""
if self.total_stubs == 0:
return 100.0
- return (self.completed_stubs + self.failed_stubs) / (self.total_stubs - self.skipped_stubs) * 100.0
+ return (
+ (self.completed_stubs + self.failed_stubs)
+ / (self.total_stubs - self.skipped_stubs)
+ * 100.0
+ )
@dataclass
class ExecutionResults:
"""Aggregated results from benchmark execution.
-
+
Attributes:
execution_summary: Summary metadata about the execution
individual_results: List of successful benchmark results
@@ -263,7 +272,7 @@ class ExecutionResults:
execution_time: Total execution time in seconds
progress: Final progress statistics
"""
-
+
execution_summary: Dict[str, Any]
individual_results: List[BenchmarkResult]
failed_executions: List[BenchmarkFailure]
@@ -283,35 +292,35 @@ class ExecutionResults:
create_benchmark_instance,
ensure_benchmarks_loaded,
validate_benchmark_config as _validate_benchmark_config_base,
- get_benchmark_subsets
+ get_benchmark_subsets,
)
def validate_benchmark_config(config: BenchmarkConfig) -> None:
"""Validate a benchmark configuration against available benchmarks.
-
+
This function checks if the benchmark exists and if specified subsets are valid.
-
+
Args:
config: BenchmarkConfig to validate
-
+
Raises:
ValueError: If benchmark or subsets are invalid
-
+
Example:
>>> config = BenchmarkConfig("longbench", ["narrativeqa"])
>>> validate_benchmark_config(config) # Passes validation
- >>>
+ >>>
>>> config = BenchmarkConfig("invalid_benchmark")
>>> validate_benchmark_config(config) # Raises ValueError
"""
try:
# Create temporary benchmark instance to validate
benchmark = create_benchmark_instance(config.benchmark_name, subsets=None)
-
+
# Validate subsets if specified
config.validate_with_benchmark_instance(benchmark)
-
+
except Exception as e:
raise ValueError(f"Invalid benchmark configuration: {e}")
@@ -320,6 +329,7 @@ def validate_benchmark_config(config: BenchmarkConfig) -> None:
# Matrix Expansion and Benchmark Stub Generation Functions
# =============================================================================
+
def generate_benchmark_stubs(
model_names: List[str],
sparse_attention_configs: List[Tuple[str, Optional[SparseAttentionConfig]]],
@@ -327,13 +337,13 @@ def generate_benchmark_stubs(
adapter_config: AdapterConfig,
generation_kwargs: Optional[Dict[str, Any]] = None,
request_kwargs: Optional[Dict[str, Any]] = None,
- base_result_dir: str = "./benchmark_results"
+ base_result_dir: str = "./benchmark_results",
) -> List[BenchmarkStub]:
"""Generate all benchmark stubs for the model ร sparse_config ร benchmark ร subset matrix.
-
+
Creates BenchmarkStub instances for every combination of:
- model_names ร sparse_attention_configs ร benchmark_configs ร subsets
-
+
Args:
model_names: List of model names to benchmark
sparse_attention_configs: List of (name, config) tuples for sparse attention
@@ -342,13 +352,13 @@ def generate_benchmark_stubs(
generation_kwargs: Optional generation parameters
request_kwargs: Optional request processing parameters
base_result_dir: Base directory for storing results
-
+
Returns:
List of BenchmarkStub instances representing all combinations
-
+
Example:
>>> stubs = generate_benchmark_stubs(
- ... model_names=["model1", "model2"],
+ ... model_names=["model1", "model2"],
... sparse_attention_configs=[("dense", None), ("sparse", config)],
... benchmark_configs=[BenchmarkConfig("longbench", ["narrativeqa"])],
... adapter_config=AdapterConfig()
@@ -356,14 +366,14 @@ def generate_benchmark_stubs(
>>> len(stubs) # 2 models ร 2 configs ร 1 benchmark ร 1 subset = 4
4
"""
-
+
if generation_kwargs is None:
generation_kwargs = {}
if request_kwargs is None:
request_kwargs = {}
-
+
stubs: List[BenchmarkStub] = []
-
+
# Expand the full matrix: model ร sparse_config ร benchmark ร subset
for model_name in model_names:
for sparse_config_name, sparse_attention_config in sparse_attention_configs:
@@ -376,9 +386,9 @@ def generate_benchmark_stubs(
model_name=model_name,
sparse_config_name=sparse_config_name,
benchmark_name=benchmark_config.benchmark_name,
- subset=subset
+ subset=subset,
)
-
+
stub = BenchmarkStub(
model_name=model_name,
sparse_config_name=sparse_config_name,
@@ -388,10 +398,10 @@ def generate_benchmark_stubs(
adapter_config=adapter_config,
generation_kwargs=generation_kwargs.copy(),
request_kwargs=request_kwargs.copy(),
- result_dir=result_dir
+ result_dir=result_dir,
)
stubs.append(stub)
-
+
else:
# Handle benchmarks with no subsets (run full benchmark)
result_dir = construct_result_directory(
@@ -399,9 +409,9 @@ def generate_benchmark_stubs(
model_name=model_name,
sparse_config_name=sparse_config_name,
benchmark_name=benchmark_config.benchmark_name,
- subset=None
+ subset=None,
)
-
+
stub = BenchmarkStub(
model_name=model_name,
sparse_config_name=sparse_config_name,
@@ -411,10 +421,10 @@ def generate_benchmark_stubs(
adapter_config=adapter_config,
generation_kwargs=generation_kwargs.copy(),
request_kwargs=request_kwargs.copy(),
- result_dir=result_dir
+ result_dir=result_dir,
)
stubs.append(stub)
-
+
return stubs
@@ -422,50 +432,50 @@ def filter_existing_results(
stubs: List[BenchmarkStub],
check_result_files: bool = True,
required_files: Optional[List[str]] = None,
- verbose: bool = True
+ verbose: bool = True,
) -> Tuple[List[BenchmarkStub], List[BenchmarkStub]]:
"""Filter benchmark stubs to skip already completed experiments (resumability).
-
+
Checks each stub's result_dir to determine if the experiment has already been
completed. This enables resuming interrupted benchmark runs.
-
+
Args:
stubs: List of benchmark stubs to filter
check_result_files: Whether to check for specific result files (not just directory)
required_files: List of filenames that must exist for completion (default: ["results.json"])
verbose: Whether to log skipped combinations
-
+
Returns:
Tuple of (pending_stubs, completed_stubs) where:
- pending_stubs: Stubs that need to be executed
- completed_stubs: Stubs that have already been completed
-
+
Example:
>>> pending, completed = filter_existing_results(all_stubs)
>>> print(f"Found {len(completed)} completed, {len(pending)} pending")
"""
if required_files is None:
required_files = ["raw_results.csv"]
-
+
pending_stubs: List[BenchmarkStub] = []
completed_stubs: List[BenchmarkStub] = []
-
+
for stub in stubs:
result_path = Path(stub.result_dir)
print(f"Checking result directory: {result_path}", flush=True)
-
+
# Check if result directory exists
if not result_path.exists():
pending_stubs.append(stub)
continue
-
+
# If not checking result files, just check directory existence
if not check_result_files:
completed_stubs.append(stub)
if verbose:
logging.info(f"Skipping completed experiment: {stub.result_dir}")
continue
-
+
# Check for required result files
all_files_exist = True
for required_file in required_files:
@@ -473,20 +483,24 @@ def filter_existing_results(
if not file_path.exists() or not file_path.is_file():
all_files_exist = False
break
-
+
if all_files_exist:
completed_stubs.append(stub)
if verbose:
- logging.info(f"Skipping completed experiment (has {required_files}): {stub.result_dir}")
+ logging.info(
+ f"Skipping completed experiment (has {required_files}): {stub.result_dir}"
+ )
else:
pending_stubs.append(stub)
if verbose:
logging.info(f"Pending experiment (missing results): {stub.result_dir}")
-
+
if verbose:
total = len(stubs)
completed_count = len(completed_stubs)
pending_count = len(completed_stubs)
- logging.info(f"Resumability check: {completed_count}/{total} completed, {pending_count}/{total} pending")
-
- return pending_stubs, completed_stubs
\ No newline at end of file
+ logging.info(
+ f"Resumability check: {completed_count}/{total} completed, {pending_count}/{total} pending"
+ )
+
+ return pending_stubs, completed_stubs
diff --git a/benchmark/infinite_bench/__init__.py b/benchmark/infinite_bench/__init__.py
index 2a5de88d..7f26a1ea 100644
--- a/benchmark/infinite_bench/__init__.py
+++ b/benchmark/infinite_bench/__init__.py
@@ -8,4 +8,4 @@
from .calculate_metrics import calculate_metrics
from .infinite_bench import InfiniteBench
-__all__ = ["calculate_metrics", "InfiniteBench"]
\ No newline at end of file
+__all__ = ["calculate_metrics", "InfiniteBench"]
diff --git a/benchmark/infinite_bench/calculate_metrics.py b/benchmark/infinite_bench/calculate_metrics.py
index e9937896..489d8fb3 100644
--- a/benchmark/infinite_bench/calculate_metrics.py
+++ b/benchmark/infinite_bench/calculate_metrics.py
@@ -365,7 +365,9 @@ def get_score_one(pred: str, label: str, task_name: str, model_name: str) -> flo
"math_calc": get_score_one_math_calc,
}
if task_name == "longbook_sum_eng":
- raise AssertionError("Longbook sum task is not supported due to a conflict with rouge library. ")
+ raise AssertionError(
+ "Longbook sum task is not supported due to a conflict with rouge library. "
+ )
assert task_name in NAME_TO_SCORE_GETTER, f"Invalid task name: {task_name}"
score = NAME_TO_SCORE_GETTER[task_name](pred, label, model_name)
return float(score)
diff --git a/benchmark/infinite_bench/create_huggingface_dataset.py b/benchmark/infinite_bench/create_huggingface_dataset.py
index 60303fa5..363b0413 100644
--- a/benchmark/infinite_bench/create_huggingface_dataset.py
+++ b/benchmark/infinite_bench/create_huggingface_dataset.py
@@ -156,7 +156,8 @@ def update_math_find_context(row):
df["context"] = df.apply(update_math_find_context, axis=1)
df["context"] = df["context"].apply(
lambda x: x.replace(
- "You should answer with only one number, no other words. The largest number of the list is:", ""
+ "You should answer with only one number, no other words. The largest number of the list is:",
+ "",
)
)
elif task == "code_run":
@@ -177,13 +178,19 @@ def update_context(row):
df["context"] = df.apply(update_context, axis=1)
else:
- df["context"] = df["context"].apply(lambda x: context_prefix[task].format(context=x))
+ df["context"] = df["context"].apply(
+ lambda x: context_prefix[task].format(context=x)
+ )
if task in ["longbook_choice_eng", "code_debug"]:
df["question"] = df["question"] + df["options"].apply(
- lambda x: question_template[task].format(OPTION_A=x[0], OPTION_B=x[1], OPTION_C=x[2], OPTION_D=x[3])
+ lambda x: question_template[task].format(
+ OPTION_A=x[0], OPTION_B=x[1], OPTION_C=x[2], OPTION_D=x[3]
+ )
+ )
+ df["answer"] = df.apply(
+ lambda row: ["ABCD"[list(row.options).index(row.answer)]], axis=1
)
- df["answer"] = df.apply(lambda row: ["ABCD"[list(row.options).index(row.answer)]], axis=1)
if task == "kv_retrieval":
# moved to answer prefix
diff --git a/benchmark/infinite_bench/infinite_bench.py b/benchmark/infinite_bench/infinite_bench.py
index 3b1bd50b..9a373e05 100644
--- a/benchmark/infinite_bench/infinite_bench.py
+++ b/benchmark/infinite_bench/infinite_bench.py
@@ -31,34 +31,37 @@ class InfiniteBench(Benchmark):
# All available InfiniteBench datasets
all_datasets: List[str] = [
"passkey",
- "kv_retrieval",
+ "kv_retrieval",
"number_string",
"code_run",
"code_debug",
"math_find",
"longbook_qa_eng",
"longdialogue_qa_eng",
- "longbook_choice_eng"
+ "longbook_choice_eng",
]
-
+
benchmark_name: str = "infinite_bench"
huggingface_dataset_id: str = "MaxJeblick/InfiniteBench"
def _load_datasets(self) -> pd.DataFrame:
"""Load InfiniteBench datasets by individual configs.
-
+
InfiniteBench requires loading each subset as a separate config.
-
+
Returns:
Combined pandas DataFrame with all samples from subsets_to_run.
"""
print(f"Loading InfiniteBench datasets: {self.subsets_to_run}")
dfs = []
-
+
for subset in self.subsets_to_run:
try:
from datasets import load_dataset
- subset_dataset = load_dataset(self.huggingface_dataset_id, subset, split="test")
+
+ subset_dataset = load_dataset(
+ self.huggingface_dataset_id, subset, split="test"
+ )
subset_df = subset_dataset.to_pandas()
subset_df["task"] = subset # Ensure task column exists
dfs.append(subset_df)
@@ -66,12 +69,13 @@ def _load_datasets(self) -> pd.DataFrame:
except Exception as subset_error:
print(f" โ Failed to load {subset}: {str(subset_error)}")
continue
-
+
if not dfs:
raise Exception("No InfiniteBench subsets could be loaded successfully")
-
+
# Combine all subset DataFrames
import pandas as pd
+
combined_df = pd.concat(dfs, ignore_index=True)
print(f"Combined {len(combined_df)} total samples from {len(dfs)} subsets")
return combined_df
@@ -97,7 +101,7 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
task_groups = results_df.groupby("task")
task_scores: Dict[str, float] = {}
all_scores: List[float] = []
-
+
for task_name, task_df in task_groups:
try:
# Use the calculate_metrics function from HashAttention evaluation
@@ -105,21 +109,23 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
task_scores[task_name] = score
all_scores.append(score)
print(f" โ {task_name}: {score:.4f}")
-
+
except Exception as e:
print(f"Error evaluating task {task_name}: {str(e)}")
task_scores[task_name] = 0.0
# Compute overall metrics
overall_score = sum(all_scores) / len(all_scores) if all_scores else 0.0
-
+
overall_metrics: Dict[str, Any] = {
"overall_score": round(overall_score, 4),
- "task_scores": {task: round(score, 4) for task, score in task_scores.items()},
+ "task_scores": {
+ task: round(score, 4) for task, score in task_scores.items()
+ },
"summary": {
"total_tasks": len(task_scores),
- "total_samples": len(results_df)
- }
+ "total_samples": len(results_df),
+ },
}
-
- return overall_metrics
\ No newline at end of file
+
+ return overall_metrics
diff --git a/benchmark/longbench/__init__.py b/benchmark/longbench/__init__.py
index f757aa48..fa3bebde 100644
--- a/benchmark/longbench/__init__.py
+++ b/benchmark/longbench/__init__.py
@@ -8,4 +8,4 @@
from .calculate_metrics import calculate_metrics, calculate_metrics_e
from .longbench import LongBench
-__all__ = ["calculate_metrics", "calculate_metrics_e", "LongBench"]
\ No newline at end of file
+__all__ = ["calculate_metrics", "calculate_metrics_e", "LongBench"]
diff --git a/benchmark/longbench/calculate_metrics.py b/benchmark/longbench/calculate_metrics.py
index 05fb322c..8c491bc5 100644
--- a/benchmark/longbench/calculate_metrics.py
+++ b/benchmark/longbench/calculate_metrics.py
@@ -43,7 +43,12 @@ def scorer_e(dataset, predictions, answers, lengths, all_classes):
if dataset in ["trec", "triviaqa", "samsum", "lsht"]:
prediction = prediction.lstrip("\n").split("\n")[0]
for ground_truth in ground_truths:
- score = max(score, dataset2metric[dataset](prediction, ground_truth, all_classes=all_classes))
+ score = max(
+ score,
+ dataset2metric[dataset](
+ prediction, ground_truth, all_classes=all_classes
+ ),
+ )
if length < 4000:
scores["0-4k"].append(score)
elif length < 8000:
@@ -62,7 +67,12 @@ def scorer(dataset, predictions, answers, all_classes):
if dataset in ["trec", "triviaqa", "samsum", "lsht"]:
prediction = prediction.lstrip("\n").split("\n")[0]
for ground_truth in ground_truths:
- score = max(score, dataset2metric[dataset](prediction, ground_truth, all_classes=all_classes))
+ score = max(
+ score,
+ dataset2metric[dataset](
+ prediction, ground_truth, all_classes=all_classes
+ ),
+ )
total_score += score
return round(100 * total_score / len(predictions), 2)
diff --git a/benchmark/longbench/create_huggingface_dataset.py b/benchmark/longbench/create_huggingface_dataset.py
index f7e2269c..84978b04 100644
--- a/benchmark/longbench/create_huggingface_dataset.py
+++ b/benchmark/longbench/create_huggingface_dataset.py
@@ -123,15 +123,27 @@
if task == "trec":
dataset = dataset.map(
- lambda x: {"input": question_template[task].format(input=x["input"].removesuffix("Type:"))}
+ lambda x: {
+ "input": question_template[task].format(
+ input=x["input"].removesuffix("Type:")
+ )
+ }
)
elif task == "triviaqa":
dataset = dataset.map(
- lambda x: {"input": question_template[task].format(input=x["input"].removesuffix("Answer:"))}
+ lambda x: {
+ "input": question_template[task].format(
+ input=x["input"].removesuffix("Answer:")
+ )
+ }
)
elif task == "samsum":
dataset = dataset.map(
- lambda x: {"input": question_template[task].format(input=x["input"].removesuffix("Summary:"))}
+ lambda x: {
+ "input": question_template[task].format(
+ input=x["input"].removesuffix("Summary:")
+ )
+ }
)
else:
dataset = dataset.map(lambda x: {"input": question_template[task].format(**x)})
@@ -170,15 +182,27 @@
if task == "trec":
dataset = dataset.map(
- lambda x: {"input": question_template[task].format(input=x["input"].removesuffix("Type:"))}
+ lambda x: {
+ "input": question_template[task].format(
+ input=x["input"].removesuffix("Type:")
+ )
+ }
)
elif task == "triviaqa":
dataset = dataset.map(
- lambda x: {"input": question_template[task].format(input=x["input"].removesuffix("Answer:"))}
+ lambda x: {
+ "input": question_template[task].format(
+ input=x["input"].removesuffix("Answer:")
+ )
+ }
)
elif task == "samsum":
dataset = dataset.map(
- lambda x: {"input": question_template[task].format(input=x["input"].removesuffix("Summary:"))}
+ lambda x: {
+ "input": question_template[task].format(
+ input=x["input"].removesuffix("Summary:")
+ )
+ }
)
else:
dataset = dataset.map(lambda x: {"input": question_template[task].format(**x)})
@@ -194,4 +218,6 @@
# Push to hub
dataset = Dataset.from_pandas(df)
- dataset.push_to_hub("xAlg-AI/att-hub-longbench", config_name=f"{task}_e", split="test")
+ dataset.push_to_hub(
+ "xAlg-AI/att-hub-longbench", config_name=f"{task}_e", split="test"
+ )
diff --git a/benchmark/longbench/longbench.py b/benchmark/longbench/longbench.py
index 081fa6fa..d2edc2ef 100644
--- a/benchmark/longbench/longbench.py
+++ b/benchmark/longbench/longbench.py
@@ -33,36 +33,64 @@ class LongBench(Benchmark):
# All available LongBench datasets (22 standard + 13 extended)
all_datasets: List[str] = [
# Standard LongBench datasets
- "narrativeqa", "qasper", "multifieldqa_en", "multifieldqa_zh",
- "hotpotqa", "2wikimqa", "musique", "dureader", "gov_report",
- "qmsum", "multi_news", "vcsum", "trec", "triviaqa", "samsum",
- "lsht", "passage_count", "passage_retrieval_en",
- "passage_retrieval_zh", "lcc", "repobench-p",
+ "narrativeqa",
+ "qasper",
+ "multifieldqa_en",
+ "multifieldqa_zh",
+ "hotpotqa",
+ "2wikimqa",
+ "musique",
+ "dureader",
+ "gov_report",
+ "qmsum",
+ "multi_news",
+ "vcsum",
+ "trec",
+ "triviaqa",
+ "samsum",
+ "lsht",
+ "passage_count",
+ "passage_retrieval_en",
+ "passage_retrieval_zh",
+ "lcc",
+ "repobench-p",
# LongBench-E (extended) datasets with longer contexts
- "qasper_e", "multifieldqa_en_e", "hotpotqa_e", "2wikimqa_e",
- "gov_report_e", "multi_news_e", "trec_e", "triviaqa_e",
- "samsum_e", "passage_count_e", "passage_retrieval_en_e",
- "lcc_e", "repobench-p_e"
+ "qasper_e",
+ "multifieldqa_en_e",
+ "hotpotqa_e",
+ "2wikimqa_e",
+ "gov_report_e",
+ "multi_news_e",
+ "trec_e",
+ "triviaqa_e",
+ "samsum_e",
+ "passage_count_e",
+ "passage_retrieval_en_e",
+ "lcc_e",
+ "repobench-p_e",
]
-
+
benchmark_name: str = "longbench"
huggingface_dataset_id: str = "Xnhyacinth/LongBench"
def _load_datasets(self) -> pd.DataFrame:
"""Load LongBench datasets by individual configs.
-
+
LongBench requires loading each subset as a separate config.
-
+
Returns:
Combined pandas DataFrame with all samples from subsets_to_run.
"""
print(f"Loading LongBench datasets: {self.subsets_to_run}")
dfs = []
-
+
for subset in self.subsets_to_run:
try:
from datasets import load_dataset
- subset_dataset = load_dataset(self.huggingface_dataset_id, subset, split="test")
+
+ subset_dataset = load_dataset(
+ self.huggingface_dataset_id, subset, split="test"
+ )
subset_df = subset_dataset.to_pandas()
subset_df["task"] = subset # Ensure task column exists
dfs.append(subset_df)
@@ -70,12 +98,13 @@ def _load_datasets(self) -> pd.DataFrame:
except Exception as subset_error:
print(f" โ Failed to load {subset}: {str(subset_error)}")
continue
-
+
if not dfs:
raise Exception("No LongBench subsets could be loaded successfully")
-
+
# Combine all subset DataFrames
import pandas as pd
+
combined_df = pd.concat(dfs, ignore_index=True)
print(f"Combined {len(combined_df)} total samples from {len(dfs)} subsets")
return combined_df
@@ -104,35 +133,39 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
task_scores: Dict[str, Any] = {}
standard_scores: List[float] = []
extended_scores: Dict[str, List[float]] = {"0-4k": [], "4-8k": [], "8k+": []}
-
+
for task_name, task_df in task_groups:
try:
if task_name.endswith("_e"):
# Extended dataset - use calculate_metrics_e for length-based breakdown
length_scores: Dict[str, float] = calculate_metrics_e(task_df)
task_scores[task_name] = length_scores
-
+
# Accumulate extended scores for overall calculation
for length_range in extended_scores:
if length_range in length_scores:
- extended_scores[length_range].append(length_scores[length_range])
+ extended_scores[length_range].append(
+ length_scores[length_range]
+ )
else:
# Standard dataset - use calculate_metrics
score: float = calculate_metrics(task_df)
task_scores[task_name] = score
standard_scores.append(score)
-
+
except Exception as e:
print(f"Error evaluating task {task_name}: {str(e)}")
task_scores[task_name] = {"error": str(e)}
# Compute overall metrics
overall_metrics: Dict[str, Any] = {"task_scores": task_scores}
-
+
# Overall score for standard datasets
if standard_scores:
- overall_metrics["standard_overall_score"] = round(sum(standard_scores) / len(standard_scores), 2)
-
+ overall_metrics["standard_overall_score"] = round(
+ sum(standard_scores) / len(standard_scores), 2
+ )
+
# Overall scores for extended datasets by length range
if any(extended_scores.values()):
extended_overall: Dict[str, float] = {}
@@ -141,21 +174,25 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
extended_overall[length_range] = round(sum(scores) / len(scores), 2)
if extended_overall:
overall_metrics["extended_overall_scores"] = extended_overall
-
+
# Compute grand overall score if we have both standard and extended results
all_scores: List[float] = standard_scores.copy()
for scores in extended_scores.values():
all_scores.extend(scores)
-
+
if all_scores:
- overall_metrics["overall_score"] = round(sum(all_scores) / len(all_scores), 2)
-
+ overall_metrics["overall_score"] = round(
+ sum(all_scores) / len(all_scores), 2
+ )
+
# Add summary statistics
overall_metrics["summary"] = {
"total_tasks": len(task_scores),
- "standard_tasks": len([t for t in task_scores.keys() if not t.endswith("_e")]),
+ "standard_tasks": len(
+ [t for t in task_scores.keys() if not t.endswith("_e")]
+ ),
"extended_tasks": len([t for t in task_scores.keys() if t.endswith("_e")]),
- "total_samples": len(results_df)
+ "total_samples": len(results_df),
}
-
- return overall_metrics
\ No newline at end of file
+
+ return overall_metrics
diff --git a/benchmark/longbenchv2/__init__.py b/benchmark/longbenchv2/__init__.py
index b829ae7c..087f8627 100644
--- a/benchmark/longbenchv2/__init__.py
+++ b/benchmark/longbenchv2/__init__.py
@@ -8,4 +8,4 @@
from .calculate_metrics import calculate_metrics
from .longbenchv2 import LongBenchv2
-__all__ = ["calculate_metrics", "LongBenchv2"]
\ No newline at end of file
+__all__ = ["calculate_metrics", "LongBenchv2"]
diff --git a/benchmark/longbenchv2/calculate_metrics.py b/benchmark/longbenchv2/calculate_metrics.py
index f85eebb3..71e91710 100644
--- a/benchmark/longbenchv2/calculate_metrics.py
+++ b/benchmark/longbenchv2/calculate_metrics.py
@@ -29,7 +29,9 @@ def scorer(predictions, answers, lengths, difficulties):
compensated = False
easy, hard, short, medium, long = 0, 0, 0, 0, 0
easy_acc, hard_acc, short_acc, medium_acc, long_acc = 0, 0, 0, 0, 0
- for pred, answer, length, difficulty in zip(predictions, answers, lengths, difficulties):
+ for pred, answer, length, difficulty in zip(
+ predictions, answers, lengths, difficulties
+ ):
acc = int(extract_answer(pred) == answer)
if compensated and pred["pred"] is None:
acc = 0.25 # type:ignore[assignment]
diff --git a/benchmark/longbenchv2/create_huggingface_dataset.py b/benchmark/longbenchv2/create_huggingface_dataset.py
index 95c4b45a..e261342a 100644
--- a/benchmark/longbenchv2/create_huggingface_dataset.py
+++ b/benchmark/longbenchv2/create_huggingface_dataset.py
@@ -31,7 +31,9 @@
# Longbench-v2
for task in ["0shot", "cot"]:
dataset = load_dataset("THUDM/LongBench-v2", split="train")
- dataset = dataset.map(lambda x: {"context": context_prefix[task].format(context=x["context"].strip())})
+ dataset = dataset.map(
+ lambda x: {"context": context_prefix[task].format(context=x["context"].strip())}
+ )
dataset = dataset.map(
lambda x: {
"question": question_template[task].format(
diff --git a/benchmark/longbenchv2/longbenchv2.py b/benchmark/longbenchv2/longbenchv2.py
index 01f56f61..566338f6 100644
--- a/benchmark/longbenchv2/longbenchv2.py
+++ b/benchmark/longbenchv2/longbenchv2.py
@@ -26,29 +26,29 @@ class LongBenchv2(Benchmark):
"""
# All available LongBenchv2 datasets
- all_datasets: List[str] = [
- "0shot",
- "cot"
- ]
-
+ all_datasets: List[str] = ["0shot", "cot"]
+
benchmark_name: str = "longbenchv2"
huggingface_dataset_id: str = "Xnhyacinth/LongBench-v2"
def _load_datasets(self) -> pd.DataFrame:
"""Load LongBenchv2 datasets by individual configs.
-
+
LongBenchv2 requires loading each subset as a separate config.
-
+
Returns:
Combined pandas DataFrame with all samples from subsets_to_run.
"""
print(f"Loading LongBenchv2 datasets: {self.subsets_to_run}")
dfs = []
-
+
for subset in self.subsets_to_run:
try:
from datasets import load_dataset
- subset_dataset = load_dataset(self.huggingface_dataset_id, subset, split="test")
+
+ subset_dataset = load_dataset(
+ self.huggingface_dataset_id, subset, split="test"
+ )
subset_df = subset_dataset.to_pandas()
subset_df["task"] = subset # Ensure task column exists
dfs.append(subset_df)
@@ -56,12 +56,13 @@ def _load_datasets(self) -> pd.DataFrame:
except Exception as subset_error:
print(f" โ Failed to load {subset}: {str(subset_error)}")
continue
-
+
if not dfs:
raise Exception("No LongBenchv2 subsets could be loaded successfully")
-
+
# Combine all subset DataFrames
import pandas as pd
+
combined_df = pd.concat(dfs, ignore_index=True)
print(f"Combined {len(combined_df)} total samples from {len(dfs)} subsets")
return combined_df
@@ -90,20 +91,20 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
task_groups = results_df.groupby("task")
task_scores: Dict[str, Any] = {}
all_scores: List[float] = []
-
+
for task_name, task_df in task_groups:
try:
# Use the calculate_metrics function from HashAttention evaluation
scores = calculate_metrics(task_df)
-
+
# Parse the scores string to extract overall accuracy
if len(scores) >= 2:
- overall_score_str = scores[1].split('\t')[0]
+ overall_score_str = scores[1].split("\t")[0]
try:
overall_score = float(overall_score_str)
task_scores[task_name] = {
"overall_accuracy": overall_score,
- "detailed_scores": scores
+ "detailed_scores": scores,
}
all_scores.append(overall_score)
print(f" โ {task_name}: {overall_score:.1f}%")
@@ -111,21 +112,21 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
task_scores[task_name] = {"error": "Could not parse score"}
else:
task_scores[task_name] = {"error": "No scores returned"}
-
+
except Exception as e:
print(f"Error evaluating task {task_name}: {str(e)}")
task_scores[task_name] = {"error": str(e)}
# Compute overall metrics
overall_score = sum(all_scores) / len(all_scores) if all_scores else 0.0
-
+
overall_metrics: Dict[str, Any] = {
"overall_score": round(overall_score, 1),
"task_scores": task_scores,
"summary": {
"total_tasks": len(task_scores),
- "total_samples": len(results_df)
- }
+ "total_samples": len(results_df),
+ },
}
-
- return overall_metrics
\ No newline at end of file
+
+ return overall_metrics
diff --git a/benchmark/loogle/calculate_metrics.py b/benchmark/loogle/calculate_metrics.py
index 7deb02bd..ab3bdd48 100644
--- a/benchmark/loogle/calculate_metrics.py
+++ b/benchmark/loogle/calculate_metrics.py
@@ -76,7 +76,10 @@ def wrapped_metric(answer, predicted_answer):
try:
return metric_fn(answer, predicted_answer)
except Exception as e:
- print(f"Cannot calculate metric: {e}" f" on answer:{answer} and predicted_answer:{predicted_answer}")
+ print(
+ f"Cannot calculate metric: {e}"
+ f" on answer:{answer} and predicted_answer:{predicted_answer}"
+ )
return {key: 0.0 for key in metric_fn("Hi there", "hi there")}
return wrapped_metric
@@ -93,21 +96,36 @@ def calculate_metrics(df: pd.DataFrame) -> dict:
("exact", get_exact_match),
("partial", get_partial_match),
]:
- match, count = zip(*df_task.apply(lambda x: metric_fn(x["answer"], x["predicted_answer"]), axis=1))
+ match, count = zip(
+ *df_task.apply(
+ lambda x: metric_fn(x["answer"], x["predicted_answer"]), axis=1
+ )
+ )
scores[task][f"{prefix}_match"] = round(sum(match) / sum(count), 4)
else:
- df["predicted_answer"] = df["predicted_answer"].apply(lambda x: x if isinstance(x, str) else "")
+ df["predicted_answer"] = df["predicted_answer"].apply(
+ lambda x: x if isinstance(x, str) else ""
+ )
for metric_fn in [get_bleu_score, get_rouge_score, get_meteor_score]: # type: ignore
metric_fn = try_except_metric(metric_fn)
- metric_scores = [metric_fn(row["answer"], row["predicted_answer"]) for _, row in df_task.iterrows()]
+ metric_scores = [
+ metric_fn(row["answer"], row["predicted_answer"])
+ for _, row in df_task.iterrows()
+ ]
scores[task].update(pd.DataFrame(metric_scores).mean().to_dict())
# BERT scores (batched)
scores[task]["bert"] = (
- score(df_task["answer"].to_list(), df_task["predicted_answer"].to_list(), lang="EN")[1].mean().item()
+ score(
+ df_task["answer"].to_list(),
+ df_task["predicted_answer"].to_list(),
+ lang="EN",
+ )[1]
+ .mean()
+ .item()
)
return scores
diff --git a/benchmark/loogle/create_huggingface_dataset.py b/benchmark/loogle/create_huggingface_dataset.py
index ef2fd15b..78eb732e 100644
--- a/benchmark/loogle/create_huggingface_dataset.py
+++ b/benchmark/loogle/create_huggingface_dataset.py
@@ -30,19 +30,30 @@
}
# Source: https://github.com/bigai-nlco/LooGLE/blob/main/config/task2maxlen.json
-max_new_tokens = {"shortdep_qa": 300, "longdep_qa": 500, "longdep_summarization": 500, "shortdep_cloze": 50}
+max_new_tokens = {
+ "shortdep_qa": 300,
+ "longdep_qa": 500,
+ "longdep_summarization": 500,
+ "shortdep_cloze": 50,
+}
for task in ["shortdep_qa", "longdep_qa", "shortdep_cloze", "longdep_summarization"]:
- df = load_dataset("bigainlco/LooGLE", task, split="test", trust_remote_code=True).to_pandas()
+ df = load_dataset(
+ "bigainlco/LooGLE", task, split="test", trust_remote_code=True
+ ).to_pandas()
if task == "longdep_summarization":
df["question"] = ""
df = df.rename(columns={"output": "answer", "input": "context"})
else:
- df["qa_pairs"] = df["qa_pairs"].apply(lambda x: eval(x) if x != "none" else [{"Q": "", "A": "", "S": [""]}])
+ df["qa_pairs"] = df["qa_pairs"].apply(
+ lambda x: eval(x) if x != "none" else [{"Q": "", "A": "", "S": [""]}]
+ )
df = df.explode("qa_pairs")
- df = pd.concat([df.drop(["qa_pairs"], axis=1), df["qa_pairs"].apply(pd.Series)], axis=1)
+ df = pd.concat(
+ [df.drop(["qa_pairs"], axis=1), df["qa_pairs"].apply(pd.Series)], axis=1
+ )
df = df.rename(columns={"A": "answer", "Q": "question", "input": "context"})
if task == "shortdep_cloze":
df["answer"] = df["answer"].apply(json.dumps, ensure_ascii=False)
diff --git a/benchmark/loogle/loogle.py b/benchmark/loogle/loogle.py
index d66b77ba..9c311a9b 100644
--- a/benchmark/loogle/loogle.py
+++ b/benchmark/loogle/loogle.py
@@ -31,27 +31,30 @@ class Loogle(Benchmark):
"shortdep_qa",
"longdep_qa",
"shortdep_cloze",
- "longdep_summarization"
+ "longdep_summarization",
]
-
+
benchmark_name: str = "loogle"
huggingface_dataset_id: str = "simonjegou/loogle"
def _load_datasets(self) -> pd.DataFrame:
"""Load Loogle datasets by individual configs.
-
+
Loogle requires loading each subset as a separate config.
-
+
Returns:
Combined pandas DataFrame with all samples from subsets_to_run.
"""
print(f"Loading Loogle datasets: {self.subsets_to_run}")
dfs = []
-
+
for subset in self.subsets_to_run:
try:
from datasets import load_dataset
- subset_dataset = load_dataset(self.huggingface_dataset_id, subset, split="test")
+
+ subset_dataset = load_dataset(
+ self.huggingface_dataset_id, subset, split="test"
+ )
subset_df = subset_dataset.to_pandas()
subset_df["task"] = subset # Ensure task column exists
dfs.append(subset_df)
@@ -59,12 +62,13 @@ def _load_datasets(self) -> pd.DataFrame:
except Exception as subset_error:
print(f" โ Failed to load {subset}: {str(subset_error)}")
continue
-
+
if not dfs:
raise Exception("No Loogle subsets could be loaded successfully")
-
+
# Combine all subset DataFrames
import pandas as pd
+
combined_df = pd.concat(dfs, ignore_index=True)
print(f"Combined {len(combined_df)} total samples from {len(dfs)} subsets")
return combined_df
@@ -88,7 +92,7 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
# Use the calculate_metrics function from HashAttention evaluation
task_scores: Dict[str, Dict[str, float]] = calculate_metrics(results_df)
-
+
# Compute overall score by averaging task scores
all_scores: List[float] = []
for task, scores in task_scores.items():
@@ -101,16 +105,16 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
# Fallback to first available metric
elif scores:
all_scores.append(list(scores.values())[0])
-
+
overall_score = sum(all_scores) / len(all_scores) if all_scores else 0.0
-
+
overall_metrics: Dict[str, Any] = {
"overall_score": round(overall_score, 4),
"task_scores": task_scores,
"summary": {
"total_tasks": len(task_scores),
- "total_samples": len(results_df)
- }
+ "total_samples": len(results_df),
+ },
}
-
- return overall_metrics
\ No newline at end of file
+
+ return overall_metrics
diff --git a/benchmark/mock_benchmark/__init__.py b/benchmark/mock_benchmark/__init__.py
index 0774219f..09d3ed56 100644
--- a/benchmark/mock_benchmark/__init__.py
+++ b/benchmark/mock_benchmark/__init__.py
@@ -1,9 +1,7 @@
-
-
"""
Mock benchmark module for testing and demonstration purposes.
"""
from .mock_benchmark import MockBenchmark
-__all__ = ["MockBenchmark"]
\ No newline at end of file
+__all__ = ["MockBenchmark"]
diff --git a/benchmark/mock_benchmark/mock_benchmark.py b/benchmark/mock_benchmark/mock_benchmark.py
index a069b919..2f0108d0 100644
--- a/benchmark/mock_benchmark/mock_benchmark.py
+++ b/benchmark/mock_benchmark/mock_benchmark.py
@@ -27,11 +27,13 @@ class MockBenchmark(Benchmark):
# Class attributes required by base Benchmark class
all_datasets: List[str] = ["reading_comprehension"]
benchmark_name: str = "mock_benchmark"
- huggingface_dataset_id: str = "mock/dataset" # Not actually used since we override _load_datasets
+ huggingface_dataset_id: str = (
+ "mock/dataset" # Not actually used since we override _load_datasets
+ )
def _load_datasets(self) -> pd.DataFrame:
"""Load mock dataset with hardcoded samples.
-
+
Returns:
DataFrame containing 5 mock samples with context, question, and answers.
"""
@@ -51,7 +53,6 @@ def _load_datasets(self) -> pd.DataFrame:
in their leaves, where specialized organelles called chloroplasts contain the
chlorophyll necessary for the process.
""".strip(),
-
"history_context": """
The Renaissance was a period of cultural, artistic, and intellectual revival
that began in Italy during the 14th century and spread throughout Europe.
@@ -67,7 +68,6 @@ def _load_datasets(self) -> pd.DataFrame:
laid the foundation for the Scientific Revolution and the Enlightenment that
would follow.
""".strip(),
-
"geography_context": """
The Amazon rainforest, often called the "lungs of the Earth," is the world's
largest tropical rainforest. Located primarily in Brazil, it also extends
@@ -81,7 +81,7 @@ def _load_datasets(self) -> pd.DataFrame:
drainage basin. Unfortunately, the rainforest faces threats from deforestation,
mining, and climate change, making its conservation critical for the health
of our planet.
- """.strip()
+ """.strip(),
}
# Create sample data - 5 samples total, 2 sharing the same context
@@ -89,46 +89,50 @@ def _load_datasets(self) -> pd.DataFrame:
{
"context": contexts["science_context"],
"question": "What are the two main stages of photosynthesis?",
- "answers": ["light-dependent reactions and light-independent reactions",
- "light-dependent reactions and Calvin cycle"],
- "task": "reading_comprehension"
+ "answers": [
+ "light-dependent reactions and light-independent reactions",
+ "light-dependent reactions and Calvin cycle",
+ ],
+ "task": "reading_comprehension",
},
{
"context": contexts["science_context"], # Same context as sample 1
"question": "What gas is produced as a byproduct of photosynthesis?",
"answers": ["oxygen"],
- "task": "reading_comprehension"
+ "task": "reading_comprehension",
},
{
"context": contexts["history_context"],
"question": "In which century did the Renaissance begin?",
"answers": ["14th century", "14th"],
- "task": "reading_comprehension"
+ "task": "reading_comprehension",
},
{
"context": contexts["geography_context"],
"question": "Why is the Amazon rainforest called the 'lungs of the Earth'?",
- "answers": ["because it absorbs carbon dioxide and produces oxygen",
- "it regulates global climate by absorbing CO2 and producing oxygen"],
- "task": "reading_comprehension"
+ "answers": [
+ "because it absorbs carbon dioxide and produces oxygen",
+ "it regulates global climate by absorbing CO2 and producing oxygen",
+ ],
+ "task": "reading_comprehension",
},
{
"context": contexts["geography_context"],
"question": "Which river flows through the Amazon rainforest?",
"answers": ["Amazon River", "the Amazon River"],
- "task": "reading_comprehension"
- }
+ "task": "reading_comprehension",
+ },
]
# Convert to DataFrame
df: pd.DataFrame = pd.DataFrame(sample_data)
-
+
# Add sample IDs for tracking
df["sample_id"] = range(1, len(df) + 1)
-
+
print(f"Loaded {len(df)} mock samples")
print(f"Unique contexts: {df['context'].nunique()}")
-
+
return df
def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
@@ -159,32 +163,38 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
predicted_answer: str = str(row["predicted_answer"]).strip().lower()
ground_truth_answers: List[str] = row["answers"]
sample_id: int = row["sample_id"]
-
+
# Check if prediction matches any ground truth answer
is_correct: bool = False
for gt_answer in ground_truth_answers:
gt_answer_normalized: str = str(gt_answer).strip().lower()
-
+
# Check exact match or substring match
- if (predicted_answer == gt_answer_normalized or
- gt_answer_normalized in predicted_answer):
+ if (
+ predicted_answer == gt_answer_normalized
+ or gt_answer_normalized in predicted_answer
+ ):
is_correct = True
break
-
+
if is_correct:
correct_predictions += 1
-
- sample_scores.append({
- "sample_id": sample_id,
- "question": row["question"],
- "predicted_answer": row["predicted_answer"],
- "ground_truth": ground_truth_answers,
- "correct": is_correct
- })
+
+ sample_scores.append(
+ {
+ "sample_id": sample_id,
+ "question": row["question"],
+ "predicted_answer": row["predicted_answer"],
+ "ground_truth": ground_truth_answers,
+ "correct": is_correct,
+ }
+ )
# Calculate metrics
- accuracy: float = correct_predictions / total_samples if total_samples > 0 else 0.0
-
+ accuracy: float = (
+ correct_predictions / total_samples if total_samples > 0 else 0.0
+ )
+
metrics: Dict[str, Any] = {
"accuracy": round(accuracy, 3),
"correct_predictions": correct_predictions,
@@ -194,8 +204,8 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
"benchmark": self.benchmark_name,
"task": "reading_comprehension",
"unique_contexts": results_df["context"].nunique(),
- "evaluation_method": "exact_match_and_substring"
- }
+ "evaluation_method": "exact_match_and_substring",
+ },
}
- return metrics
\ No newline at end of file
+ return metrics
diff --git a/benchmark/ruler/calculate_metrics.py b/benchmark/ruler/calculate_metrics.py
index 16d1b924..4fa78ec9 100644
--- a/benchmark/ruler/calculate_metrics.py
+++ b/benchmark/ruler/calculate_metrics.py
@@ -9,7 +9,12 @@
def string_match_part(preds, refs):
score = (
- sum([max([1.0 if r.lower() in pred.lower() else 0.0 for r in ref]) for pred, ref in zip(preds, refs)])
+ sum(
+ [
+ max([1.0 if r.lower() in pred.lower() else 0.0 for r in ref])
+ for pred, ref in zip(preds, refs)
+ ]
+ )
/ len(preds)
* 100
)
@@ -19,7 +24,10 @@ def string_match_part(preds, refs):
def string_match_all(preds, refs):
score = (
sum(
- [sum([1.0 if r.lower() in pred.lower() else 0.0 for r in ref]) / len(ref) for pred, ref in zip(preds, refs)]
+ [
+ sum([1.0 if r.lower() in pred.lower() else 0.0 for r in ref]) / len(ref)
+ for pred, ref in zip(preds, refs)
+ ]
)
/ len(preds)
* 100
@@ -31,7 +39,9 @@ def calculate_metrics(df: pd.DataFrame) -> dict:
scores = {}
np_pattern = re.compile(r"[\x00-\x1f]")
- df["predicted_answer"] = df["predicted_answer"].apply(lambda x: np_pattern.sub("", x.strip()).strip())
+ df["predicted_answer"] = df["predicted_answer"].apply(
+ lambda x: np_pattern.sub("", x.strip()).strip()
+ )
for task, df_task in df.groupby("task"):
task_category = task.split("_")[0]
diff --git a/benchmark/ruler/create_huggingface_dataset.py b/benchmark/ruler/create_huggingface_dataset.py
index b31f14bf..74628967 100644
--- a/benchmark/ruler/create_huggingface_dataset.py
+++ b/benchmark/ruler/create_huggingface_dataset.py
@@ -13,7 +13,9 @@
QUESTION_PATTERNS = {
"niah": re.compile(r"What (?:is|are all) the special magic"),
"vt": re.compile(r"Question: Find all variables that are assigned the value"),
- "cwe": re.compile(r"Question: What are the 10 most common words in the above list\?"),
+ "cwe": re.compile(
+ r"Question: What are the 10 most common words in the above list\?"
+ ),
"fwe": re.compile(r"Question: Do not provide any explanation\."),
"qa": re.compile(r"Answer the question based on the given documents\."),
}
@@ -41,7 +43,9 @@ def get_dataframe(path):
Parse the data from the provided path and return a DataFrame with the context, question, answers and task
"""
- assert re.match(r".*\/\d+$", str(path)), "The path should must ends with the context length (e.g. with /4096)"
+ assert re.match(
+ r".*\/\d+$", str(path)
+ ), "The path should must ends with the context length (e.g. with /4096)"
df_list = []
for task_path in Path(path).glob("**/*.jsonl"):
@@ -60,21 +64,27 @@ def split_context_question(text):
question, answer = qa[:idx], qa[idx:]
return context, question, answer
- df["context"], df["question"], df["answer_prefix"] = zip(*df["input"].apply(split_context_question))
+ df["context"], df["question"], df["answer_prefix"] = zip(
+ *df["input"].apply(split_context_question)
+ )
df["task"] = task
df["max_new_tokens"] = MAX_NEW_TOKENS[task.split("_")[0]]
df_list.append(df)
# Concatenate all the dataframes
df = pd.concat(df_list)
- df = df[["context", "question", "answer_prefix", "outputs", "task", "max_new_tokens"]]
+ df = df[
+ ["context", "question", "answer_prefix", "outputs", "task", "max_new_tokens"]
+ ]
df = df.rename(columns={"outputs": "answer"}).reset_index(drop=True)
return df
if __name__ == "__main__":
- data_dir = Path("/mnt/workspace/projects/RULER/scripts/data/data/") # output of the generate.sh script
+ data_dir = Path(
+ "/mnt/workspace/projects/RULER/scripts/data/data/"
+ ) # output of the generate.sh script
repo_id = "simonjegou/ruler"
# Loop over all the context lengths
diff --git a/benchmark/ruler/ruler.py b/benchmark/ruler/ruler.py
index afe770d5..1035c666 100644
--- a/benchmark/ruler/ruler.py
+++ b/benchmark/ruler/ruler.py
@@ -27,31 +27,29 @@ class Ruler(Benchmark):
"""
# All available Ruler datasets (context lengths)
- all_datasets: List[str] = [
- "4096",
- "8192",
- "16384",
- "32768"
- ]
-
+ all_datasets: List[str] = ["4096", "8192", "16384", "32768"]
+
benchmark_name: str = "ruler"
huggingface_dataset_id: str = "simonjegou/ruler"
def _load_datasets(self) -> pd.DataFrame:
"""Load Ruler datasets by individual configs.
-
+
Ruler requires loading each context length as a separate config.
-
+
Returns:
Combined pandas DataFrame with all samples from subsets_to_run.
"""
print(f"Loading Ruler datasets: {self.subsets_to_run}")
dfs = []
-
+
for subset in self.subsets_to_run:
try:
from datasets import load_dataset
- subset_dataset = load_dataset(self.huggingface_dataset_id, subset, split="test")
+
+ subset_dataset = load_dataset(
+ self.huggingface_dataset_id, subset, split="test"
+ )
subset_df = subset_dataset.to_pandas()
# Add context length as a column for analysis
subset_df["context_length"] = subset
@@ -60,12 +58,13 @@ def _load_datasets(self) -> pd.DataFrame:
except Exception as subset_error:
print(f" โ Failed to load {subset}: {str(subset_error)}")
continue
-
+
if not dfs:
raise Exception("No Ruler subsets could be loaded successfully")
-
+
# Combine all subset DataFrames
import pandas as pd
+
combined_df = pd.concat(dfs, ignore_index=True)
print(f"Combined {len(combined_df)} total samples from {len(dfs)} subsets")
return combined_df
@@ -91,15 +90,15 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
# Use the calculate_metrics function from HashAttention evaluation
task_scores: Dict[str, Dict[str, float]] = calculate_metrics(results_df)
-
+
# Extract string_match scores and compute overall
all_scores: List[float] = []
for task, scores in task_scores.items():
if "string_match" in scores:
all_scores.append(scores["string_match"])
-
+
overall_score = sum(all_scores) / len(all_scores) if all_scores else 0.0
-
+
# Group by context length if available
context_length_scores: Dict[str, float] = {}
if "context_length" in results_df.columns:
@@ -107,9 +106,19 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
length_df = results_df[results_df["context_length"] == context_length]
if len(length_df) > 0:
length_scores = calculate_metrics(length_df)
- length_overall = sum(score.get("string_match", 0) for score in length_scores.values()) / len(length_scores) if length_scores else 0.0
- context_length_scores[str(context_length)] = round(length_overall, 2)
-
+ length_overall = (
+ sum(
+ score.get("string_match", 0)
+ for score in length_scores.values()
+ )
+ / len(length_scores)
+ if length_scores
+ else 0.0
+ )
+ context_length_scores[str(context_length)] = round(
+ length_overall, 2
+ )
+
overall_metrics: Dict[str, Any] = {
"overall_score": round(overall_score, 2),
"task_scores": task_scores,
@@ -117,8 +126,8 @@ def post_run_evaluate(self, results_df: pd.DataFrame) -> Dict[str, Any]:
"summary": {
"total_tasks": len(task_scores),
"total_samples": len(results_df),
- "context_lengths": list(context_length_scores.keys())
- }
+ "context_lengths": list(context_length_scores.keys()),
+ },
}
-
- return overall_metrics
\ No newline at end of file
+
+ return overall_metrics
diff --git a/benchmark/scripts/benchmark.py b/benchmark/scripts/benchmark.py
index 289fa0f1..e69e0802 100644
--- a/benchmark/scripts/benchmark.py
+++ b/benchmark/scripts/benchmark.py
@@ -19,9 +19,12 @@
from benchmark.executor import BenchmarkExecutor
from benchmark.executor_config import BenchmarkConfig, AdapterConfig
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+from sparse_attention_hub.sparse_attention.research_attention import (
+ ResearchAttentionConfig,
+)
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
- LocalMaskerConfig, SinkMaskerConfig
+ LocalMaskerConfig,
+ SinkMaskerConfig,
)
# ============================================================================
@@ -29,80 +32,63 @@
# ============================================================================
# GPU Configuration
-GPUS = [0,2,7] # Use all available GPUs
+GPUS = [0, 2, 7] # Use all available GPUs
MAX_CONCURRENT_RUNS = 3 # One per GPU
# Model List
MODELS = [
- "microsoft/Phi-4-mini-instruct",
- "meta-llama/Llama-3.2-1B-Instruct",
+ "microsoft/Phi-4-mini-instruct",
+ "meta-llama/Llama-3.2-1B-Instruct",
]
# Sparse Attention Configurations
SPARSE_CONFIGS = [
# Dense baseline (no sparse attention)
("dense", None),
-
# StreamingLLM configurations
- ("streaming_conservative", ResearchAttentionConfig(masker_configs=[
- SinkMaskerConfig(sink_size=4),
- LocalMaskerConfig(window_size=16)
- ])),
+ (
+ "streaming_conservative",
+ ResearchAttentionConfig(
+ masker_configs=[
+ SinkMaskerConfig(sink_size=4),
+ LocalMaskerConfig(window_size=16),
+ ]
+ ),
+ ),
]
# Benchmark List
# 1. InfiniteBench - using passkey task
infinite_bench_config = BenchmarkConfig(
- benchmark_name="infinite_bench",
- subsets=["passkey"]
+ benchmark_name="infinite_bench", subsets=["passkey"]
)
# 2. Ruler - using 4096 context length
-ruler_config = BenchmarkConfig(
- benchmark_name="ruler",
- subsets=["4096"]
-)
+ruler_config = BenchmarkConfig(benchmark_name="ruler", subsets=["4096"])
# 3. Loogle - using shortdep_qa task
-loogle_config = BenchmarkConfig(
- benchmark_name="loogle",
- subsets=["shortdep_qa"]
-)
+loogle_config = BenchmarkConfig(benchmark_name="loogle", subsets=["shortdep_qa"])
# 4. ZeroScrolls - using gov_report task
zero_scrolls_config = BenchmarkConfig(
- benchmark_name="zero_scrolls",
- subsets=["default"]
+ benchmark_name="zero_scrolls", subsets=["default"]
)
# 5. LongBenchv2 - using 0shot task
-longbenchv2_config = BenchmarkConfig(
- benchmark_name="longbenchv2",
- subsets=["0shot"]
-)
+longbenchv2_config = BenchmarkConfig(benchmark_name="longbenchv2", subsets=["0shot"])
# 6. AIME2024 - using single task
-aime2024_config = BenchmarkConfig(
- benchmark_name="aime2024",
- subsets=["aime2024"]
-)
+aime2024_config = BenchmarkConfig(benchmark_name="aime2024", subsets=["aime2024"])
# 7. AIME2025 - using single task
-aime2025_config = BenchmarkConfig(
- benchmark_name="aime2025",
- subsets=["aime2025"]
-)
+aime2025_config = BenchmarkConfig(benchmark_name="aime2025", subsets=["aime2025"])
# 8. LongBench (existing) - using narrativeqa task
-longbench_config = BenchmarkConfig(
- benchmark_name="longbench",
- subsets=["narrativeqa"]
-)
+longbench_config = BenchmarkConfig(benchmark_name="longbench", subsets=["narrativeqa"])
# 9. Mock Benchmark (existing) - using single task
mock_benchmark_config = BenchmarkConfig(
- benchmark_name="mock_benchmark",
- subsets=["reading_comprehension"]
+ benchmark_name="mock_benchmark", subsets=["reading_comprehension"]
)
# List of all sample configurations
@@ -115,7 +101,7 @@
aime2024_config,
aime2025_config,
longbench_config,
- mock_benchmark_config
+ mock_benchmark_config,
]
@@ -127,7 +113,7 @@
},
tokenizer_kwargs={
"padding_side": "left",
- }
+ },
)
# Generation Parameters
@@ -157,7 +143,7 @@
if __name__ == "__main__":
print("๐ Starting Minimalistic Benchmark Suite")
print("=" * 50)
-
+
print(f"๐ง Configuration:")
print(f" - GPUs: {GPUS}")
print(f" - Models: {len(MODELS)}")
@@ -174,25 +160,29 @@
print(f" - Benchmarks: {len(BENCHMARKS)}")
for i, benchmark in enumerate(BENCHMARKS, 1):
if benchmark.subsets:
- print(f" {i}. {benchmark.benchmark_name}: {len(benchmark.subsets)} subsets")
+ print(
+ f" {i}. {benchmark.benchmark_name}: {len(benchmark.subsets)} subsets"
+ )
else:
print(f" {i}. {benchmark.benchmark_name}: all subsets")
print(f" - Max concurrent: {MAX_CONCURRENT_RUNS}")
print(f" - Result dir: {RESULT_DIR}")
print(f" - Resumability: {'enabled' if ENABLE_RESUMABILITY else 'disabled'}")
-
+
# Calculate total combinations
total_models = len(MODELS)
total_configs = len(SPARSE_CONFIGS)
total_benchmarks = sum(len(b.subsets) if b.subsets else 1 for b in BENCHMARKS)
total_combinations = total_models * total_configs * total_benchmarks
-
+
print(f"\n๐ Experiment Matrix: {total_combinations} total combinations")
print(f" - Models: {total_models}")
print(f" - Sparse configs: {total_configs}")
print(f" - Benchmark-subsets: {total_benchmarks}")
- print(f" - Estimated time: {total_combinations * TIMEOUT_PER_BENCHMARK / 3600:.1f} hours (worst case)")
-
+ print(
+ f" - Estimated time: {total_combinations * TIMEOUT_PER_BENCHMARK / 3600:.1f} hours (worst case)"
+ )
+
# Create executor
print(f"\n๐ง Initializing BenchmarkExecutor...")
executor = BenchmarkExecutor(
@@ -202,9 +192,9 @@
enable_resumability=ENABLE_RESUMABILITY,
required_result_files=["raw_results.csv"],
timeout_per_benchmark=TIMEOUT_PER_BENCHMARK,
- verbose=True
+ verbose=True,
)
-
+
# Run benchmarks
print(f"\n๐ฏ Running Benchmark Matrix...")
try:
@@ -214,9 +204,9 @@
benchmark_configs=BENCHMARKS,
adapter_config=ADAPTER_CONFIG,
generation_kwargs=GENERATION_KWARGS,
- request_kwargs=REQUEST_KWARGS
+ request_kwargs=REQUEST_KWARGS,
)
-
+
# Print summary
print(f"\nโ
Benchmark Execution Completed!")
print(f" - Total: {results.progress.total_stubs}")
@@ -224,10 +214,10 @@
print(f" - Failed: {results.progress.failed_stubs}")
print(f" - Skipped: {results.progress.skipped_stubs}")
print(f" - Results saved to: {RESULT_DIR}")
-
+
except KeyboardInterrupt:
print(f"\nโ ๏ธ Interrupted by user")
print(f" Partial results in: {RESULT_DIR}")
except Exception as e:
print(f"\nโ Execution failed: {e}")
- raise
+ raise
diff --git a/benchmark/scripts/executor_example.py b/benchmark/scripts/executor_example.py
index 092e4467..835581e9 100644
--- a/benchmark/scripts/executor_example.py
+++ b/benchmark/scripts/executor_example.py
@@ -20,12 +20,19 @@
from benchmark.executor import BenchmarkExecutor
from benchmark.executor_config import BenchmarkConfig, AdapterConfig
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+from sparse_attention_hub.sparse_attention.research_attention import (
+ ResearchAttentionConfig,
+)
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
- LocalMaskerConfig, SinkMaskerConfig, OracleTopKConfig, OracleTopPMaskerConfig
+ LocalMaskerConfig,
+ SinkMaskerConfig,
+ OracleTopKConfig,
+ OracleTopPMaskerConfig,
)
from sparse_attention_hub.sparse_attention.research_attention.maskers.sampling.implementations import (
- AdaptiveSamplingMaskerConfig, RandomSamplingMaskerConfig, MagicPigConfig
+ AdaptiveSamplingMaskerConfig,
+ RandomSamplingMaskerConfig,
+ MagicPigConfig,
)
# ============================================================================
@@ -40,121 +47,140 @@
# Model List
MODELS = [
- "meta-llama/Llama-3.1-8B-Instruct",
+ "meta-llama/Llama-3.1-8B-Instruct",
]
# Sparse Attention Configurations
SPARSE_CONFIGS = [
# Dense baseline (no sparse attention)
("dense", None),
-
# StreamingLLM configurations
- ("streaming_conservative", ResearchAttentionConfig(masker_configs=[
- SinkMaskerConfig(sink_size=128),
- LocalMaskerConfig(window_size=INTENDED_SPARSITY)
- ])),
- #Oracle-TopK
- ("streaming_oracle_topk", ResearchAttentionConfig(masker_configs=[
- SinkMaskerConfig(sink_size=128),
- LocalMaskerConfig(window_size=128),
- OracleTopKConfig(heavy_size=INTENDED_SPARSITY)
- ])),
+ (
+ "streaming_conservative",
+ ResearchAttentionConfig(
+ masker_configs=[
+ SinkMaskerConfig(sink_size=128),
+ LocalMaskerConfig(window_size=INTENDED_SPARSITY),
+ ]
+ ),
+ ),
+ # Oracle-TopK
+ (
+ "streaming_oracle_topk",
+ ResearchAttentionConfig(
+ masker_configs=[
+ SinkMaskerConfig(sink_size=128),
+ LocalMaskerConfig(window_size=128),
+ OracleTopKConfig(heavy_size=INTENDED_SPARSITY),
+ ]
+ ),
+ ),
# Oracle-TopP
- ("streaming_oracle_topp", ResearchAttentionConfig(masker_configs=[
- SinkMaskerConfig(sink_size=128),
- LocalMaskerConfig(window_size=128),
- OracleTopPMaskerConfig(top_p=0.85)
- ])),
+ (
+ "streaming_oracle_topp",
+ ResearchAttentionConfig(
+ masker_configs=[
+ SinkMaskerConfig(sink_size=128),
+ LocalMaskerConfig(window_size=128),
+ OracleTopPMaskerConfig(top_p=0.85),
+ ]
+ ),
+ ),
# Adaptive Sampling
- ("streaming_adaptive_sampling", ResearchAttentionConfig(masker_configs=[
- SinkMaskerConfig(sink_size=128),
- LocalMaskerConfig(window_size=128),
- OracleTopKConfig(heavy_size=128),
- AdaptiveSamplingMaskerConfig(base_rate_sampling=0.05, epsilon=0.25, delta=0.25, init_offset=128, local_offset=128)
- ])),
- # Random Sampling
- ("streaming_random_sampling", ResearchAttentionConfig(masker_configs=[
- SinkMaskerConfig(sink_size=128),
- LocalMaskerConfig(window_size=128),
- RandomSamplingMaskerConfig(sampling_rate=0.1)
- ])),
- # MagicPig
- ("streaming_magicpig", ResearchAttentionConfig(masker_configs=[
- SinkMaskerConfig(sink_size=128),
- LocalMaskerConfig(window_size=128),
- MagicPigConfig(lsh_l=8, lsh_k=8)
- ])),
+ (
+ "streaming_adaptive_sampling",
+ ResearchAttentionConfig(
+ masker_configs=[
+ SinkMaskerConfig(sink_size=128),
+ LocalMaskerConfig(window_size=128),
+ OracleTopKConfig(heavy_size=128),
+ AdaptiveSamplingMaskerConfig(
+ base_rate_sampling=0.05,
+ epsilon=0.25,
+ delta=0.25,
+ init_offset=128,
+ local_offset=128,
+ ),
+ ]
+ ),
+ ),
+ # Random Sampling
+ (
+ "streaming_random_sampling",
+ ResearchAttentionConfig(
+ masker_configs=[
+ SinkMaskerConfig(sink_size=128),
+ LocalMaskerConfig(window_size=128),
+ RandomSamplingMaskerConfig(sampling_rate=0.1),
+ ]
+ ),
+ ),
+ # MagicPig
+ (
+ "streaming_magicpig",
+ ResearchAttentionConfig(
+ masker_configs=[
+ SinkMaskerConfig(sink_size=128),
+ LocalMaskerConfig(window_size=128),
+ MagicPigConfig(lsh_l=8, lsh_k=8),
+ ]
+ ),
+ ),
]
# Benchmark List
# 1. InfiniteBench - using passkey task
infinite_bench_config = BenchmarkConfig(
- benchmark_name="infinite_bench",
- subsets=["passkey"]
+ benchmark_name="infinite_bench", subsets=["passkey"]
)
# 2. Ruler - using 4096 context length
-ruler_config = BenchmarkConfig(
- benchmark_name="ruler",
- subsets=["4096"]
-)
+ruler_config = BenchmarkConfig(benchmark_name="ruler", subsets=["4096"])
# 3. Loogle - using shortdep_qa task
loogle_config = BenchmarkConfig(
benchmark_name="loogle",
subsets=["shortdep_qa"],
- #subsets=["longdep_qa"],
- #subsets=["shortdep_cloze"],
- #subsets=["longdep_summarization"],
+ # subsets=["longdep_qa"],
+ # subsets=["shortdep_cloze"],
+ # subsets=["longdep_summarization"],
)
# 4. ZeroScrolls - using gov_report task
zero_scrolls_config = BenchmarkConfig(
- benchmark_name="zero_scrolls",
- subsets=["default"]
+ benchmark_name="zero_scrolls", subsets=["default"]
)
# 5. LongBenchv2 - using 0shot task
-longbenchv2_config = BenchmarkConfig(
- benchmark_name="longbenchv2",
- subsets=["0shot"]
-)
+longbenchv2_config = BenchmarkConfig(benchmark_name="longbenchv2", subsets=["0shot"])
# 6. AIME2024 - using single task
-aime2024_config = BenchmarkConfig(
- benchmark_name="aime2024",
- subsets=["aime2024"]
-)
+aime2024_config = BenchmarkConfig(benchmark_name="aime2024", subsets=["aime2024"])
# 7. AIME2025 - using single task
-aime2025_config = BenchmarkConfig(
- benchmark_name="aime2025",
- subsets=["aime2025"]
-)
+aime2025_config = BenchmarkConfig(benchmark_name="aime2025", subsets=["aime2025"])
# 8. LongBench (existing) - using narrativeqa task
longbench_config = BenchmarkConfig(
- benchmark_name="longbench",
- subsets=["passage_retrieval_en"]
+ benchmark_name="longbench", subsets=["passage_retrieval_en"]
)
# 9. Mock Benchmark (existing) - using single task
mock_benchmark_config = BenchmarkConfig(
- benchmark_name="mock_benchmark",
- subsets=["reading_comprehension"]
+ benchmark_name="mock_benchmark", subsets=["reading_comprehension"]
)
# List of all sample configurations
BENCHMARKS = [
- #infinite_bench_config,
- #ruler_config,
+ # infinite_bench_config,
+ # ruler_config,
loogle_config,
- #zero_scrolls_config,
- #longbenchv2_config,
- #aime2024_config,
- #aime2025_config,
- #longbench_config,
- #mock_benchmark_config
+ # zero_scrolls_config,
+ # longbenchv2_config,
+ # aime2024_config,
+ # aime2025_config,
+ # longbench_config,
+ # mock_benchmark_config
]
@@ -167,7 +193,7 @@
},
tokenizer_kwargs={
"padding_side": "left",
- }
+ },
)
# Generation Parameters
@@ -196,7 +222,7 @@
if __name__ == "__main__":
print("๐ Starting Minimalistic Benchmark Suite")
print("=" * 50)
-
+
print(f"๐ง Configuration:")
print(f" - GPUs: {GPUS}")
print(f" - Models: {len(MODELS)}")
@@ -213,25 +239,29 @@
print(f" - Benchmarks: {len(BENCHMARKS)}")
for i, benchmark in enumerate(BENCHMARKS, 1):
if benchmark.subsets:
- print(f" {i}. {benchmark.benchmark_name}: {len(benchmark.subsets)} subsets")
+ print(
+ f" {i}. {benchmark.benchmark_name}: {len(benchmark.subsets)} subsets"
+ )
else:
print(f" {i}. {benchmark.benchmark_name}: all subsets")
print(f" - Max concurrent: {MAX_CONCURRENT_RUNS}")
print(f" - Result dir: {RESULT_DIR}")
print(f" - Resumability: {'enabled' if ENABLE_RESUMABILITY else 'disabled'}")
-
+
# Calculate total combinations
total_models = len(MODELS)
total_configs = len(SPARSE_CONFIGS)
total_benchmarks = sum(len(b.subsets) if b.subsets else 1 for b in BENCHMARKS)
total_combinations = total_models * total_configs * total_benchmarks
-
+
print(f"\n๐ Experiment Matrix: {total_combinations} total combinations")
print(f" - Models: {total_models}")
print(f" - Sparse configs: {total_configs}")
print(f" - Benchmark-subsets: {total_benchmarks}")
- print(f" - Estimated time: {total_combinations * TIMEOUT_PER_BENCHMARK / 3600:.1f} hours (worst case)")
-
+ print(
+ f" - Estimated time: {total_combinations * TIMEOUT_PER_BENCHMARK / 3600:.1f} hours (worst case)"
+ )
+
# Create executor
print(f"\n๐ง Initializing BenchmarkExecutor...")
executor = BenchmarkExecutor(
@@ -241,9 +271,9 @@
enable_resumability=ENABLE_RESUMABILITY,
required_result_files=["raw_results.csv"],
timeout_per_benchmark=TIMEOUT_PER_BENCHMARK,
- verbose=True
+ verbose=True,
)
-
+
# Run benchmarks
print(f"\n๐ฏ Running Benchmark Matrix...")
try:
@@ -253,9 +283,9 @@
benchmark_configs=BENCHMARKS,
adapter_config=ADAPTER_CONFIG,
generation_kwargs=GENERATION_KWARGS,
- request_kwargs=REQUEST_KWARGS
+ request_kwargs=REQUEST_KWARGS,
)
-
+
# Print summary
print(f"\nโ
Benchmark Execution Completed!")
print(f" - Total: {results.progress.total_stubs}")
@@ -263,10 +293,10 @@
print(f" - Failed: {results.progress.failed_stubs}")
print(f" - Skipped: {results.progress.skipped_stubs}")
print(f" - Results saved to: {RESULT_DIR}")
-
+
except KeyboardInterrupt:
print(f"\nโ ๏ธ Interrupted by user")
print(f" Partial results in: {RESULT_DIR}")
except Exception as e:
print(f"\nโ Execution failed: {e}")
- raise
+ raise
diff --git a/benchmark/scripts/full_benchmarking/full_benchmark.py b/benchmark/scripts/full_benchmarking/full_benchmark.py
index d37d4b7b..1fe318bf 100644
--- a/benchmark/scripts/full_benchmarking/full_benchmark.py
+++ b/benchmark/scripts/full_benchmarking/full_benchmark.py
@@ -42,7 +42,9 @@
from sparse_attention_hub.adapters.model_servers.huggingface import ModelServerHF
## spin a model server ###
-model_server = ModelServerHF(ModelServerConfig(enable_stats_logging=True, delete_on_zero_reference=False))
+model_server = ModelServerHF(
+ ModelServerConfig(enable_stats_logging=True, delete_on_zero_reference=False)
+)
# ============================================================================
# CONFIGURATION
@@ -503,7 +505,7 @@ def get_random_sampling_config_name(sink_size, window_size, sampling_rate):
),
]
-#SPARSE_CONFIGS = [("dense", None)]
+# SPARSE_CONFIGS = [("dense", None)]
# ==========================================================================
# Benchmark List
@@ -578,7 +580,7 @@ def get_random_sampling_config_name(sink_size, window_size, sampling_rate):
# mock_benchmark_config,
# niah1, niah2, niah3, cwe, fwe, vt, qa1, qa2, multikey1, multikey2, multikey3, multikey, multivalue
cwe,
- niah1
+ niah1,
]
diff --git a/benchmark/scripts/full_benchmarking/generate_results_table.py b/benchmark/scripts/full_benchmarking/generate_results_table.py
index 701ace74..8f4fbc51 100755
--- a/benchmark/scripts/full_benchmarking/generate_results_table.py
+++ b/benchmark/scripts/full_benchmarking/generate_results_table.py
@@ -22,40 +22,40 @@
def find_dataset_files(base_folder: str, dataset_name: str) -> List[Path]:
"""
Find all files corresponding to the given dataset in the base folder.
-
+
Args:
base_folder: Path to the base folder containing benchmark results
dataset_name: Name of the dataset to search for
-
+
Returns:
List of paths to dataset directories
"""
base_path = Path(base_folder)
dataset_paths = []
-
+
# Walk through the directory structure
for model_dir in base_path.iterdir():
if not model_dir.is_dir():
continue
-
+
for config_dir in model_dir.iterdir():
if not config_dir.is_dir():
continue
-
+
dataset_dir = config_dir / dataset_name
if dataset_dir.exists():
dataset_paths.append(dataset_dir)
-
+
return dataset_paths
def extract_config_name(dataset_path: Path) -> str:
"""
Extract the configuration name from the dataset path.
-
+
Args:
dataset_path: Path to the dataset directory
-
+
Returns:
Configuration name
"""
@@ -66,49 +66,51 @@ def extract_config_name(dataset_path: Path) -> str:
def parse_metrics_json(metrics_file: Path) -> Dict[str, Any]:
"""
Parse the metrics.json file to extract macro metrics.
-
+
Args:
metrics_file: Path to the metrics.json file
-
+
Returns:
Dictionary containing the parsed metrics
"""
try:
- with open(metrics_file, 'r') as f:
+ with open(metrics_file, "r") as f:
data = json.load(f)
-
+
metrics = {}
-
+
# Extract overall score
- if 'overall_score' in data:
- metrics['overall_score'] = data['overall_score']
-
+ if "overall_score" in data:
+ metrics["overall_score"] = data["overall_score"]
+
# Extract task-specific scores
- if 'task_scores' in data:
- for task_name, task_metrics in data['task_scores'].items():
+ if "task_scores" in data:
+ for task_name, task_metrics in data["task_scores"].items():
for metric_name, metric_value in task_metrics.items():
# Create a unique metric name if there are multiple tasks
- if len(data['task_scores']) > 1:
+ if len(data["task_scores"]) > 1:
metric_key = f"{task_name}_{metric_name}"
else:
metric_key = metric_name
metrics[metric_key] = metric_value
-
+
return metrics
-
+
except Exception as e:
print(f"Warning: Could not parse {metrics_file}: {e}")
return {}
-def parse_micro_metrics_jsonl(micro_metrics_file: Path, max_lines: int = 1000) -> Dict[str, float]:
+def parse_micro_metrics_jsonl(
+ micro_metrics_file: Path, max_lines: int = 1000
+) -> Dict[str, float]:
"""
Parse the micro_metrics.jsonl file to extract average density and attention error.
-
+
Args:
micro_metrics_file: Path to the micro_metrics.jsonl file
max_lines: Maximum number of lines to read
-
+
Returns:
Dictionary containing average density and attention error
"""
@@ -116,37 +118,37 @@ def parse_micro_metrics_jsonl(micro_metrics_file: Path, max_lines: int = 1000) -
density_values = []
error_values = []
line_count = 0
-
- with open(micro_metrics_file, 'r') as f:
+
+ with open(micro_metrics_file, "r") as f:
for line in f:
if line_count >= max_lines:
break
-
+
try:
data = json.loads(line.strip())
- metric_name = data.get('metric', '')
- value = data.get('value', 0.0)
-
- if metric_name == 'research_attention_density':
+ metric_name = data.get("metric", "")
+ value = data.get("value", 0.0)
+
+ if metric_name == "research_attention_density":
density_values.append(value)
- elif metric_name == 'research_attention_output_error':
+ elif metric_name == "research_attention_output_error":
error_values.append(value)
-
+
except json.JSONDecodeError:
continue
-
+
line_count += 1
-
+
micro_metrics = {}
-
+
if density_values:
- micro_metrics['avg_density'] = sum(density_values) / len(density_values)
-
+ micro_metrics["avg_density"] = sum(density_values) / len(density_values)
+
if error_values:
- micro_metrics['avg_attention_error'] = sum(error_values) / len(error_values)
-
+ micro_metrics["avg_attention_error"] = sum(error_values) / len(error_values)
+
return micro_metrics
-
+
except Exception as e:
print(f"Warning: Could not parse {micro_metrics_file}: {e}")
return {}
@@ -155,97 +157,97 @@ def parse_micro_metrics_jsonl(micro_metrics_file: Path, max_lines: int = 1000) -
def get_all_available_metrics(dataset_paths: List[Path]) -> List[str]:
"""
Get all available metrics across all dataset paths.
-
+
Args:
dataset_paths: List of paths to dataset directories
-
+
Returns:
List of all available metric names
"""
all_metrics = set()
-
+
for dataset_path in dataset_paths:
metrics_file = dataset_path / "metrics.json"
if metrics_file.exists():
metrics = parse_metrics_json(metrics_file)
all_metrics.update(metrics.keys())
-
+
micro_metrics_file = dataset_path / "micro_metrics.jsonl"
if micro_metrics_file.exists():
micro_metrics = parse_micro_metrics_jsonl(micro_metrics_file)
all_metrics.update(micro_metrics.keys())
-
+
return sorted(list(all_metrics))
def generate_results_table(dataset_name: str, base_folder: str) -> pd.DataFrame:
"""
Generate a results table from benchmark data.
-
+
Args:
dataset_name: Name of the dataset to analyze
base_folder: Path to the base folder containing benchmark results
-
+
Returns:
DataFrame with configs as rows and metrics as columns
"""
print(f"๐ Finding dataset files for '{dataset_name}' in '{base_folder}'...")
dataset_paths = find_dataset_files(base_folder, dataset_name)
-
+
if not dataset_paths:
print(f"โ No dataset files found for '{dataset_name}' in '{base_folder}'")
return pd.DataFrame()
-
+
print(f"โ
Found {len(dataset_paths)} dataset directories")
-
+
# Get all available metrics
print("๐ Discovering available metrics...")
all_metrics = get_all_available_metrics(dataset_paths)
print(f"โ
Found {len(all_metrics)} metrics: {', '.join(all_metrics)}")
-
+
# Initialize results dictionary
results = []
-
+
print("๐ Processing each configuration...")
for dataset_path in dataset_paths:
config_name = extract_config_name(dataset_path)
print(f" Processing: {config_name}")
-
- row_data = {'config': config_name}
-
+
+ row_data = {"config": config_name}
+
# Parse macro metrics
metrics_file = dataset_path / "metrics.json"
if metrics_file.exists():
macro_metrics = parse_metrics_json(metrics_file)
row_data.update(macro_metrics)
-
+
# Parse micro metrics
micro_metrics_file = dataset_path / "micro_metrics.jsonl"
if micro_metrics_file.exists():
micro_metrics = parse_micro_metrics_jsonl(micro_metrics_file)
row_data.update(micro_metrics)
-
+
results.append(row_data)
-
+
# Create DataFrame
df = pd.DataFrame(results)
-
+
if df.empty:
print("โ No results found")
return df
-
+
# Reorder columns to put config first, then metrics
- if 'config' in df.columns:
- other_cols = [col for col in df.columns if col != 'config']
- df = df[['config'] + other_cols]
-
+ if "config" in df.columns:
+ other_cols = [col for col in df.columns if col != "config"]
+ df = df[["config"] + other_cols]
+
return df
def save_results_table(df: pd.DataFrame, dataset_name: str, base_folder: str) -> None:
"""
Save the results table to files.
-
+
Args:
df: DataFrame containing the results
dataset_name: Name of the dataset
@@ -254,21 +256,21 @@ def save_results_table(df: pd.DataFrame, dataset_name: str, base_folder: str) ->
if df.empty:
print("โ No data to save")
return
-
+
# Create output directory
output_dir = Path(base_folder) / "results_tables"
output_dir.mkdir(exist_ok=True)
-
+
# Save as CSV
csv_file = output_dir / f"{dataset_name}_results.csv"
df.to_csv(csv_file, index=False)
print(f"๐พ Saved CSV results to: {csv_file}")
-
+
# Save as Excel
excel_file = output_dir / f"{dataset_name}_results.xlsx"
df.to_excel(excel_file, index=False)
print(f"๐พ Saved Excel results to: {excel_file}")
-
+
# Print summary
print(f"\n๐ Results Summary:")
print(f" - Configurations: {len(df)}")
@@ -285,43 +287,42 @@ def main():
Examples:
python generate_results_table.py loogle_shortdep_cloze full_benchmark.matrix
python generate_results_table.py longbench_passage_retrieval_en ./results/
- """
+ """,
)
-
+
parser.add_argument(
"dataset_name",
- help="Name of the dataset to analyze (e.g., loogle_shortdep_cloze)"
+ help="Name of the dataset to analyze (e.g., loogle_shortdep_cloze)",
)
-
+
parser.add_argument(
- "base_folder",
- help="Path to the base folder containing benchmark results"
+ "base_folder", help="Path to the base folder containing benchmark results"
)
-
+
args = parser.parse_args()
-
+
print("๐ Starting Results Table Generation")
print("=" * 50)
print(f"Dataset: {args.dataset_name}")
print(f"Base folder: {args.base_folder}")
print()
-
+
# Generate the results table
df = generate_results_table(args.dataset_name, args.base_folder)
-
+
if not df.empty:
# Display the table
print("\n๐ Results Table:")
print("=" * 80)
print(df.to_string(index=False))
-
+
# Save the results
save_results_table(df, args.dataset_name, args.base_folder)
-
+
print(f"\nโ
Results table generation completed successfully!")
else:
print(f"\nโ No results found for dataset '{args.dataset_name}'")
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/benchmark/scripts/micro_benchmarking/analyse_stress_tests.py b/benchmark/scripts/micro_benchmarking/analyse_stress_tests.py
index ba3e3444..8529d271 100644
--- a/benchmark/scripts/micro_benchmarking/analyse_stress_tests.py
+++ b/benchmark/scripts/micro_benchmarking/analyse_stress_tests.py
@@ -20,7 +20,7 @@
def parse_config_name(config_name: str) -> Dict[str, float]:
"""Parse configuration name to extract parameters.
-
+
Args:
config_name: Configuration name like "adaptive_sampling.sink_0.001_window_0.001_heavy_0.005_base_0.01_epsilon_0.01_delta_0.01"
or "oracle_top_k_0.5.sink_0.02_window_0.02"
@@ -28,16 +28,16 @@ def parse_config_name(config_name: str) -> Dict[str, float]:
or "hashattention.sink_0.001_window_0.001_top_k_0.17"
or "adaptive_sampling_hat.sink_0.01_window_0.01_heavy_0.02_base_0.01_epsilon_0.5_delta_0.5"
or "random_sampling.sink_0.001_window_0.001_sampling_rate_0.01"
-
+
Returns:
Dictionary with parsed parameters
"""
# Extract parameters using regex for different configuration types
-
+
# Pattern for adaptive_sampling_hat (must come before adaptive_sampling to avoid conflicts)
adaptive_hat_pattern = r"adaptive_sampling_hat\.sink_([\d.]+)_window_([\d.]+)_heavy_([\d.]+)_base_([\d.]+)_epsilon_([\d.]+)_delta_([\d.]+)"
adaptive_hat_match = re.match(adaptive_hat_pattern, config_name)
-
+
if adaptive_hat_match:
return {
"config_type": "adaptive_sampling_hat",
@@ -46,13 +46,13 @@ def parse_config_name(config_name: str) -> Dict[str, float]:
"heavy_size": float(adaptive_hat_match.group(3)),
"base_rate_sampling": float(adaptive_hat_match.group(4)),
"epsilon": float(adaptive_hat_match.group(5)),
- "delta": float(adaptive_hat_match.group(6))
+ "delta": float(adaptive_hat_match.group(6)),
}
-
+
# Pattern for adaptive_sampling
adaptive_pattern = r"adaptive_sampling\.sink_([\d.]+)_window_([\d.]+)_heavy_([\d.]+)_base_([\d.]+)_epsilon_([\d.]+)_delta_([\d.]+)"
adaptive_match = re.match(adaptive_pattern, config_name)
-
+
if adaptive_match:
return {
"config_type": "adaptive_sampling",
@@ -61,97 +61,103 @@ def parse_config_name(config_name: str) -> Dict[str, float]:
"heavy_size": float(adaptive_match.group(3)),
"base_rate_sampling": float(adaptive_match.group(4)),
"epsilon": float(adaptive_match.group(5)),
- "delta": float(adaptive_match.group(6))
+ "delta": float(adaptive_match.group(6)),
}
-
+
# Pattern for oracle_top_k
top_k_pattern = r"oracle_top_k_([\d.]+)\.sink_([\d.]+)_window_([\d.]+)"
top_k_match = re.match(top_k_pattern, config_name)
-
+
if top_k_match:
return {
"config_type": "oracle_top_k",
"top_k": float(top_k_match.group(1)),
"sink_size": float(top_k_match.group(2)),
- "window_size": float(top_k_match.group(3))
+ "window_size": float(top_k_match.group(3)),
}
-
+
# Pattern for oracle_top_p
top_p_pattern = r"oracle_top_p_([\d.]+)\.sink_([\d.]+)_window_([\d.]+)"
top_p_match = re.match(top_p_pattern, config_name)
-
+
if top_p_match:
return {
"config_type": "oracle_top_p",
"top_p": float(top_p_match.group(1)),
"sink_size": float(top_p_match.group(2)),
- "window_size": float(top_p_match.group(3))
+ "window_size": float(top_p_match.group(3)),
}
-
+
# Pattern for hashattention
- hashattention_pattern = r"hashattention\.sink_([\d.]+)_window_([\d.]+)_top_k_([\d.]+)"
+ hashattention_pattern = (
+ r"hashattention\.sink_([\d.]+)_window_([\d.]+)_top_k_([\d.]+)"
+ )
hashattention_match = re.match(hashattention_pattern, config_name)
-
+
if hashattention_match:
return {
"config_type": "hashattention",
"sink_size": float(hashattention_match.group(1)),
"window_size": float(hashattention_match.group(2)),
- "hat_top_k": float(hashattention_match.group(3))
+ "hat_top_k": float(hashattention_match.group(3)),
}
-
+
# Pattern for random_sampling
- random_sampling_pattern = r"random_sampling\.sink_([\d.]+)_window_([\d.]+)_sampling_rate_([\d.]+)"
+ random_sampling_pattern = (
+ r"random_sampling\.sink_([\d.]+)_window_([\d.]+)_sampling_rate_([\d.]+)"
+ )
random_sampling_match = re.match(random_sampling_pattern, config_name)
-
+
if random_sampling_match:
return {
"config_type": "random_sampling",
"sink_size": float(random_sampling_match.group(1)),
"window_size": float(random_sampling_match.group(2)),
- "sampling_rate": float(random_sampling_match.group(3))
+ "sampling_rate": float(random_sampling_match.group(3)),
}
-
+
# If no pattern matches, return empty dict
return {"config_type": "unknown"}
def load_config_file(config_path: Path) -> Dict[str, Any]:
"""Load configuration from JSON file.
-
+
Args:
config_path: Path to config.json file
-
+
Returns:
Configuration dictionary
"""
- with open(config_path, 'r') as f:
+ with open(config_path, "r") as f:
return json.load(f)
def load_micro_metrics(metrics_path: Path) -> List[Dict[str, Any]]:
"""Load micro metrics from JSONL file.
-
+
Args:
metrics_path: Path to micro_metrics.jsonl file
-
+
Returns:
List of metric entries
"""
metrics = []
- with open(metrics_path, 'r') as f:
+ with open(metrics_path, "r") as f:
for line in f:
if line.strip():
metrics.append(json.loads(line))
return metrics
-def process_experiment_directory(exp_dir: Path) -> List[Tuple[List[Dict[str, Any]], Dict[str, Any], str]]:
+def process_experiment_directory(
+ exp_dir: Path,
+) -> List[Tuple[List[Dict[str, Any]], Dict[str, Any], str]]:
"""Process a single experiment directory.
-
+
Args:
exp_dir: Path to experiment directory
-
+
Returns:
List of tuples (metrics_data, config_data, dataset_name) for each benchmark directory
"""
@@ -159,45 +165,45 @@ def process_experiment_directory(exp_dir: Path) -> List[Tuple[List[Dict[str, Any
benchmark_dirs = [d for d in exp_dir.iterdir() if d.is_dir()]
if not benchmark_dirs:
return []
-
+
results = []
-
+
for benchmark_dir in benchmark_dirs:
dataset_name = benchmark_dir.name
-
+
# Load configuration
config_path = benchmark_dir / "config.json"
if not config_path.exists():
continue
-
+
config = load_config_file(config_path)
-
+
# Load micro metrics
metrics_path = benchmark_dir / "micro_metrics.jsonl"
if not metrics_path.exists():
continue
-
+
metrics = load_micro_metrics(metrics_path)
-
+
results.append((metrics, config, dataset_name))
-
+
return results
def extract_sparse_config_params(config: Dict[str, Any]) -> Dict[str, Any]:
"""Extract sparse attention configuration parameters.
-
+
Args:
config: Configuration dictionary
-
+
Returns:
Dictionary with sparse attention parameters
"""
sparse_config = config.get("sparse_attention_config", {})
masker_configs = sparse_config.get("masker_configs", [])
-
+
params = {}
-
+
# Extract parameters from masker configs
for masker_config in masker_configs:
if "sink_size" in masker_config:
@@ -226,55 +232,61 @@ def extract_sparse_config_params(config: Dict[str, Any]) -> Dict[str, Any]:
elif "sampling_rate" in masker_config:
# Random sampling parameters
params["sampling_rate"] = masker_config["sampling_rate"]
-
+
return params
-def organize_metrics_by_layer(metrics: List[Dict[str, Any]]) -> Dict[int, Dict[str, float]]:
+def organize_metrics_by_layer(
+ metrics: List[Dict[str, Any]]
+) -> Dict[int, Dict[str, float]]:
"""Organize metrics by layer index and average multiple measurements.
-
+
Args:
metrics: List of metric entries
-
+
Returns:
Dictionary mapping layer_idx to averaged metrics
"""
layer_metrics = {}
-
+
# First pass: collect all values for each layer
for metric in metrics:
layer_idx = metric.get("metadata", {}).get("layer_idx")
if layer_idx is None:
continue
-
+
if layer_idx not in layer_metrics:
layer_metrics[layer_idx] = {"density": [], "error": []}
-
+
metric_name = metric.get("metric")
value = metric.get("value")
-
+
if metric_name == "research_attention_density":
layer_metrics[layer_idx]["density"].append(value)
elif metric_name == "research_attention_output_error":
layer_metrics[layer_idx]["error"].append(value)
-
+
# Second pass: average the collected values
averaged_metrics = {}
for layer_idx, values in layer_metrics.items():
averaged_metrics[layer_idx] = {}
-
+
if values["density"]:
- averaged_metrics[layer_idx]["density"] = sum(values["density"]) / len(values["density"])
-
+ averaged_metrics[layer_idx]["density"] = sum(values["density"]) / len(
+ values["density"]
+ )
+
if values["error"]:
- averaged_metrics[layer_idx]["error"] = sum(values["error"]) / len(values["error"])
-
+ averaged_metrics[layer_idx]["error"] = sum(values["error"]) / len(
+ values["error"]
+ )
+
return averaged_metrics
def analyze_stress_tests(results_dir: str, output_dir: str) -> None:
"""Analyze stress test results and generate TSV files.
-
+
Args:
results_dir: Path to stress_test_adaptive.matrix directory
output_dir: Output directory for TSV files
@@ -282,39 +294,39 @@ def analyze_stress_tests(results_dir: str, output_dir: str) -> None:
results_path = Path(results_dir)
output_path = Path(output_dir)
output_path.mkdir(exist_ok=True)
-
+
# Find model directories
model_dirs = [d for d in results_path.iterdir() if d.is_dir()]
-
+
all_vector_data = []
all_metadata = []
-
+
for model_dir in model_dirs:
model_name = model_dir.name
-
+
# Find configuration directories
config_dirs = [d for d in model_dir.iterdir() if d.is_dir()]
-
+
for config_dir in config_dirs:
config_name = config_dir.name
-
+
# Parse configuration name
parsed_params = parse_config_name(config_name)
-
+
# Process experiment directory (now returns list of results for each benchmark)
benchmark_results = process_experiment_directory(config_dir)
-
+
if not benchmark_results:
continue
-
+
# Process each benchmark result
for metrics, config, dataset_name in benchmark_results:
# Extract sparse attention parameters
sparse_params = extract_sparse_config_params(config)
-
+
# Organize metrics by layer
layer_metrics = organize_metrics_by_layer(metrics)
-
+
# Generate vector data
for layer_idx, layer_data in layer_metrics.items():
if "density" in layer_data and "error" in layer_data:
@@ -323,10 +335,10 @@ def analyze_stress_tests(results_dir: str, output_dir: str) -> None:
"config": config_name,
"layer_idx": layer_idx,
"density": layer_data["density"],
- "error": layer_data["error"]
+ "error": layer_data["error"],
}
all_vector_data.append(vector_entry)
-
+
# Generate metadata entry
metadata_entry = {
"model": model_name,
@@ -334,46 +346,44 @@ def analyze_stress_tests(results_dir: str, output_dir: str) -> None:
"dataset": dataset_name,
"layer_idx": "all", # This will be expanded for each layer
**parsed_params,
- **sparse_params
+ **sparse_params,
}
-
+
# Add metadata for each layer
for layer_idx in layer_metrics.keys():
layer_metadata = metadata_entry.copy()
layer_metadata["layer_idx"] = layer_idx
all_metadata.append(layer_metadata)
-
+
# Write vector.tsv
vector_path = output_path / "vector.tsv"
metadata_path = output_path / "metadata.tsv"
-
+
# Write vector data
- with open(vector_path, 'w') as f:
+ with open(vector_path, "w") as f:
f.write("density\terror\n")
for entry in all_vector_data:
f.write(f"{entry['density']}\t{entry['error']}\n")
-
+
# Write metadata
- with open(metadata_path, 'w') as f:
+ with open(metadata_path, "w") as f:
if all_metadata:
# Get all unique keys from all metadata entries
all_keys = set()
for entry in all_metadata:
all_keys.update(entry.keys())
-
+
# Sort keys for consistent output
sorted_keys = sorted(all_keys)
-
+
# Write header
f.write("\t".join(sorted_keys) + "\n")
-
+
# Write data
for entry in all_metadata:
row = [str(entry.get(key, "")) for key in sorted_keys]
f.write("\t".join(row) + "\n")
-
-
print(f"Analysis complete!")
print(f"Vector data written to: {vector_path}")
print(f"Metadata written to: {metadata_path}")
@@ -385,25 +395,25 @@ def main():
"""Main function."""
parser = argparse.ArgumentParser(description="Analyze stress test results")
parser.add_argument(
- "--results-dir",
+ "--results-dir",
default="./stress_test_adaptive.matrix",
- help="Path to stress test results directory"
+ help="Path to stress test results directory",
)
parser.add_argument(
- "--output-dir",
+ "--output-dir",
default="./sparse-attention-hub-share/docs/micro_tests/",
- help="Output directory for TSV files"
+ help="Output directory for TSV files",
)
-
+
args = parser.parse_args()
-
+
# Check if results directory exists
if not os.path.exists(args.results_dir):
print(f"Error: Results directory '{args.results_dir}' does not exist")
return
-
+
analyze_stress_tests(args.results_dir, args.output_dir)
if __name__ == "__main__":
- main()
+ main()
diff --git a/benchmark/scripts/micro_benchmarking/plot_stress_test_results.py b/benchmark/scripts/micro_benchmarking/plot_stress_test_results.py
index 001bcc8e..1fa50339 100644
--- a/benchmark/scripts/micro_benchmarking/plot_stress_test_results.py
+++ b/benchmark/scripts/micro_benchmarking/plot_stress_test_results.py
@@ -20,29 +20,31 @@
def load_tsv_data(vector_path: str, metadata_path: str) -> pd.DataFrame:
"""Load and merge vector and metadata TSV files.
-
+
Args:
vector_path: Path to vector.tsv file
metadata_path: Path to metadata.tsv file
-
+
Returns:
Merged DataFrame with all data
"""
# Load the data
- vector_df = pd.read_csv(vector_path, sep='\t')
- metadata_df = pd.read_csv(metadata_path, sep='\t')
-
+ vector_df = pd.read_csv(vector_path, sep="\t")
+ metadata_df = pd.read_csv(metadata_path, sep="\t")
+
# Merge the dataframes
# Since both files have the same number of rows in the same order,
# we can simply concatenate them
merged_df = pd.concat([vector_df, metadata_df], axis=1)
-
+
return merged_df
-def create_interactive_scatter_plot(df: pd.DataFrame, output_path: str, dataset_name: str = "") -> None:
+def create_interactive_scatter_plot(
+ df: pd.DataFrame, output_path: str, dataset_name: str = ""
+) -> None:
"""Create an interactive scatter plot of error vs density.
-
+
Args:
df: DataFrame containing the data
output_path: Path to save the HTML plot
@@ -50,21 +52,28 @@ def create_interactive_scatter_plot(df: pd.DataFrame, output_path: str, dataset_
"""
# Create the scatter plot
fig = go.Figure()
-
+
# Define colors for different configuration types
config_colors = {
- 'adaptive_sampling': '#1f77b4', # Blue
- 'oracle_top_k': '#ff7f0e', # Orange
- 'oracle_top_p': '#2ca02c', # Green
- 'hashattention': '#d62728', # Red
- 'adaptive_sampling_hat': '#9467bd', # Purple
- 'random_sampling': '#8c564b' # Brown
+ "adaptive_sampling": "#1f77b4", # Blue
+ "oracle_top_k": "#ff7f0e", # Orange
+ "oracle_top_p": "#2ca02c", # Green
+ "hashattention": "#d62728", # Red
+ "adaptive_sampling_hat": "#9467bd", # Purple
+ "random_sampling": "#8c564b", # Brown
}
-
+
# Add scatter traces for each configuration type
- for config_type in ['adaptive_sampling', 'oracle_top_k', 'oracle_top_p', 'hashattention', 'adaptive_sampling_hat', 'random_sampling']:
- config_data = df[df['config_type'] == config_type]
-
+ for config_type in [
+ "adaptive_sampling",
+ "oracle_top_k",
+ "oracle_top_p",
+ "hashattention",
+ "adaptive_sampling_hat",
+ "random_sampling",
+ ]:
+ config_data = df[df["config_type"] == config_type]
+
if len(config_data) > 0:
# Create hover text for this configuration type
hover_text = []
@@ -77,20 +86,20 @@ def create_interactive_scatter_plot(df: pd.DataFrame, output_path: str, dataset_
Sink Size: {row.get('sink_size', 'N/A')}Window Size: {row.get('window_size', 'N/A')}Heavy Size: {row.get('heavy_size', 'N/A')}Base Rate: {row.get('base_rate_sampling', 'N/A')}Epsilon: {row.get('epsilon', 'N/A')}Delta: {row.get('delta', 'N/A')}Top-K: {row.get('top_k', 'N/A')}Top-P: {row.get('top_p', 'N/A')}Hash Top-K: {row.get('hat_top_k', 'N/A')}Hash Heavy Size: {row.get('hat_heavy_size', 'N/A')}Hash MLP Hidden Size: {row.get('hat_mlp_hidden_size', 'N/A')}Hash MLP Activation: {row.get('hat_mlp_activation', 'N/A')}Heavy Size: {row.get('heavy_size', 'N/A')}Base Rate: {row.get('base_rate_sampling', 'N/A')}Hash MLP Hidden Size: {row.get('hat_mlp_hidden_size', 'N/A')}Hash MLP Activation: {row.get('hat_mlp_activation', 'N/A')}Sampling Rate: {row.get('sampling_rate', 'N/A')}Sampling Rate: {row.get('sampling_rate', 'N/A')}Density: {row.get('density', 'N/A'):.4f}Error: {row.get('error', 'N/A'):.4f}
"""
hover_text.append(config_info)
-
+
fig.add_trace(
go.Scatter(
- x=config_data['density'],
- y=config_data['error'],
- mode='markers',
- marker=dict(
- size=6,
- color=config_colors[config_type],
- opacity=0.7
- ),
+ x=config_data["density"],
+ y=config_data["error"],
+ mode="markers",
+ marker=dict(size=6, color=config_colors[config_type], opacity=0.7),
text=hover_text,
- hoverinfo='text',
- hovertemplate='%{text}", "AlexCuadron ", "Kumar Krishna Agrawal ", "Luis Gaspar Schroeder "]
readme = "README.md"
-license = {text = "MIT"}
-requires-python = ">=3.9"
+license = "MIT"
+homepage = "https://github.com/xAlg-ai/sparse-attention-hub"
+repository = "https://github.com/xAlg-ai/sparse-attention-hub"
+documentation = "https://sparse-attention-hub.readthedocs.io"
+keywords = ["attention", "sparse", "transformer", "deep-learning", "pytorch"]
classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Developers",
@@ -19,86 +20,75 @@ classifiers = [
"License :: OSI Approved :: MIT License",
"Operating System :: OS Independent",
"Programming Language :: Python :: 3",
- "Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
"Topic :: Software Development :: Libraries :: Python Modules",
]
-keywords = ["attention", "sparse", "transformer", "deep-learning", "pytorch"]
+packages = [{include = "sparse_attention_hub"}]
-dependencies = [
- "torch>=1.9.0",
- "numpy>=1.21.0",
- "matplotlib>=3.5.0",
- "seaborn>=0.11.0",
- "datasets>=2.0.0",
- "transformers>=4.20.0",
- "scikit-learn>=1.0.0",
- "rouge>=1.0.0",
- "nltk>=3.8.0",
- "bert-score>=0.3.0",
- "jieba>=0.42.0",
- "hf-transfer>=0.1.9",
- "fuzzywuzzy>=0.18.0",
- "pip>=25.1.1",
- "pytest>=8.4.1",
- "psutil>=7.0.0",
- "plotly>=6.2.0",
- "pandas>=2.3.1",
- "pynvml>=12.0.0",
- "colorama>=0.4.6",
-]
+[tool.poetry.dependencies]
+python = ">=3.9,<4.0"
+torch = ">=1.9.0"
+numpy = ">=1.21.0"
+matplotlib = ">=3.5.0"
+seaborn = ">=0.11.0"
+datasets = ">=2.0.0"
+transformers = ">=4.20.0"
+scikit-learn = ">=1.0.0"
+rouge = ">=1.0.0"
+nltk = ">=3.8.0"
+bert-score = ">=0.3.0"
+jieba = ">=0.42.0"
+hf-transfer = ">=0.1.9"
+fuzzywuzzy = ">=0.18.0"
+psutil = ">=7.0.0"
+plotly = ">=6.2.0"
+pandas = ">=2.3.1"
+pynvml = ">=12.0.0"
+colorama = ">=0.4.6"
-[project.optional-dependencies]
-dev = [
- # Testing
- "pytest>=7.4.0",
- "pytest-cov>=4.1.0",
- "pytest-xdist>=3.3.1",
- "pytest-mock>=3.11.1",
- "mock>=5.2.0",
- # Linting and formatting
- "black==22.10.0",
- "isort==5.12.0",
- "flake8==5.0.4",
- "pylint==2.17.4",
- "mypy==1.4.0",
- "bandit>=1.7.5",
- "pyupgrade>=3.10.1",
- "docformatter>=1.7.5",
- # Pre-commit hooks
- "pre-commit>=3.3.3",
- # Type stubs
- "types-PyYAML>=6.0.12.11",
- "types-requests>=2.31.0.2",
- "types-setuptools>=68.0.0.3",
-]
-docs = [
- "sphinx>=5.0.0",
- "sphinx-rtd-theme>=1.0.0",
- "myst-parser>=0.18.0",
-]
-benchmarks = [
- "wandb>=0.13.0",
- "tqdm>=4.64.0",
- "pandas>=1.5.0",
-]
+[tool.poetry.group.dev.dependencies]
+# Testing
+pytest = ">=7.4.0"
+pytest-cov = ">=4.1.0"
+pytest-xdist = ">=3.3.1"
+pytest-mock = ">=3.11.1"
+mock = ">=5.2.0"
+# Linting and formatting
+black = "22.10.0"
+isort = "5.12.0"
+flake8 = "5.0.4"
+pylint = "2.17.4"
+mypy = "1.4.0"
+bandit = ">=1.7.5"
+pyupgrade = ">=3.10.1"
+docformatter = ">=1.7.5"
+# Pre-commit hooks
+pre-commit = ">=3.3.3"
+# Type stubs
+types-PyYAML = ">=6.0.12.11"
+types-requests = ">=2.31.0.2"
+types-setuptools = ">=68.0.0.3"
-[project.urls]
-Homepage = "https://github.com/xAlg-ai/sparse-attention-hub"
-Repository = "https://github.com/xAlg-ai/sparse-attention-hub"
-Documentation = "https://sparse-attention-hub.readthedocs.io"
-"Bug Tracker" = "https://github.com/xAlg-ai/sparse-attention-hub/issues"
+[tool.poetry.group.docs]
+optional = true
-[project.scripts]
-sparse-attention-benchmark = "sparse_attention_hub.benchmark.executor:main"
+[tool.poetry.group.docs.dependencies]
+sphinx = ">=5.0.0"
+sphinx-rtd-theme = ">=1.0.0"
+myst-parser = ">=0.18.0"
+
+[tool.poetry.group.benchmarks]
+optional = true
-[tool.setuptools.packages.find]
-where = ["."]
-include = ["sparse_attention_hub*"]
-exclude = ["tests*"]
+[tool.poetry.group.benchmarks.dependencies]
+wandb = ">=0.13.0"
+tqdm = ">=4.64.0"
+
+[tool.poetry.scripts]
+sparse-attention-benchmark = "sparse_attention_hub.benchmark.executor:main"
[tool.pytest.ini_options]
testpaths = ["tests"]
@@ -120,7 +110,7 @@ markers = [
[tool.black]
line-length = 88
-target-version = ["py38", "py39", "py310", "py311"]
+target-version = ["py39", "py310", "py311"]
include = '\.pyi?$'
extend-exclude = '''
/(
@@ -143,7 +133,7 @@ line_length = 88
known_first_party = ["sparse_attention_hub"]
[tool.mypy]
-python_version = "3.8"
+python_version = "3.9"
warn_return_any = true
warn_unused_configs = true
disallow_untyped_defs = true
@@ -156,6 +146,7 @@ warn_unused_ignores = true
warn_no_return = true
warn_unreachable = true
strict_equality = true
+
[[tool.mypy.overrides]]
module = [
"torch.*",
@@ -166,4 +157,4 @@ module = [
"wandb.*",
]
ignore_missing_imports = true
-ignore_errors = true
+ignore_errors = true
\ No newline at end of file
diff --git a/requirements-dev.txt b/requirements-dev.txt
deleted file mode 100644
index 424422ea..00000000
--- a/requirements-dev.txt
+++ /dev/null
@@ -1,38 +0,0 @@
-# Development dependencies for sparse-attention-hub
-# Linting and formatting
-black==22.10.0
-isort==5.12.0
-flake8==5.0.4
-pylint==2.17.4
-mypy==1.4.0
-bandit==1.7.5
-pyupgrade==3.10.1
-docformatter==1.7.5
-
-# Testing
-pytest==7.4.0
-pytest-cov==4.1.0
-pytest-xdist==3.3.1
-pytest-mock==3.11.1
-mock==5.2.0
-
-# Pre-commit hooks
-pre-commit==3.3.3
-
-# Type stubs
-types-PyYAML==6.0.12.11
-types-requests==2.31.0.2
-types-setuptools==68.0.0.3
-
-# Documentation
-sphinx==7.1.2
-sphinx-rtd-theme==1.3.0
-myst-parser==2.0.0
-
-# Jupyter for examples
-jupyter==1.0.0
-ipykernel==6.25.0
-
-# Additional development tools
-twine==4.0.2
-build==0.10.0
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 23116853..00000000
--- a/requirements.txt
+++ /dev/null
@@ -1,20 +0,0 @@
-# Core dependencies
-torch>=1.9.0
-numpy>=1.21.0
-matplotlib>=3.5.0
-seaborn>=0.11.0
-datasets>=2.0.0
-transformers>=4.20.0
-
-# Development dependencies
-pytest>=7.0.0
-pytest-cov>=4.0.0
-black>=22.0.0
-isort>=5.10.0
-flake8>=5.0.0
-mypy>=0.991
-
-# Optional dependencies for benchmarks
-wandb>=0.13.0
-tqdm>=4.64.0
-pandas>=1.5.0
\ No newline at end of file
diff --git a/scripts/format.sh b/scripts/format.sh
index 7d3fdfd2..0fbd9c50 100755
--- a/scripts/format.sh
+++ b/scripts/format.sh
@@ -23,7 +23,7 @@ builtin cd "$ROOT" || exit 1
# Check if development dependencies are installed
check_tool_installed() {
if ! command -v "$1" &> /dev/null; then
- echo "Error: $1 is not installed. Please run: pip install -r requirements-dev.txt"
+ echo "Error: $1 is not installed. Please run: poetry install --with dev"
exit 1
fi
}
@@ -43,7 +43,7 @@ tool_version_check() {
if [[ "$installed_version" != "$required_version" ]]; then
echo "Warning: $tool_name version mismatch. Required: $required_version, Installed: $installed_version"
- echo "Consider running: pip install -r requirements-dev.txt"
+ echo "Consider running: poetry install --with dev"
fi
}
@@ -54,13 +54,13 @@ FLAKE8_VERSION=$(flake8 --version | head -n 1 | awk '{print $1}')
MYPY_VERSION=$(mypy --version | awk '{print $2}')
PYLINT_VERSION=$(pylint --version | head -n 1 | awk '{print $2}')
-# Check versions against requirements-dev.txt
-if [[ -f "requirements-dev.txt" ]]; then
- tool_version_check "black" "$BLACK_VERSION" "$(grep "black==" requirements-dev.txt | cut -d'=' -f3)"
- tool_version_check "isort" "$ISORT_VERSION" "$(grep "isort==" requirements-dev.txt | cut -d'=' -f3)"
- tool_version_check "flake8" "$FLAKE8_VERSION" "$(grep "flake8==" requirements-dev.txt | cut -d'=' -f3)"
- tool_version_check "mypy" "$MYPY_VERSION" "$(grep "mypy==" requirements-dev.txt | cut -d'=' -f3)"
- tool_version_check "pylint" "$PYLINT_VERSION" "$(grep "pylint==" requirements-dev.txt | cut -d'=' -f3)"
+# Check versions against pyproject.toml
+if [[ -f "pyproject.toml" ]]; then
+ tool_version_check "black" "$BLACK_VERSION" "$(grep 'black = ' pyproject.toml | sed 's/.*"\(.*\)".*/\1/')"
+ tool_version_check "isort" "$ISORT_VERSION" "$(grep 'isort = ' pyproject.toml | sed 's/.*"\(.*\)".*/\1/')"
+ tool_version_check "flake8" "$FLAKE8_VERSION" "$(grep 'flake8 = ' pyproject.toml | sed 's/.*"\(.*\)".*/\1/')"
+ tool_version_check "mypy" "$MYPY_VERSION" "$(grep 'mypy = ' pyproject.toml | sed 's/.*"\(.*\)".*/\1/')"
+ tool_version_check "pylint" "$PYLINT_VERSION" "$(grep 'pylint = ' pyproject.toml | sed 's/.*"\(.*\)".*/\1/')"
fi
# Formatting flags
diff --git a/scripts/lint.sh b/scripts/lint.sh
index dba86cdf..0aa9e2af 100755
--- a/scripts/lint.sh
+++ b/scripts/lint.sh
@@ -29,7 +29,7 @@ NC='\033[0m' # No Color
# Check if tool is installed
check_tool() {
if ! command -v "$1" &> /dev/null; then
- echo -e "${RED}Error: $1 is not installed. Please run: pip install -r requirements-dev.txt${NC}"
+ echo -e "${RED}Error: $1 is not installed. Please run: poetry install --with dev${NC}"
return 1
fi
return 0
diff --git a/sparse_attention_hub/metric_logging/logger.py b/sparse_attention_hub/metric_logging/logger.py
index 8663135c..a6788008 100644
--- a/sparse_attention_hub/metric_logging/logger.py
+++ b/sparse_attention_hub/metric_logging/logger.py
@@ -13,23 +13,41 @@
@dataclass
class LogEvent:
- """Log event data structure."""
+ """Log event data structure for metric logging.
+
+ Attributes:
+ timestamp: When the event was logged.
+ metric: Metric identifier string.
+ value: The actual metric value (can be None).
+ metadata: Additional context information like layer, head, etc.
+ location: Auto-inferred location as "module.function" or "class.method".
+ """
timestamp: datetime
- metric: str # Metric identifier string
+ metric: str
value: Union[None, Any]
- metadata: Dict[str, Any] # Additional context (layer, head, etc.)
- location: str # Auto-inferred: "module.function" or "class.method"
+ metadata: Dict[str, Any]
+ location: str
class MicroMetricLogger:
- """Singleton logger for micro metrics with queue-based architecture."""
+ """Singleton logger for micro metrics with queue-based architecture.
+
+ This class provides a singleton pattern for logging micro metrics during
+ sparse attention operations. It supports metric registration, sampling,
+ and configurable limits on record count and flush behavior.
+
+ Attributes:
+ log_path: Optional directory path where log files will be written.
+ flush_every: Number of events after which to flush to disk.
+ flush_interval: Time interval in seconds for automatic flushing.
+ max_records: Maximum number of records to log (None for unlimited).
+ sampling_factor: Probability of logging each event (0.0-1.0).
+ """
_instance: Optional["MicroMetricLogger"] = None
_initialized: bool = False
-
- # Class-level storage for registered metrics (works without initialization)
- _registered_metrics: Dict[str, type] = {} # identifier -> dtype mapping
+ _registered_metrics: Dict[str, type] = {}
def __new__(cls, *args, **kwargs) -> "MicroMetricLogger":
if cls._instance is None:
@@ -39,46 +57,123 @@ def __new__(cls, *args, **kwargs) -> "MicroMetricLogger":
def __init__(
self,
log_path: Optional[str] = None,
- flush_every: int = 1000, # Flush every N events
- flush_interval: float = 60.0, # Flush every N seconds
- enabled_metrics: Union[List[str], str] = None,
- max_records: Optional[int] = None, # Maximum number of events to log
- sampling_factor: float = 1.0, # Probability of logging each event (0.0-1.0)
- ): # List of string identifiers to enable, or "all"
+ flush_every: int = 1000,
+ flush_interval: float = 60.0,
+ enabled_metrics: Union[List[str], str, None] = None,
+ max_records: Optional[int] = None,
+ sampling_factor: float = 1.0,
+ ) -> None:
+ """Initialize the MicroMetricLogger.
+
+ Args:
+ log_path: Optional directory path where log files will be written.
+ flush_every: Number of events after which to flush to disk.
+ flush_interval: Time interval in seconds for automatic flushing.
+ enabled_metrics: List of metric identifiers to enable, "all", or None.
+ max_records: Maximum number of events to log (None for unlimited).
+ sampling_factor: Probability of logging each event (0.0-1.0).
+ """
if not self._initialized:
- self.log_path = log_path
- self.flush_every = flush_every
- self.flush_interval = flush_interval
- self.max_records = max_records
- self.sampling_factor = max(
- 0.0, min(1.0, sampling_factor)
- ) # Clamp to [0.0, 1.0]
-
- # Internal state
- self.log_queue: deque = deque(maxlen=10000) # Circular buffer
- self.enabled_metrics: set = set()
- self.last_flush_time = time.time()
- self._total_records_logged: int = 0 # Track total events logged
-
- # Enable metrics if log_path is provided
- if self.log_path is not None:
- self._ensure_log_directory()
- self.enable_metrics(enabled_metrics)
-
- MicroMetricLogger._initialized = True
+ self._initialize_logger(
+ log_path,
+ flush_every,
+ flush_interval,
+ enabled_metrics,
+ max_records,
+ sampling_factor,
+ )
else:
- if self.log_path and log_path and self.log_path != log_path:
- print(
- f"Warning: MicroMetricLogger already initialized with log_path: {self.log_path}"
- )
+ self._handle_reinitialization(log_path)
+
+ def _initialize_logger(
+ self,
+ log_path: Optional[str],
+ flush_every: int,
+ flush_interval: float,
+ enabled_metrics: Union[List[str], str, None],
+ max_records: Optional[int],
+ sampling_factor: float,
+ ) -> None:
+ """Initialize logger instance for the first time.
+
+ Args:
+ log_path: Directory path for log files.
+ flush_every: Number of events before flushing.
+ flush_interval: Time interval for flushing.
+ enabled_metrics: Metrics to enable initially.
+ max_records: Maximum records limit.
+ sampling_factor: Event sampling probability.
+ """
+ self._set_configuration(
+ log_path, flush_every, flush_interval, max_records, sampling_factor
+ )
+ self._initialize_state()
+ self._configure_logging_if_needed(log_path, enabled_metrics)
+ MicroMetricLogger._initialized = True
+
+ def _set_configuration(
+ self,
+ log_path: Optional[str],
+ flush_every: int,
+ flush_interval: float,
+ max_records: Optional[int],
+ sampling_factor: float,
+ ) -> None:
+ """Set logger configuration parameters.
+
+ Args:
+ log_path: Directory path for log files.
+ flush_every: Number of events before flushing.
+ flush_interval: Time interval for flushing.
+ max_records: Maximum records limit.
+ sampling_factor: Event sampling probability.
+ """
+ self.log_path = log_path
+ self.flush_every = flush_every
+ self.flush_interval = flush_interval
+ self.max_records = max_records
+ self.sampling_factor = max(0.0, min(1.0, sampling_factor))
- # main registration function
+ def _initialize_state(self) -> None:
+ """Initialize internal state variables."""
+ self.log_queue: deque = deque(maxlen=10000)
+ self.enabled_metrics: set = set()
+ self.last_flush_time = time.time()
+ self._total_records_logged: int = 0
+
+ def _configure_logging_if_needed(
+ self, log_path: Optional[str], enabled_metrics: Union[List[str], str, None]
+ ) -> None:
+ """Configure logging if log path is provided.
+
+ Args:
+ log_path: Directory path for log files.
+ enabled_metrics: Metrics to enable initially.
+ """
+ if log_path is not None:
+ self._ensure_log_directory()
+ self.enable_metrics(enabled_metrics)
+
+ def _handle_reinitialization(self, log_path: Optional[str]) -> None:
+ """Handle warning when logger is reinitialized.
+
+ Args:
+ log_path: New log path being requested.
+ """
+ if self.log_path and log_path and self.log_path != log_path:
+ print(
+ f"Warning: MicroMetricLogger already initialized with log_path: {self.log_path}"
+ )
@classmethod
def register_metric(cls, identifier: str, dtype: type) -> None:
"""Register a metric with its string identifier and expected data type.
- This works at class level and doesn't require initialization.
+ This works at class level and doesn't require logger initialization.
+
+ Args:
+ identifier: Unique string identifier for the metric.
+ dtype: Expected data type for the metric values.
"""
if identifier in cls._registered_metrics:
print(f"Warning: Metric '{identifier}' is being re-registered")
@@ -86,10 +181,12 @@ def register_metric(cls, identifier: str, dtype: type) -> None:
@classmethod
def get_registered_metrics(cls) -> Dict[str, type]:
- """Get all registered metrics at class level."""
- return cls._registered_metrics.copy()
+ """Get all registered metrics at class level.
- # helper methods
+ Returns:
+ Dictionary mapping metric identifiers to their expected data types.
+ """
+ return cls._registered_metrics.copy()
def _ensure_log_directory(self) -> None:
"""Ensure the log directory exists."""
@@ -97,24 +194,20 @@ def _ensure_log_directory(self) -> None:
os.makedirs(self.log_path, exist_ok=True)
def _get_calling_location(self) -> str:
- """Get the calling location using inspect module."""
+ """Get the calling location using inspect module.
+
+ Returns:
+ String representation of the calling location in format
+ "module.class.method" or "module.function".
+ """
try:
- # Get the calling frame (skip this method and the log method)
caller_frame = inspect.currentframe().f_back.f_back
if caller_frame is None:
return "unknown"
- # Get module name
- module = inspect.getmodule(caller_frame)
- module_name = module.__name__ if module else "unknown"
-
- # Get function/class name
+ module_name = self._get_module_name(caller_frame)
function_name = caller_frame.f_code.co_name
-
- # Try to get class name if it's a method
- class_name = None
- if "self" in caller_frame.f_locals:
- class_name = caller_frame.f_locals["self"].__class__.__name__
+ class_name = self._get_class_name(caller_frame)
if class_name:
return f"{module_name}.{class_name}.{function_name}"
@@ -123,66 +216,146 @@ def _get_calling_location(self) -> str:
except Exception:
return "unknown"
- def __del__(self):
- """Cleanup when logger is destroyed."""
- self.flush() # Final flush
+ def _get_module_name(self, frame: Any) -> str:
+ """Get module name from frame.
- # api
+ Args:
+ frame: Stack frame object.
- def enable_metrics(self, metrics: Union[List[str], str] = None) -> None:
+ Returns:
+ Module name or "unknown" if not available.
+ """
+ module = inspect.getmodule(frame)
+ return module.__name__ if module else "unknown"
+
+ def _get_class_name(self, frame: Any) -> Optional[str]:
+ """Get class name from frame if it's a method call.
+
+ Args:
+ frame: Stack frame object.
+
+ Returns:
+ Class name if available, None otherwise.
+ """
+ if "self" in frame.f_locals:
+ return frame.f_locals["self"].__class__.__name__
+ return None
+
+ def __del__(self) -> None:
+ """Cleanup when logger is destroyed."""
+ self.flush()
+
+ def enable_metrics(self, metrics: Union[List[str], str, None] = None) -> None:
"""Enable logging for specific metrics.
Args:
- metrics: List of metric identifiers to enable, or "all" for all registered metrics.
- If None, enables no metrics (empty list).
+ metrics: List of metric identifiers to enable, "all" for all registered
+ metrics, or None to disable all metrics.
"""
if metrics == "all":
self.enabled_metrics = set(self._registered_metrics.keys())
elif isinstance(metrics, (list, set)):
- # Only enable metrics that are registered
- valid_metrics = set(metrics) & set(self._registered_metrics.keys())
- invalid_metrics = set(metrics) - set(self._registered_metrics.keys())
- if invalid_metrics:
- print(
- f"Warning: Attempting to enable unregistered metrics: {invalid_metrics}"
- )
- self.enabled_metrics = valid_metrics
+ self._enable_selected_metrics(metrics)
else:
- # Default to empty set
self.enabled_metrics = set()
- def log(self, identifier: str, value: Any, metadata: Dict[str, Any] = None) -> None:
- """Log a metric value with optional metadata. Location is auto-inferred.
+ def _enable_selected_metrics(self, metrics: Union[List[str], set]) -> None:
+ """Enable only the selected metrics that are registered.
- This only works if log_path is defined.
+ Args:
+ metrics: List or set of metric identifiers to enable.
"""
- # Check if logging is configured
- if self.log_path is None:
+ valid_metrics = set(metrics) & set(self._registered_metrics.keys())
+ invalid_metrics = set(metrics) - set(self._registered_metrics.keys())
+
+ if invalid_metrics:
print(
- f"Warning: Cannot log metric '{identifier}' - log_path not defined. Use configure_logging() first."
+ f"Warning: Attempting to enable unregistered metrics: {invalid_metrics}"
)
+
+ self.enabled_metrics = valid_metrics
+
+ def log(
+ self, identifier: str, value: Any, metadata: Optional[Dict[str, Any]] = None
+ ) -> None:
+ """Log a metric value with optional metadata.
+
+ Location is auto-inferred from the calling context. This only works if log_path is defined.
+
+ Args:
+ identifier: Unique string identifier for the metric.
+ value: The metric value to log.
+ metadata: Optional additional context information.
+ """
+ if not self._should_log_metric(identifier):
+ return
+
+ if not self._is_within_limits():
return
- # Check if metric is enabled
+ if not self._passes_sampling():
+ return
+
+ event = self._create_log_event(identifier, value, metadata)
+ self._add_event_and_flush_if_needed(event)
+
+ def _should_log_metric(self, identifier: str) -> bool:
+ """Check if the metric should be logged based on configuration.
+
+ Args:
+ identifier: Metric identifier to check.
+
+ Returns:
+ True if the metric should be logged, False otherwise.
+ """
+ if self.log_path is None:
+ print(
+ f"Warning: Cannot log metric '{identifier}' - log_path not defined. "
+ "Use configure_logging() first."
+ )
+ return False
+
if identifier not in self.enabled_metrics:
print(
f"Warning: Attempting to log metric '{identifier}' which is not enabled"
)
- return
+ return False
- # Check if max_records limit has been reached
- if (
+ return True
+
+ def _is_within_limits(self) -> bool:
+ """Check if logging is within configured limits.
+
+ Returns:
+ True if within limits, False if max records reached.
+ """
+ return not (
self.max_records is not None
and self._total_records_logged >= self.max_records
- ):
- return # Silently ignore - we've hit the limit
+ )
+
+ def _passes_sampling(self) -> bool:
+ """Check if event passes sampling filter.
+
+ Returns:
+ True if event should be logged based on sampling factor.
+ """
+ return self.sampling_factor >= 1.0 or random.random() <= self.sampling_factor
- # Apply sampling factor
- if self.sampling_factor < 1.0 and random.random() > self.sampling_factor:
- return # Event not selected for logging due to sampling
+ def _create_log_event(
+ self, identifier: str, value: Any, metadata: Optional[Dict[str, Any]]
+ ) -> LogEvent:
+ """Create a log event from the provided data.
- # Create log event
- event = LogEvent(
+ Args:
+ identifier: Metric identifier.
+ value: Metric value.
+ metadata: Optional metadata.
+
+ Returns:
+ LogEvent instance ready for logging.
+ """
+ return LogEvent(
timestamp=datetime.now(),
metric=identifier,
value=value,
@@ -190,75 +363,131 @@ def log(self, identifier: str, value: Any, metadata: Dict[str, Any] = None) -> N
location=self._get_calling_location(),
)
- # Add to queue and increment counter
+ def _add_event_and_flush_if_needed(self, event: LogEvent) -> None:
+ """Add event to queue and flush if threshold is reached.
+
+ Args:
+ event: LogEvent to add to the queue.
+ """
self.log_queue.append(event)
self._total_records_logged += 1
- # Check if we should flush
if len(self.log_queue) >= self.flush_every:
self.flush()
def configure_logging(
self,
log_path: str,
- enabled_metrics: Union[List[str], str] = None,
+ enabled_metrics: Union[List[str], str, None] = None,
max_records: Optional[int] = None,
sampling_factor: float = 1.0,
) -> None:
"""Configure logging with a log path and optionally enable metrics.
This must be called before logging can work. Resets the total records counter.
+
+ Args:
+ log_path: Directory path where log files will be written.
+ enabled_metrics: List of metric identifiers to enable, "all", or None.
+ max_records: Maximum number of records to log (None for unlimited).
+ sampling_factor: Probability of logging each event (0.0-1.0).
"""
self.log_path = log_path
self._ensure_log_directory()
self.enable_metrics(enabled_metrics)
+ self._update_logging_limits(max_records, sampling_factor)
+ self._total_records_logged = 0
- # Update limits if provided
+ def _update_logging_limits(
+ self, max_records: Optional[int], sampling_factor: float
+ ) -> None:
+ """Update logging limits and sampling factor.
+
+ Args:
+ max_records: Maximum records limit.
+ sampling_factor: Event sampling probability.
+ """
self.max_records = max_records
self.sampling_factor = max(0.0, min(1.0, sampling_factor))
- # Reset the total records counter when reconfiguring
- self._total_records_logged = 0
-
def flush(self) -> None:
- """Force flush the current queue to disk."""
+ """Force flush the current queue to disk.
+
+ Writes all queued events to the log file in JSONL format.
+ """
if not self.log_queue or self.log_path is None:
return
- # Get current timestamp for filename
- filename = "micro_metrics.jsonl"
- filepath = os.path.join(self.log_path, filename)
+ filepath = os.path.join(self.log_path, "micro_metrics.jsonl")
+ self._write_events_to_file(filepath)
+ self.last_flush_time = time.time()
+
+ def _write_events_to_file(self, filepath: str) -> None:
+ """Write all queued events to the specified file.
- # Write events to file
+ Args:
+ filepath: Path to the log file.
+ """
with open(filepath, "a", encoding="utf-8") as f:
while self.log_queue:
event = self.log_queue.popleft()
- # Convert dataclass to dict and serialize
- event_dict = asdict(event)
- # Convert datetime to ISO format string
- event_dict["timestamp"] = event_dict["timestamp"].isoformat()
+ event_dict = self._serialize_event(event)
f.write(json.dumps(event_dict) + "\n")
- self.last_flush_time = time.time()
+ def _serialize_event(self, event: LogEvent) -> Dict[str, Any]:
+ """Serialize a log event to a dictionary.
+
+ Args:
+ event: LogEvent to serialize.
+
+ Returns:
+ Dictionary representation of the event.
+ """
+ event_dict = asdict(event)
+ event_dict["timestamp"] = event_dict["timestamp"].isoformat()
+ return event_dict
def is_metric_enabled(self, identifier: str) -> bool:
- """Check if a specific metric is requested for logging."""
+ """Check if a specific metric is enabled for logging.
+
+ Args:
+ identifier: Metric identifier to check.
+
+ Returns:
+ True if the metric is enabled, False otherwise.
+ """
return identifier in self.enabled_metrics
def get_enabled_metrics(self) -> set:
- """Get currently enabled metrics."""
+ """Get currently enabled metrics.
+
+ Returns:
+ Set of enabled metric identifiers.
+ """
return self.enabled_metrics.copy()
def is_logging_configured(self) -> bool:
- """Check if logging is configured (log_path is set)."""
+ """Check if logging is configured.
+
+ Returns:
+ True if log_path is set, False otherwise.
+ """
return self.log_path is not None
def get_total_records_logged(self) -> int:
- """Get the total number of records logged since initialization or last configure_logging call."""
+ """Get the total number of records logged.
+
+ Returns:
+ Number of records logged since initialization or last configure_logging call.
+ """
return getattr(self, "_total_records_logged", 0)
def is_max_records_reached(self) -> bool:
- """Check if the maximum number of records has been reached."""
+ """Check if the maximum number of records has been reached.
+
+ Returns:
+ True if max records limit has been reached, False otherwise.
+ """
if self.max_records is None:
return False
return self.get_total_records_logged() >= self.max_records
@@ -274,9 +503,17 @@ def get_records_remaining(self) -> Optional[int]:
return max(0, self.max_records - self.get_total_records_logged())
def get_sampling_factor(self) -> float:
- """Get the current sampling factor."""
+ """Get the current sampling factor.
+
+ Returns:
+ Current sampling factor (0.0-1.0).
+ """
return getattr(self, "sampling_factor", 1.0)
def get_max_records(self) -> Optional[int]:
- """Get the current max_records limit."""
+ """Get the current max_records limit.
+
+ Returns:
+ Maximum records limit, or None if no limit is set.
+ """
return getattr(self, "max_records", None)
diff --git a/sparse_attention_hub/plotting/generator.py b/sparse_attention_hub/plotting/generator.py
index 61b317ca..1847b134 100644
--- a/sparse_attention_hub/plotting/generator.py
+++ b/sparse_attention_hub/plotting/generator.py
@@ -4,19 +4,35 @@
from typing import Any, Dict, Optional
import matplotlib.pyplot as plt
+import numpy as np
import seaborn as sns
from .granularity import Granularity
class PlotGenerator:
- """Generates plots for sparse attention analysis."""
+ """Generates plots for sparse attention analysis.
- def __init__(self, storage_path: str = "./plots"):
+ This class provides functionality to generate various types of plots
+ for analyzing sparse attention patterns at different granularities.
+
+ Attributes:
+ storage_path: Directory path where generated plots will be saved.
+ """
+
+ def __init__(self, storage_path: str = "./plots") -> None:
+ """Initialize the PlotGenerator.
+
+ Args:
+ storage_path: Directory path for storing generated plots.
+ Defaults to "./plots".
+ """
self.storage_path = storage_path
self._ensure_storage_directory()
+ self._setup_plotting_style()
- # Set up plotting style
+ def _setup_plotting_style(self) -> None:
+ """Configure matplotlib and seaborn plotting styles."""
plt.style.use("default")
sns.set_palette("husl")
@@ -34,105 +50,160 @@ def generate_plot(
"""Generate a plot with specified granularity.
Args:
- granularity: Level of granularity for the plot
- data: Data to plot (if None, generates sample data)
- plot_type: Type of plot to generate
- **kwargs: Additional plotting parameters
+ granularity: Level of granularity for the plot (per token, head, or layer).
+ data: Optional data to plot. If None, generates sample data for demonstration.
+ plot_type: Type of plot to generate. Currently unused but reserved for future extensions.
+ **kwargs: Additional plotting parameters passed to the plot generation methods.
Returns:
- Path to the generated plot file
+ Absolute path to the generated plot file.
+
+ Raises:
+ ValueError: If the specified granularity is not supported.
"""
if granularity == Granularity.PER_TOKEN:
- return self._generate_plot_1(granularity, data, **kwargs)
+ return self._generate_line_plot(granularity, data, **kwargs)
elif granularity == Granularity.PER_HEAD:
- return self._generate_plot_2(granularity, data, **kwargs)
+ return self._generate_heatmap_plot(granularity, data, **kwargs)
elif granularity == Granularity.PER_LAYER:
- return self._generate_plot_1(granularity, data, **kwargs)
+ return self._generate_line_plot(granularity, data, **kwargs)
else:
raise ValueError(f"Unsupported granularity: {granularity}")
- def _generate_plot_1(
+ def _generate_line_plot(
self,
granularity: Granularity,
data: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> str:
- """Generate plot type 1 (line plots, attention patterns).
+ """Generate line plots for attention patterns.
Args:
- granularity: Level of granularity
- data: Data to plot
- **kwargs: Additional parameters
+ granularity: Level of granularity (per token or per layer).
+ data: Optional data to plot. If None, generates sample data.
+ **kwargs: Additional plotting parameters.
Returns:
- Path to generated plot
+ Path to the generated plot file.
"""
- # TODO: Implement plot generation
fig, ax = plt.subplots(figsize=(10, 6))
if data is None:
- # Generate sample data for demonstration
- import numpy as np
-
- x = np.linspace(0, 10, 100)
- y = np.sin(x) + np.random.normal(0, 0.1, 100)
- ax.plot(x, y, label=f"Sample {granularity.value} data")
+ self._plot_sample_line_data(ax, granularity)
else:
- # Plot actual data
- # Implementation depends on data structure
- pass
+ self._plot_actual_line_data(ax, data)
+
+ self._configure_line_plot_axes(ax, granularity)
+ return self._save_plot(fig, "plot", granularity, data)
+
+ def _plot_sample_line_data(self, ax: plt.Axes, granularity: Granularity) -> None:
+ """Plot sample line data for demonstration.
+
+ Args:
+ ax: Matplotlib axes object to plot on.
+ granularity: Granularity level for labeling.
+ """
+ x = np.linspace(0, 10, 100)
+ y = np.sin(x) + np.random.normal(0, 0.1, 100)
+ ax.plot(x, y, label=f"Sample {granularity.value} data")
+
+ def _plot_actual_line_data(self, ax: plt.Axes, data: Dict[str, Any]) -> None:
+ """Plot actual line data from provided data dictionary.
+
+ Args:
+ ax: Matplotlib axes object to plot on.
+ data: Data dictionary containing plot information.
+ """
+ # Implementation depends on data structure
+ # This is a placeholder for future actual data plotting
+ pass
+
+ def _configure_line_plot_axes(self, ax: plt.Axes, granularity: Granularity) -> None:
+ """Configure axes for line plots.
- ax.set_title(
+ Args:
+ ax: Matplotlib axes object to configure.
+ granularity: Granularity level for title formatting.
+ """
+ title = (
f"Sparse Attention Analysis - {granularity.value.replace('_', ' ').title()}"
)
+ ax.set_title(title)
ax.set_xlabel("Position")
ax.set_ylabel("Attention Weight")
ax.legend()
ax.grid(True, alpha=0.3)
- # Save plot
- filename = f"plot_{granularity.value}_{hash(str(data))}.png"
- filepath = os.path.join(self.storage_path, filename)
- plt.savefig(filepath, dpi=300, bbox_inches="tight")
- plt.close()
-
- return filepath
-
- def _generate_plot_2(
+ def _generate_heatmap_plot(
self,
granularity: Granularity,
data: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> str:
- """Generate plot type 2 (heatmaps, attention matrices).
+ """Generate heatmaps for attention matrices.
Args:
- granularity: Level of granularity
- data: Data to plot
- **kwargs: Additional parameters
+ granularity: Level of granularity (typically per head).
+ data: Optional data containing attention matrix. If None, generates sample data.
+ **kwargs: Additional plotting parameters.
Returns:
- Path to generated plot
+ Path to the generated plot file.
+
+ Raises:
+ ValueError: If data is provided but doesn't contain required attention_matrix.
"""
- # TODO: Implement heatmap generation
fig, ax = plt.subplots(figsize=(8, 8))
+ attention_matrix = self._get_attention_matrix(granularity, data)
+ self._create_heatmap(ax, attention_matrix)
+ self._configure_heatmap_axes(ax, granularity)
+
+ return self._save_plot(fig, "heatmap", granularity, data)
+
+ def _get_attention_matrix(
+ self, granularity: Granularity, data: Optional[Dict[str, Any]]
+ ) -> np.ndarray:
+ """Get attention matrix from data or generate sample data.
+
+ Args:
+ granularity: Granularity level for determining matrix size.
+ data: Optional data dictionary containing attention matrix.
+
+ Returns:
+ Attention matrix as numpy array.
+
+ Raises:
+ ValueError: If data is provided but doesn't contain attention_matrix.
+ """
if data is None:
- # Generate sample attention matrix
- import numpy as np
-
- size = 12 if granularity == Granularity.PER_HEAD else 8
- attention_matrix = np.random.rand(size, size)
- attention_matrix = (
- attention_matrix + attention_matrix.T
- ) / 2 # Make symmetric
- else:
- # Use actual data
- attention_matrix = data.get("attention_matrix", None)
- if attention_matrix is None:
- raise ValueError("attention_matrix required in data for heatmap")
+ return self._generate_sample_attention_matrix(granularity)
+
+ attention_matrix = data.get("attention_matrix")
+ if attention_matrix is None:
+ raise ValueError("attention_matrix required in data for heatmap")
+ return attention_matrix
- # Create heatmap
+ def _generate_sample_attention_matrix(self, granularity: Granularity) -> np.ndarray:
+ """Generate sample attention matrix for demonstration.
+
+ Args:
+ granularity: Granularity level for determining matrix size.
+
+ Returns:
+ Symmetric attention matrix.
+ """
+ size = 12 if granularity == Granularity.PER_HEAD else 8
+ attention_matrix = np.random.rand(size, size)
+ return (attention_matrix + attention_matrix.T) / 2 # Make symmetric
+
+ def _create_heatmap(self, ax: plt.Axes, attention_matrix: np.ndarray) -> None:
+ """Create heatmap visualization.
+
+ Args:
+ ax: Matplotlib axes object to plot on.
+ attention_matrix: Attention matrix to visualize.
+ """
sns.heatmap(
attention_matrix,
annot=True,
@@ -142,18 +213,40 @@ def _generate_plot_2(
cbar_kws={"label": "Attention Weight"},
)
- ax.set_title(
- f"Attention Matrix - {granularity.value.replace('_', ' ').title()}"
- )
+ def _configure_heatmap_axes(self, ax: plt.Axes, granularity: Granularity) -> None:
+ """Configure axes for heatmap plots.
+
+ Args:
+ ax: Matplotlib axes object to configure.
+ granularity: Granularity level for title formatting.
+ """
+ title = f"Attention Matrix - {granularity.value.replace('_', ' ').title()}"
+ ax.set_title(title)
ax.set_xlabel("Key Position")
ax.set_ylabel("Query Position")
- # Save plot
- filename = f"heatmap_{granularity.value}_{hash(str(data))}.png"
- filepath = os.path.join(self.storage_path, filename)
- plt.savefig(filepath, dpi=300, bbox_inches="tight")
- plt.close()
+ def _save_plot(
+ self,
+ fig: plt.Figure,
+ plot_type: str,
+ granularity: Granularity,
+ data: Optional[Dict[str, Any]],
+ ) -> str:
+ """Save plot to file and close figure.
+ Args:
+ fig: Matplotlib figure object to save.
+ plot_type: Type of plot for filename.
+ granularity: Granularity level for filename.
+ data: Data used for generating unique filename hash.
+
+ Returns:
+ Absolute path to the saved plot file.
+ """
+ filename = f"{plot_type}_{granularity.value}_{hash(str(data))}.png"
+ filepath = os.path.join(self.storage_path, filename)
+ fig.savefig(filepath, dpi=300, bbox_inches="tight")
+ plt.close(fig)
return filepath
def generate_comparison_plot(
@@ -162,36 +255,50 @@ def generate_comparison_plot(
"""Generate comparison plot for multiple datasets.
Args:
- data_dict: Dictionary mapping labels to data
- granularity: Level of granularity
+ data_dict: Dictionary mapping dataset labels to their corresponding data.
+ granularity: Level of granularity for the comparison plot.
Returns:
- Path to generated comparison plot
+ Path to the generated comparison plot file.
"""
fig, ax = plt.subplots(figsize=(12, 8))
+ self._plot_comparison_data(ax, data_dict)
+ self._configure_comparison_axes(ax, granularity)
+
+ return self._save_plot(fig, "comparison", granularity, data_dict)
+
+ def _plot_comparison_data(self, ax: plt.Axes, data_dict: Dict[str, Any]) -> None:
+ """Plot comparison data for multiple datasets.
+
+ Args:
+ ax: Matplotlib axes object to plot on.
+ data_dict: Dictionary mapping labels to data.
+ """
for label, data in data_dict.items():
- # Plot each dataset
- # Implementation depends on data structure
+ # Plot each dataset - implementation depends on data structure
+ # This is a placeholder for future actual data plotting
pass
- ax.set_title(f"Comparison - {granularity.value.replace('_', ' ').title()}")
+ def _configure_comparison_axes(
+ self, ax: plt.Axes, granularity: Granularity
+ ) -> None:
+ """Configure axes for comparison plots.
+
+ Args:
+ ax: Matplotlib axes object to configure.
+ granularity: Granularity level for title formatting.
+ """
+ title = f"Comparison - {granularity.value.replace('_', ' ').title()}"
+ ax.set_title(title)
ax.legend()
ax.grid(True, alpha=0.3)
- # Save plot
- filename = f"comparison_{granularity.value}_{hash(str(data_dict))}.png"
- filepath = os.path.join(self.storage_path, filename)
- plt.savefig(filepath, dpi=300, bbox_inches="tight")
- plt.close()
-
- return filepath
-
def set_storage_path(self, path: str) -> None:
- """Set the storage path for plots.
+ """Set the storage path for generated plots.
Args:
- path: New storage path
+ path: New directory path where plots will be saved.
"""
self.storage_path = path
self._ensure_storage_directory()
diff --git a/sparse_attention_hub/sparse_attention/efficient_attention/base.py b/sparse_attention_hub/sparse_attention/efficient_attention/base.py
index 0d1b5f01..06be8e79 100644
--- a/sparse_attention_hub/sparse_attention/efficient_attention/base.py
+++ b/sparse_attention_hub/sparse_attention/efficient_attention/base.py
@@ -1,9 +1,8 @@
"""Base classes for efficient attention mechanisms."""
-
from abc import abstractmethod
from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple, Type, cast
import torch
from torch import nn
@@ -13,13 +12,22 @@
@dataclass
class EfficientAttentionConfig(SparseAttentionConfig):
- """Configuration class for efficient attention mechanisms."""
+ """Configuration class for efficient attention mechanisms.
+
+ This is a base configuration class for production-ready sparse attention
+ implementations that prioritize efficiency and performance.
+ """
pass
class EfficientAttention(SparseAttention):
- """Abstract base class for efficient attention mechanisms."""
+ """Abstract base class for efficient attention mechanisms.
+
+ This class serves as the base for production-ready sparse attention
+ implementations that are optimized for performance and memory efficiency,
+ such as HashAttention and DoubleSparsity.
+ """
def __init__(self, sparse_attention_config: SparseAttentionConfig) -> None:
"""Initialize efficient attention mechanism.
@@ -54,19 +62,32 @@ def create_from_config(cls, config: SparseAttentionConfig) -> "EfficientAttentio
Args:
config: Configuration for the efficient attention mechanism.
+ Must be an instance of EfficientAttentionConfig.
Returns:
Instance of the efficient attention mechanism.
Raises:
TypeError: If config is not an EfficientAttentionConfig.
+ ValueError: If no implementation is found for the config type.
"""
if not isinstance(config, EfficientAttentionConfig):
raise TypeError(f"Expected EfficientAttentionConfig, got {type(config)}")
- # Import here to avoid circular imports
- from typing import Type, cast
+ registry = cls._get_implementation_registry()
+ concrete_class = cls._get_concrete_class(config, registry)
+ return concrete_class.create_from_config(config)
+ @classmethod
+ def _get_implementation_registry(
+ cls,
+ ) -> Dict[Type[EfficientAttentionConfig], Type["EfficientAttention"]]:
+ """Get the registry mapping config types to implementation classes.
+
+ Returns:
+ Dictionary mapping config types to their corresponding implementation classes.
+ """
+ # Import here to avoid circular imports
from .implementations import (
DoubleSparsity,
DoubleSparsityConfig,
@@ -74,26 +95,32 @@ def create_from_config(cls, config: SparseAttentionConfig) -> "EfficientAttentio
HashAttentionConfig,
)
- # Registry mapping config types to concrete efficient attention classes
- _EFFICIENT_ATTENTION_REGISTRY: Dict[
- Type[EfficientAttentionConfig], Type[EfficientAttention]
- ] = {
+ return {
DoubleSparsityConfig: DoubleSparsity,
HashAttentionConfig: HashAttention,
}
- # Look up the concrete class based on the config type
- concrete_class: Optional[
- Type[EfficientAttention]
- ] = _EFFICIENT_ATTENTION_REGISTRY.get(type(config))
+ @classmethod
+ def _get_concrete_class(
+ cls,
+ config: EfficientAttentionConfig,
+ registry: Dict[Type[EfficientAttentionConfig], Type["EfficientAttention"]],
+ ) -> Type["EfficientAttention"]:
+ """Get the concrete implementation class for the given configuration.
+
+ Args:
+ config: Configuration instance.
+ registry: Registry mapping config types to implementation classes.
+
+ Returns:
+ Concrete implementation class.
+
+ Raises:
+ ValueError: If no implementation is found for the config type.
+ """
+ concrete_class = registry.get(type(config))
if concrete_class is None:
raise ValueError(
f"No efficient attention class found for config type: {type(config)}"
)
-
- # Cast to help mypy understand the type
- concrete_class: Type[EfficientAttention] = cast(
- Type[EfficientAttention], concrete_class
- )
- # Call the concrete class's create_from_config method
- return concrete_class.create_from_config(config)
+ return cast(Type["EfficientAttention"], concrete_class)
diff --git a/sparse_attention_hub/sparse_attention/research_attention/base.py b/sparse_attention_hub/sparse_attention/research_attention/base.py
index 34189068..0fed826a 100644
--- a/sparse_attention_hub/sparse_attention/research_attention/base.py
+++ b/sparse_attention_hub/sparse_attention/research_attention/base.py
@@ -23,13 +23,28 @@
@dataclass
class ResearchAttentionConfig(SparseAttentionConfig):
- """Configuration class for research attention mechanisms."""
+ """Configuration class for research attention mechanisms.
+
+ This configuration specifies the masker components that will be applied
+ sequentially to create sparse attention patterns for research purposes.
+
+ Attributes:
+ masker_configs: List of masker configurations to apply in sequence.
+ """
masker_configs: List[MaskerConfig]
class ResearchAttention(SparseAttention):
- """Base class for research attention mechanisms with maskers."""
+ """Base class for research attention mechanisms with configurable maskers.
+
+ This class implements sparse attention by applying a sequence of maskers
+ to create custom attention patterns. It supports metrics logging and
+ validation of masker configurations.
+
+ Attributes:
+ maskers: List of research maskers to apply sequentially.
+ """
maskers: List[ResearchMasker]
@@ -42,24 +57,33 @@ def __init__(
Args:
sparse_attention_config: Configuration for the sparse attention mechanism.
- maskers: List of research maskers to apply.
+ maskers: List of research maskers to apply in sequence.
Raises:
ValueError: If more than one sampling masker is provided.
"""
super().__init__(sparse_attention_config)
+ self._validate_masker_configuration(maskers)
+ self.maskers = maskers
+
+ def _validate_masker_configuration(self, maskers: List[ResearchMasker]) -> None:
+ """Validate the masker configuration.
+
+ Args:
+ maskers: List of maskers to validate.
- # Validate that there's at most one sampling masker
- sampling_masker_count: int = sum(
+ Raises:
+ ValueError: If more than one sampling masker is provided.
+ """
+ sampling_masker_count = sum(
1 for masker in maskers if isinstance(masker, SamplingMasker)
)
if sampling_masker_count > 1:
raise ValueError(
- "Only one sampling masker supported for efficiency; consider implementing all sampling logic in one masker"
+ "Only one sampling masker supported for efficiency; "
+ "consider implementing all sampling logic in one masker"
)
- self.maskers = maskers
-
def custom_attention(
self,
module: nn.Module,
@@ -75,30 +99,106 @@ def custom_attention(
"""Compute research attention mechanism with masking.
Args:
- module: The attention module
- queries: Query tensor of shape (b, h, sk, d)
- keys: Key tensor of shape (b, h, sq, d)
- values: Value tensor of shape (b, h, sq, d)
- attention_mask: Optional attention mask of shape (b, h, sq, sk)
- scaling: Scaling factor for attention weights
- dropout: Dropout probability
- **kwargs: Additional keyword arguments
+ module: The attention module.
+ queries: Query tensor of shape (b, h, sk, d).
+ keys: Key tensor of shape (b, h, sq, d).
+ values: Value tensor of shape (b, h, sq, d).
+ attention_mask: Optional attention mask of shape (b, h, sq, sk).
+ scaling: Scaling factor for attention weights.
+ dropout: Dropout probability.
+ sparse_meta_data: Additional metadata for sparse attention computation.
+ **kwargs: Additional keyword arguments.
Returns:
Tuple of attention output and optional attention weights.
"""
- # Create an empty Mask object
- mask_shape: Tuple[int, int, int, int] = (
+ sparse_attention_mask = self._create_initial_mask(queries, keys)
+ sparse_attention_mask = self._apply_maskers(
+ sparse_attention_mask=sparse_attention_mask,
+ keys=keys,
+ queries=queries,
+ values=values,
+ attention_mask=attention_mask,
+ scaling=scaling,
+ dropout=dropout,
+ sparse_meta_data=sparse_meta_data,
+ **kwargs,
+ )
+
+ self._log_attention_density(sparse_attention_mask, kwargs)
+
+ attention_output, attention_weights = self._compute_masked_attention(
+ module=module,
+ queries=queries,
+ keys=keys,
+ values=values,
+ attention_mask=attention_mask,
+ scaling=scaling,
+ dropout=dropout,
+ sparse_attention_mask=sparse_attention_mask,
+ **kwargs,
+ )
+
+ self._log_attention_error(
+ module=module,
+ queries=queries,
+ keys=keys,
+ values=values,
+ attention_mask=attention_mask,
+ scaling=scaling,
+ dropout=dropout,
+ attention_output=attention_output,
+ **kwargs,
+ )
+
+ return attention_output, attention_weights
+
+ def _create_initial_mask(self, queries: torch.Tensor, keys: torch.Tensor) -> Mask:
+ """Create an initial empty mask for the attention computation.
+
+ Args:
+ queries: Query tensor.
+ keys: Key tensor.
+
+ Returns:
+ Empty mask with appropriate shape.
+ """
+ mask_shape = (
queries.shape[0],
queries.shape[1],
queries.shape[2],
keys.shape[2],
)
- sparse_attention_mask: Mask = Mask.create_empty_mask(
- mask_shape, dtype=queries.dtype
- )
+ return Mask.create_empty_mask(mask_shape, dtype=queries.dtype)
- # Apply all maskers sequentially, each one on the output of the previous one
+ def _apply_maskers(
+ self,
+ sparse_attention_mask: Mask,
+ keys: torch.Tensor,
+ queries: torch.Tensor,
+ values: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float,
+ sparse_meta_data: Dict[Any, Any],
+ **kwargs: Dict[str, Any],
+ ) -> Mask:
+ """Apply all maskers sequentially to create the final sparse attention mask.
+
+ Args:
+ sparse_attention_mask: Initial mask to start with.
+ keys: Key tensor.
+ queries: Query tensor.
+ values: Value tensor.
+ attention_mask: Optional attention mask.
+ scaling: Scaling factor.
+ dropout: Dropout probability.
+ sparse_meta_data: Additional metadata.
+ **kwargs: Additional keyword arguments.
+
+ Returns:
+ Final sparse attention mask after applying all maskers.
+ """
for masker in self.maskers:
sparse_attention_mask = masker.add_mask(
keys=keys,
@@ -111,19 +211,53 @@ def custom_attention(
previous_mask=sparse_attention_mask,
**kwargs,
)
+ return sparse_attention_mask
+ def _log_attention_density(
+ self, sparse_attention_mask: Mask, kwargs: Dict[str, Any]
+ ) -> None:
+ """Log attention density metric if enabled.
+
+ Args:
+ sparse_attention_mask: The sparse attention mask.
+ kwargs: Keyword arguments containing layer information.
+ """
if MicroMetricLogger().is_metric_enabled("research_attention_density"):
MicroMetricLogger().log(
"research_attention_density",
sparse_attention_mask.get_density(),
- metadata={"layer_idx": kwargs["layer_idx"]},
+ metadata={"layer_idx": kwargs.get("layer_idx")},
)
- # Call compute_masked_attention_output on the result of the last mask
- # Always request attention weights to match the expected return signature
- attention_output: torch.Tensor
- attention_weights: torch.Tensor
- attention_output, attention_weights = get_masked_attention_output(
+ def _compute_masked_attention(
+ self,
+ module: nn.Module,
+ queries: torch.Tensor,
+ keys: torch.Tensor,
+ values: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float,
+ sparse_attention_mask: Mask,
+ **kwargs: Dict[str, Any],
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Compute the masked attention output and weights.
+
+ Args:
+ module: The attention module.
+ queries: Query tensor.
+ keys: Key tensor.
+ values: Value tensor.
+ attention_mask: Optional attention mask.
+ scaling: Scaling factor.
+ dropout: Dropout probability.
+ sparse_attention_mask: Sparse attention mask.
+ **kwargs: Additional keyword arguments.
+
+ Returns:
+ Tuple of attention output and attention weights.
+ """
+ return get_masked_attention_output(
module=module,
queries=queries,
keys=keys,
@@ -136,6 +270,31 @@ def custom_attention(
**kwargs,
)
+ def _log_attention_error(
+ self,
+ module: nn.Module,
+ queries: torch.Tensor,
+ keys: torch.Tensor,
+ values: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ scaling: float,
+ dropout: float,
+ attention_output: torch.Tensor,
+ **kwargs: Dict[str, Any],
+ ) -> None:
+ """Log attention output error metric if enabled.
+
+ Args:
+ module: The attention module.
+ queries: Query tensor.
+ keys: Key tensor.
+ values: Value tensor.
+ attention_mask: Optional attention mask.
+ scaling: Scaling factor.
+ dropout: Dropout probability.
+ attention_output: Computed attention output.
+ **kwargs: Additional keyword arguments.
+ """
if MicroMetricLogger().is_metric_enabled("research_attention_output_error"):
true_attention_output, _ = get_true_attention_output(
module,
@@ -153,20 +312,19 @@ def custom_attention(
MicroMetricLogger().log(
"research_attention_output_error",
float(error.item()),
- metadata={"layer_idx": kwargs["layer_idx"]},
+ metadata={"layer_idx": kwargs.get("layer_idx")},
)
- return attention_output, attention_weights
-
@classmethod
def create_from_config(cls, config: SparseAttentionConfig) -> "ResearchAttention":
"""Create research attention instance from configuration.
Args:
config: Configuration for the research attention mechanism.
+ Must be an instance of ResearchAttentionConfig.
Returns:
- Instance of the research attention mechanism.
+ Instance of the research attention mechanism with configured maskers.
Raises:
TypeError: If config is not a ResearchAttentionConfig.
@@ -174,12 +332,23 @@ def create_from_config(cls, config: SparseAttentionConfig) -> "ResearchAttention
if not isinstance(config, ResearchAttentionConfig):
raise TypeError(f"Expected ResearchAttentionConfig, got {type(config)}")
- # Create ResearchMasker objects from the configs using the factory method
- maskers: List[ResearchMasker] = []
- for masker_config in config.masker_configs:
- masker: ResearchMasker = ResearchMasker.create_masker_from_config(
- masker_config
- )
- maskers.append(masker)
-
+ maskers = cls._create_maskers_from_config(config.masker_configs)
return cls(config, maskers)
+
+ @classmethod
+ def _create_maskers_from_config(
+ cls, masker_configs: List[MaskerConfig]
+ ) -> List[ResearchMasker]:
+ """Create research masker objects from their configurations.
+
+ Args:
+ masker_configs: List of masker configurations.
+
+ Returns:
+ List of configured research masker instances.
+ """
+ maskers = []
+ for masker_config in masker_configs:
+ masker = ResearchMasker.create_masker_from_config(masker_config)
+ maskers.append(masker)
+ return maskers
diff --git a/sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/random_sampling.py b/sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/random_sampling.py
index ed72b255..a5b06044 100644
--- a/sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/random_sampling.py
+++ b/sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/random_sampling.py
@@ -37,7 +37,9 @@ class RandomSamplingMaskerConfig(SamplingMaskerConfig):
no indices are sampled, while 1.0 means all indices are sampled.
"""
- sampling_rate: float # Float in range [0,1] representing fraction of indices to sample
+ sampling_rate: (
+ float # Float in range [0,1] representing fraction of indices to sample
+ )
def __post_init__(self) -> None:
"""Validate sampling_rate after initialization."""
diff --git a/sparse_attention_hub/sparse_attention/utils/hashattention_utils.py b/sparse_attention_hub/sparse_attention/utils/hashattention_utils.py
index ba8b5aea..ad30c215 100644
--- a/sparse_attention_hub/sparse_attention/utils/hashattention_utils.py
+++ b/sparse_attention_hub/sparse_attention/utils/hashattention_utils.py
@@ -1,10 +1,18 @@
+"""Utilities for converting and loading HashAttention weights.
+
+This module provides functions to convert USA (Universal Sparse Attention)
+checkpoint weights to HashAttention format and load them for model use.
+
+Note:
+ Currently only supports 3-layered MLPs. Support for other configurations
+ needs to be implemented.
+"""
+
import pickle
from typing import Dict, List
import torch
-""" This only works for 3 layered MLPs. Need to fix for other cases"""
-
def convert_usa_weights_to_hash_attention(
usa_checkpoint_path: str,
@@ -16,90 +24,24 @@ def convert_usa_weights_to_hash_attention(
"""Convert USA module weights to HashAttentionTopKMasker format.
Args:
- usa_checkpoint_path: Path to USA checkpoint file
- num_layers: Number of layers in the model
- num_heads: Number of attention heads
- num_mlp_layers: Number of MLP layers (default: 3)
- device: Device to load weights on
+ usa_checkpoint_path: Path to the USA checkpoint file.
+ num_layers: Number of transformer layers in the model. Defaults to 32.
+ num_heads: Number of attention heads per layer. Defaults to 32.
+ num_mlp_layers: Number of MLP layers in the hash transformation. Defaults to 3.
+ device: Device to load the weights on ("cpu", "cuda", etc.). Defaults to "cpu".
Returns:
- Dictionary of layer-wise weights in HashAttention format
+ Dictionary mapping layer indices to their corresponding weights in HashAttention format.
+ Each layer contains "query_matrix", "query_bias", "key_matrix", and "key_bias" lists.
"""
-
print(f"Loading USA weights from {usa_checkpoint_path}")
usa_state_dict = torch.load(usa_checkpoint_path, map_location=device)
hat_weights = {}
-
for layer_idx in range(num_layers):
- layer_weights = {
- "query_matrix": [],
- "query_bias": [],
- "key_matrix": [],
- "key_bias": [],
- }
-
- # Collect weights for all heads in this layer
- query_matrices_per_layer = [[] for _ in range(num_mlp_layers)]
- query_biases_per_layer = [[] for _ in range(num_mlp_layers)]
- key_matrices_per_layer = [[] for _ in range(num_mlp_layers)]
- key_biases_per_layer = [[] for _ in range(num_mlp_layers)]
-
- for head_idx in range(num_heads):
- query_prefix = f"{layer_idx}.learning_to_hash_transformation_q.{head_idx}"
- key_prefix = f"{layer_idx}.learning_to_hash_transformation_k.{head_idx}"
-
- # Extract weights from MLP layers (linear layers at indices 0, 2, 4, ...)
- linear_indices = [i * 2 for i in range(num_mlp_layers)]
- for i, linear_idx in enumerate(linear_indices):
- # Query weights
- weight_key = f"{query_prefix}.{linear_idx}.weight"
- bias_key = f"{query_prefix}.{linear_idx}.bias"
-
- if weight_key in usa_state_dict:
- weight = usa_state_dict[
- weight_key
- ].t() # Transpose to (in_features, out_features)
- query_matrices_per_layer[i].append(weight)
-
- if bias_key in usa_state_dict:
- query_biases_per_layer[i].append(usa_state_dict[bias_key])
- else:
- query_biases_per_layer[i].append(
- torch.zeros(usa_state_dict[weight_key].shape[0])
- )
-
- # Key weights
- weight_key = f"{key_prefix}.{linear_idx}.weight"
- bias_key = f"{key_prefix}.{linear_idx}.bias"
-
- if weight_key in usa_state_dict:
- weight = usa_state_dict[
- weight_key
- ].t() # Transpose to (in_features, out_features)
- key_matrices_per_layer[i].append(weight)
-
- if bias_key in usa_state_dict:
- key_biases_per_layer[i].append(usa_state_dict[bias_key])
- else:
- key_biases_per_layer[i].append(
- torch.zeros(usa_state_dict[weight_key].shape[0])
- )
-
- # Stack all heads for each layer
- for i in range(num_mlp_layers):
- if query_matrices_per_layer[i]:
- layer_weights["query_matrix"].append(
- torch.stack(query_matrices_per_layer[i])
- )
- layer_weights["query_bias"].append(
- torch.stack(query_biases_per_layer[i])
- )
- layer_weights["key_matrix"].append(
- torch.stack(key_matrices_per_layer[i])
- )
- layer_weights["key_bias"].append(torch.stack(key_biases_per_layer[i]))
-
+ layer_weights = _convert_layer_weights(
+ usa_state_dict, layer_idx, num_heads, num_mlp_layers
+ )
hat_weights[layer_idx] = layer_weights
print(
@@ -108,6 +50,204 @@ def convert_usa_weights_to_hash_attention(
return hat_weights
+def _convert_layer_weights(
+ usa_state_dict: Dict[str, torch.Tensor],
+ layer_idx: int,
+ num_heads: int,
+ num_mlp_layers: int,
+) -> Dict[str, List[torch.Tensor]]:
+ """Convert weights for a single layer.
+
+ Args:
+ usa_state_dict: USA checkpoint state dictionary.
+ layer_idx: Index of the layer to process.
+ num_heads: Number of attention heads.
+ num_mlp_layers: Number of MLP layers.
+
+ Returns:
+ Dictionary containing converted weights for the layer.
+ """
+ layer_weights = {
+ "query_matrix": [],
+ "query_bias": [],
+ "key_matrix": [],
+ "key_bias": [],
+ }
+
+ # Collect weights for all heads in this layer
+ query_matrices_per_layer = [[] for _ in range(num_mlp_layers)]
+ query_biases_per_layer = [[] for _ in range(num_mlp_layers)]
+ key_matrices_per_layer = [[] for _ in range(num_mlp_layers)]
+ key_biases_per_layer = [[] for _ in range(num_mlp_layers)]
+
+ for head_idx in range(num_heads):
+ _extract_head_weights(
+ usa_state_dict=usa_state_dict,
+ layer_idx=layer_idx,
+ head_idx=head_idx,
+ num_mlp_layers=num_mlp_layers,
+ query_matrices_per_layer=query_matrices_per_layer,
+ query_biases_per_layer=query_biases_per_layer,
+ key_matrices_per_layer=key_matrices_per_layer,
+ key_biases_per_layer=key_biases_per_layer,
+ )
+
+ # Stack all heads for each MLP layer
+ _stack_head_weights(
+ layer_weights=layer_weights,
+ num_mlp_layers=num_mlp_layers,
+ query_matrices_per_layer=query_matrices_per_layer,
+ query_biases_per_layer=query_biases_per_layer,
+ key_matrices_per_layer=key_matrices_per_layer,
+ key_biases_per_layer=key_biases_per_layer,
+ )
+
+ return layer_weights
+
+
+def _extract_head_weights(
+ usa_state_dict: Dict[str, torch.Tensor],
+ layer_idx: int,
+ head_idx: int,
+ num_mlp_layers: int,
+ query_matrices_per_layer: List[List[torch.Tensor]],
+ query_biases_per_layer: List[List[torch.Tensor]],
+ key_matrices_per_layer: List[List[torch.Tensor]],
+ key_biases_per_layer: List[List[torch.Tensor]],
+) -> None:
+ """Extract weights for a single attention head.
+
+ Args:
+ usa_state_dict: USA checkpoint state dictionary.
+ layer_idx: Index of the current layer.
+ head_idx: Index of the current attention head.
+ num_mlp_layers: Number of MLP layers.
+ query_matrices_per_layer: Storage for query weight matrices.
+ query_biases_per_layer: Storage for query bias vectors.
+ key_matrices_per_layer: Storage for key weight matrices.
+ key_biases_per_layer: Storage for key bias vectors.
+ """
+ query_prefix = f"{layer_idx}.learning_to_hash_transformation_q.{head_idx}"
+ key_prefix = f"{layer_idx}.learning_to_hash_transformation_k.{head_idx}"
+
+ # Extract weights from MLP layers (linear layers at indices 0, 2, 4, ...)
+ linear_indices = [i * 2 for i in range(num_mlp_layers)]
+
+ for i, linear_idx in enumerate(linear_indices):
+ _extract_query_weights(
+ usa_state_dict,
+ query_prefix,
+ linear_idx,
+ i,
+ query_matrices_per_layer,
+ query_biases_per_layer,
+ )
+ _extract_key_weights(
+ usa_state_dict,
+ key_prefix,
+ linear_idx,
+ i,
+ key_matrices_per_layer,
+ key_biases_per_layer,
+ )
+
+
+def _extract_query_weights(
+ usa_state_dict: Dict[str, torch.Tensor],
+ query_prefix: str,
+ linear_idx: int,
+ mlp_layer_idx: int,
+ query_matrices_per_layer: List[List[torch.Tensor]],
+ query_biases_per_layer: List[List[torch.Tensor]],
+) -> None:
+ """Extract query weights for a specific MLP layer.
+
+ Args:
+ usa_state_dict: USA checkpoint state dictionary.
+ query_prefix: Prefix for query weight keys.
+ linear_idx: Index of the linear layer.
+ mlp_layer_idx: Index of the MLP layer.
+ query_matrices_per_layer: Storage for query weight matrices.
+ query_biases_per_layer: Storage for query bias vectors.
+ """
+ weight_key = f"{query_prefix}.{linear_idx}.weight"
+ bias_key = f"{query_prefix}.{linear_idx}.bias"
+
+ if weight_key in usa_state_dict:
+ # Transpose to (in_features, out_features)
+ weight = usa_state_dict[weight_key].t()
+ query_matrices_per_layer[mlp_layer_idx].append(weight)
+
+ if bias_key in usa_state_dict:
+ query_biases_per_layer[mlp_layer_idx].append(usa_state_dict[bias_key])
+ else:
+ # Create zero bias if not present
+ zero_bias = torch.zeros(usa_state_dict[weight_key].shape[0])
+ query_biases_per_layer[mlp_layer_idx].append(zero_bias)
+
+
+def _extract_key_weights(
+ usa_state_dict: Dict[str, torch.Tensor],
+ key_prefix: str,
+ linear_idx: int,
+ mlp_layer_idx: int,
+ key_matrices_per_layer: List[List[torch.Tensor]],
+ key_biases_per_layer: List[List[torch.Tensor]],
+) -> None:
+ """Extract key weights for a specific MLP layer.
+
+ Args:
+ usa_state_dict: USA checkpoint state dictionary.
+ key_prefix: Prefix for key weight keys.
+ linear_idx: Index of the linear layer.
+ mlp_layer_idx: Index of the MLP layer.
+ key_matrices_per_layer: Storage for key weight matrices.
+ key_biases_per_layer: Storage for key bias vectors.
+ """
+ weight_key = f"{key_prefix}.{linear_idx}.weight"
+ bias_key = f"{key_prefix}.{linear_idx}.bias"
+
+ if weight_key in usa_state_dict:
+ # Transpose to (in_features, out_features)
+ weight = usa_state_dict[weight_key].t()
+ key_matrices_per_layer[mlp_layer_idx].append(weight)
+
+ if bias_key in usa_state_dict:
+ key_biases_per_layer[mlp_layer_idx].append(usa_state_dict[bias_key])
+ else:
+ # Create zero bias if not present
+ zero_bias = torch.zeros(usa_state_dict[weight_key].shape[0])
+ key_biases_per_layer[mlp_layer_idx].append(zero_bias)
+
+
+def _stack_head_weights(
+ layer_weights: Dict[str, List[torch.Tensor]],
+ num_mlp_layers: int,
+ query_matrices_per_layer: List[List[torch.Tensor]],
+ query_biases_per_layer: List[List[torch.Tensor]],
+ key_matrices_per_layer: List[List[torch.Tensor]],
+ key_biases_per_layer: List[List[torch.Tensor]],
+) -> None:
+ """Stack weights from all heads for each MLP layer.
+
+ Args:
+ layer_weights: Dictionary to store the stacked weights.
+ num_mlp_layers: Number of MLP layers.
+ query_matrices_per_layer: Query weight matrices for all heads.
+ query_biases_per_layer: Query bias vectors for all heads.
+ key_matrices_per_layer: Key weight matrices for all heads.
+ key_biases_per_layer: Key bias vectors for all heads.
+ """
+ for i in range(num_mlp_layers):
+ if query_matrices_per_layer[i]:
+ layer_weights["query_matrix"].append(
+ torch.stack(query_matrices_per_layer[i])
+ )
+ layer_weights["query_bias"].append(torch.stack(query_biases_per_layer[i]))
+ layer_weights["key_matrix"].append(torch.stack(key_matrices_per_layer[i]))
+ layer_weights["key_bias"].append(torch.stack(key_biases_per_layer[i]))
+
+
def create_hat_weights_file_from_usa(
usa_checkpoint_path: str,
target_hat_path: str,
@@ -116,19 +256,21 @@ def create_hat_weights_file_from_usa(
num_mlp_layers: int = 3,
device: str = "cpu",
) -> None:
- """Create HAT weights file from USA checkpoint.
+ """Create HashAttention weights file from USA checkpoint.
+
+ This function converts USA checkpoint weights to HashAttention format
+ and saves them as a pickle file for later use.
Args:
- usa_checkpoint_path: Path to USA checkpoint file
- target_hat_path: Path where HAT weights file will be saved
- num_layers: Number of layers in the model
- num_heads: Number of attention heads
- num_mlp_layers: Number of MLP layers
- device: Device to load weights on
+ usa_checkpoint_path: Path to the input USA checkpoint file.
+ target_hat_path: Path where the HashAttention weights file will be saved.
+ num_layers: Number of transformer layers in the model. Defaults to 32.
+ num_heads: Number of attention heads per layer. Defaults to 32.
+ num_mlp_layers: Number of MLP layers in the hash transformation. Defaults to 3.
+ device: Device to load the weights on ("cpu", "cuda", etc.). Defaults to "cpu".
"""
print("Creating HAT weights file from USA checkpoint...")
- # Convert USA weights to HAT format
hat_weights = convert_usa_weights_to_hash_attention(
usa_checkpoint_path=usa_checkpoint_path,
num_layers=num_layers,
@@ -137,34 +279,70 @@ def create_hat_weights_file_from_usa(
device=device,
)
- # Save to pickle file
- with open(target_hat_path, "wb") as f:
- pickle.dump(hat_weights, f)
-
+ _save_weights_to_file(hat_weights, target_hat_path)
print(f"โ
HAT weights saved to {target_hat_path}")
+def _save_weights_to_file(
+ hat_weights: Dict[int, Dict[str, List[torch.Tensor]]], target_path: str
+) -> None:
+ """Save HashAttention weights to a pickle file.
+
+ Args:
+ hat_weights: Dictionary of HashAttention weights to save.
+ target_path: Path where the weights file will be saved.
+ """
+ with open(target_path, "wb") as f:
+ pickle.dump(hat_weights, f)
+
+
def load_hat_weights(
hat_weights_path: str, device: str = "cpu"
) -> Dict[int, Dict[str, List[torch.Tensor]]]:
- """Load HAT weights from pickle file.
+ """Load HashAttention weights from a pickle file.
+
+ This function loads previously saved HashAttention weights and moves
+ them to the specified device.
Args:
- hat_weights_path: Path to HAT weights pickle file
- device: Device to load weights on
+ hat_weights_path: Path to the HashAttention weights pickle file.
+ device: Device to move the loaded weights to ("cpu", "cuda", etc.). Defaults to "cpu".
Returns:
- Dictionary of layer-wise weights in HashAttention format
+ Dictionary mapping layer indices to their corresponding weights in HashAttention format.
+ Each layer contains "query_matrix", "query_bias", "key_matrix", and "key_bias" lists.
"""
print(f"Loading HAT weights from {hat_weights_path}")
- with open(hat_weights_path, "rb") as f:
- hat_weights = pickle.load(f)
-
- # Move weights to specified device
- for layer_idx, layer_weights in hat_weights.items():
- for key, value in layer_weights.items():
- hat_weights[layer_idx][key] = [tensor.to(device) for tensor in value]
+ hat_weights = _load_weights_from_file(hat_weights_path)
+ _move_weights_to_device(hat_weights, device)
print(f"โ
Loaded HAT weights for {len(hat_weights)} layers")
return hat_weights
+
+
+def _load_weights_from_file(file_path: str) -> Dict[int, Dict[str, List[torch.Tensor]]]:
+ """Load weights from a pickle file.
+
+ Args:
+ file_path: Path to the pickle file.
+
+ Returns:
+ Dictionary of loaded weights.
+ """
+ with open(file_path, "rb") as f:
+ return pickle.load(f)
+
+
+def _move_weights_to_device(
+ hat_weights: Dict[int, Dict[str, List[torch.Tensor]]], device: str
+) -> None:
+ """Move all weights to the specified device.
+
+ Args:
+ hat_weights: Dictionary of HashAttention weights to move.
+ device: Target device for the weights.
+ """
+ for layer_idx, layer_weights in hat_weights.items():
+ for key, tensor_list in layer_weights.items():
+ hat_weights[layer_idx][key] = [tensor.to(device) for tensor in tensor_list]
diff --git a/sparse_attention_hub/sparse_attention/utils/kv_utils.py b/sparse_attention_hub/sparse_attention/utils/kv_utils.py
index afc52d46..81f41def 100644
--- a/sparse_attention_hub/sparse_attention/utils/kv_utils.py
+++ b/sparse_attention_hub/sparse_attention/utils/kv_utils.py
@@ -1,4 +1,5 @@
""" Utility functions for common kv manipulation. """
+
import torch
diff --git a/tests/integration/test_adapter_integration.py b/tests/integration/test_adapter_integration.py
index 0e88542d..b5dc85a4 100644
--- a/tests/integration/test_adapter_integration.py
+++ b/tests/integration/test_adapter_integration.py
@@ -48,14 +48,18 @@ class TestAdapterIntegration:
def setup_method(self) -> None:
"""Set up test fixtures - reset ModelServer singleton."""
from sparse_attention_hub.adapters.model_servers.base import ModelServer
+
ModelServer._instance = None
def teardown_method(self) -> None:
"""Clean up after each test - reset ModelServer singleton."""
from sparse_attention_hub.adapters.model_servers.base import ModelServer
+
ModelServer._instance = None
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_full_request_processing_flow(
self, mock_tokenizer, mock_model, sparse_attention_config
@@ -141,7 +145,9 @@ def test_full_request_processing_flow(
assert mock_tokenizer_instance.encode.call_count == 2
mock_tokenizer_instance.decode.assert_called_once()
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_multiple_questions_processing(
self, mock_tokenizer, mock_model, sparse_attention_config
@@ -246,7 +252,9 @@ def test_multiple_questions_processing(
assert mock_tokenizer_instance.encode.call_count == 3
assert mock_tokenizer_instance.decode.call_count == 2
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_mode_switching_integration(
self, mock_tokenizer, mock_model, sparse_attention_config
@@ -287,7 +295,9 @@ def test_mode_switching_integration(
# Should not raise any errors
pass
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_custom_attention_function_integration(
self, mock_tokenizer, mock_model, sparse_attention_config
@@ -350,7 +360,9 @@ def test_custom_attention_function_integration(
mock_custom_attention.assert_called_once()
assert result is not None
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_error_handling_integration(
self, mock_tokenizer, mock_model, sparse_attention_config
@@ -399,7 +411,9 @@ def test_error_handling_integration(
# missing sparse_meta_data
)
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_adapter_with_device_configuration(
self, mock_tokenizer, mock_model, sparse_attention_config
@@ -453,7 +467,9 @@ def test_adapter_with_device_configuration(
"test-model", device_map="cuda", torch_dtype=torch.float16
)
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_adapter_cleanup_on_exception(
self, mock_tokenizer, mock_model, sparse_attention_config
diff --git a/tests/integration/test_benchmark.py b/tests/integration/test_benchmark.py
index c67b75a3..0c2807d1 100644
--- a/tests/integration/test_benchmark.py
+++ b/tests/integration/test_benchmark.py
@@ -50,7 +50,9 @@ def sparse_attention_config(masker_config):
class TestBenchmarkAdapterIntegration:
"""Test benchmark integration with real adapter interfaces."""
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_benchmark_with_real_adapter_interface(
self,
@@ -118,7 +120,9 @@ def mock_process_request(
assert (result_path / "raw_results.csv").exists()
assert (result_path / "metrics.json").exists()
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_dense_only_adapter_integration(
self, mock_tokenizer_class, mock_model_class, temp_result_dir
@@ -382,6 +386,7 @@ class TestErrorHandlingIntegration:
def test_adapter_failure_recovery(self, temp_result_dir):
"""Test benchmark recovery when adapter fails intermittently."""
+
# Create adapter that fails on certain contexts
def failing_process_request(
request: Request,
diff --git a/tests/unit/adapters/model_servers/test_huggingface.py b/tests/unit/adapters/model_servers/test_huggingface.py
index 5ea48060..256d32c0 100644
--- a/tests/unit/adapters/model_servers/test_huggingface.py
+++ b/tests/unit/adapters/model_servers/test_huggingface.py
@@ -485,8 +485,10 @@ def test_validate_gpu_availability_success(
mock_cuda_available.return_value = True
mock_device_count.return_value = 4
- with patch("torch.cuda.device"), patch("torch.zeros") as mock_zeros, patch(
- "torch.cuda.empty_cache"
+ with (
+ patch("torch.cuda.device"),
+ patch("torch.zeros") as mock_zeros,
+ patch("torch.cuda.empty_cache"),
):
mock_zeros.return_value = Mock()
result = self.server.validate_gpu_availability(0)
diff --git a/tests/unit/adapters/test_adapters.py b/tests/unit/adapters/test_adapters.py
index 60b36b0d..78c07509 100644
--- a/tests/unit/adapters/test_adapters.py
+++ b/tests/unit/adapters/test_adapters.py
@@ -20,6 +20,7 @@
)
from sparse_attention_hub.adapters.model_servers.base import ModelServer
+
@pytest.mark.unit
class TestRequest:
"""Test the Request class."""
@@ -147,12 +148,14 @@ def setup_method(self) -> None:
masker_configs=[self.masker_config]
)
ModelServer._instance = None
-
+
def teardown_method(self) -> None:
"""Clean up after each test."""
ModelServer._instance = None
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_create_model(self, mock_tokenizer, mock_model) -> None:
"""Test model creation."""
@@ -186,7 +189,9 @@ def test_create_model(self, mock_tokenizer, mock_model) -> None:
"test-model", torch_dtype=torch.float16
)
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_create_model_with_torch_dtype(self, mock_tokenizer, mock_model) -> None:
"""Test model creation with torch_dtype parameter."""
@@ -217,7 +222,9 @@ def test_create_model_with_torch_dtype(self, mock_tokenizer, mock_model) -> None
# Check that adapter was created successfully
assert adapter is not None
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_create_model_with_existing_pad_token(
self, mock_tokenizer, mock_model
@@ -243,7 +250,9 @@ def test_create_model_with_existing_pad_token(
# Check that pad_token was not changed
assert adapter.tokenizer.pad_token == ""
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_get_custom_attention_function(self, mock_tokenizer, mock_model) -> None:
"""Test get_custom_attention_function returns a callable."""
@@ -265,7 +274,9 @@ def test_get_custom_attention_function(self, mock_tokenizer, mock_model) -> None
custom_fn = adapter.get_custom_attention_function(adapter.sparse_attention)
assert callable(custom_fn)
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_generate_unique_attention_name(self, mock_tokenizer, mock_model) -> None:
"""Test unique attention name generation."""
@@ -290,7 +301,9 @@ def test_generate_unique_attention_name(self, mock_tokenizer, mock_model) -> Non
assert name2.startswith("sparse_attention_")
assert name1 != name2 # Should be unique
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_enable_sparse_mode_when_not_available(
self, mock_tokenizer, mock_model
@@ -317,7 +330,9 @@ def test_enable_sparse_mode_when_not_available(
exc_info.value
)
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_enable_sparse_and_dense_modes(self, mock_tokenizer, mock_model) -> None:
"""Test enable_sparse_mode and enable_dense_mode context managers."""
@@ -352,7 +367,9 @@ def test_enable_sparse_and_dense_modes(self, mock_tokenizer, mock_model) -> None
assert first_name == second_name
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_inheritance(self, mock_tokenizer, mock_model) -> None:
"""Test that ModelAdapterHF properly inherits from ModelAdapter."""
@@ -372,7 +389,9 @@ def test_inheritance(self, mock_tokenizer, mock_model) -> None:
assert isinstance(adapter, ModelHubAdapterInterface)
assert isinstance(adapter, SparseAttentionAdapterInterface)
- @patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM")
+ @patch(
+ "sparse_attention_hub.adapters.model_servers.huggingface.AutoModelForCausalLM"
+ )
@patch("sparse_attention_hub.adapters.model_servers.huggingface.AutoTokenizer")
def test_adapter_properties(self, mock_tokenizer, mock_model) -> None:
"""Test adapter properties are set correctly."""
diff --git a/tests/unit/benchmark/test_longbench.py b/tests/unit/benchmark/test_longbench.py
index 3ff12940..2810ec00 100644
--- a/tests/unit/benchmark/test_longbench.py
+++ b/tests/unit/benchmark/test_longbench.py
@@ -314,7 +314,7 @@ def test_longbench_post_run_evaluate_missing_length_column(self):
"task": ["qasper_e"],
"predicted_answer": ["Answer 1"],
"answers": [["Truth 1"]],
- "all_classes": [[]]
+ "all_classes": [[]],
# Missing 'length' column
}
)
diff --git a/tests/unit/sparse_attention/research_attention/maskers/fixed/implementations/test_hashattention_top_k.py b/tests/unit/sparse_attention/research_attention/maskers/fixed/implementations/test_hashattention_top_k.py
index 680c6090..acc49ed1 100644
--- a/tests/unit/sparse_attention/research_attention/maskers/fixed/implementations/test_hashattention_top_k.py
+++ b/tests/unit/sparse_attention/research_attention/maskers/fixed/implementations/test_hashattention_top_k.py
@@ -5,6 +5,7 @@
:date: 2025-06-29
:summary: Tests for HashAttentionTopKMasker implementation.
"""
+
import os
import pickle
import tempfile
diff --git a/tests/unit/sparse_attention/test_configs_and_factories.py b/tests/unit/sparse_attention/test_configs_and_factories.py
index 095ed934..040730ae 100644
--- a/tests/unit/sparse_attention/test_configs_and_factories.py
+++ b/tests/unit/sparse_attention/test_configs_and_factories.py
@@ -1,4 +1,5 @@
"""Tests for sparse attention config classes and create_from_config methods."""
+
import pytest
diff --git a/tests/unit/sparse_attention/utils/test_mask_attention_utils.py b/tests/unit/sparse_attention/utils/test_mask_attention_utils.py
index 9ef1b93d..13b1756a 100644
--- a/tests/unit/sparse_attention/utils/test_mask_attention_utils.py
+++ b/tests/unit/sparse_attention/utils/test_mask_attention_utils.py
@@ -6,7 +6,6 @@
:summary: Tests for sparse attention. This file is part of the Sparse Attention Hub project.
"""
-
import mock
import numpy as np
import pytest
diff --git a/tutorials/py_examples/02_streaming_llm_demo_hf.py b/tutorials/py_examples/02_streaming_llm_demo_hf.py
index dfcbb084..afa4c7ee 100644
--- a/tutorials/py_examples/02_streaming_llm_demo_hf.py
+++ b/tutorials/py_examples/02_streaming_llm_demo_hf.py
@@ -1,31 +1,35 @@
#!/usr/bin/env python3
-'''
+"""
You can also choose to have more control instead calling process_request directly.
-'''
+"""
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import time
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+from sparse_attention_hub.sparse_attention.research_attention import (
+ ResearchAttentionConfig,
+)
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
- LocalMaskerConfig, SinkMaskerConfig
+ LocalMaskerConfig,
+ SinkMaskerConfig,
)
from sparse_attention_hub.adapters.huggingface import ModelAdapterHF
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_name = "meta-llama/Llama-3.1-8B-Instruct"
-#D model_name = "microsoft/DialoGPT-medium"
-#D model_name = "google/gemma-3n-E4B-it"
+# D model_name = "microsoft/DialoGPT-medium"
+# D model_name = "google/gemma-3n-E4B-it"
model_name = "microsoft/Phi-4-mini-instruct"
model_name = "mistralai/Mistral-7B-Instruct-v0.3"
model_name = "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B"
import os
-os.chdir('/data/apdesai/code/sparse-attention-hub')
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+os.chdir("/data/apdesai/code/sparse-attention-hub")
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
# Create StreamingLLM configuration: Sink (4 tokens) + Local (64 tokens)
sink_config = SinkMaskerConfig(sink_size=4)
@@ -34,10 +38,12 @@
print(f"โ
StreamingLLM config: Sink(4) + Local(64)")
-adapter = ModelAdapterHF(model_name="meta-llama/Llama-3.1-8B-Instruct",
- sparse_attention_config=research_config,
- model_kwargs={"torch_dtype": torch.bfloat16},
- device="cuda:0")
+adapter = ModelAdapterHF(
+ model_name="meta-llama/Llama-3.1-8B-Instruct",
+ sparse_attention_config=research_config,
+ model_kwargs={"torch_dtype": torch.bfloat16},
+ device="cuda:0",
+)
test_content = """
@@ -51,11 +57,13 @@
print("Running Full Attention")
-messages = [
- {"role": "user", "content": test_content}
-]
-test_content = adapter.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-input_ids = adapter.tokenizer.encode(test_content, return_tensors="pt", add_special_tokens=False).to(device)
+messages = [{"role": "user", "content": test_content}]
+test_content = adapter.tokenizer.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True
+)
+input_ids = adapter.tokenizer.encode(
+ test_content, return_tensors="pt", add_special_tokens=False
+).to(device)
output_ids = adapter.model.generate(input_ids, max_new_tokens=50, do_sample=False)
@@ -66,10 +74,10 @@
with adapter.enable_sparse_mode():
- output_ids = adapter.model.generate(input_ids, max_new_tokens=50, do_sample=False, sparse_meta_data={})
+ output_ids = adapter.model.generate(
+ input_ids, max_new_tokens=50, do_sample=False, sparse_meta_data={}
+ )
# model.generate() validates kwargs and will not let sparse_meta_data to be passed in which is required
# Locally, I removed the validation for this to work. TODO(we can write our own generate function for utility)
output_text = adapter.tokenizer.decode(output_ids[0], skip_special_tokens=True)
print("Streaming Attention Response: ", output_text)
-
-
diff --git a/tutorials/py_examples/02_streaming_llm_tutorial.py b/tutorials/py_examples/02_streaming_llm_tutorial.py
index 3e95ae57..aee56ebf 100644
--- a/tutorials/py_examples/02_streaming_llm_tutorial.py
+++ b/tutorials/py_examples/02_streaming_llm_tutorial.py
@@ -1,17 +1,21 @@
import os
-os.chdir('/data/apdesai/code/sparse-attention-hub')
+
+os.chdir("/data/apdesai/code/sparse-attention-hub")
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import time
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+from sparse_attention_hub.sparse_attention.research_attention import (
+ ResearchAttentionConfig,
+)
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
- LocalMaskerConfig, SinkMaskerConfig
+ LocalMaskerConfig,
+ SinkMaskerConfig,
)
from sparse_attention_hub.adapters.huggingface import ModelAdapterHF
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
sink_config = SinkMaskerConfig(sink_size=4)
@@ -25,10 +29,12 @@
# model_name = "deepseek-ai/DeepSeek-R1-0528-Qwen3-8B"
model_name = "meta-llama/Llama-3.1-8B-Instruct"
-adapter = ModelAdapterHF(model_name=model_name,
- sparse_attention_config=research_config,
- model_kwargs={"torch_dtype": torch.bfloat16, "device_map": "cuda"},
- device="cuda")
+adapter = ModelAdapterHF(
+ model_name=model_name,
+ sparse_attention_config=research_config,
+ model_kwargs={"torch_dtype": torch.bfloat16, "device_map": "cuda"},
+ device="cuda",
+)
from sparse_attention_hub.adapters import Request
@@ -41,18 +47,16 @@
This approach maintains performance while reducing computational costs for long sequences.
"""
test_questions = [
- "Summarize the above in a single title with less than 10 words. Given only the title.",
- "What are other attention mechanisms that are used in the field of LLMs?"
+ "Summarize the above in a single title with less than 10 words. Given only the title.",
+ "What are other attention mechanisms that are used in the field of LLMs?",
]
request = Request(
context=test_context,
questions=test_questions,
- )
+)
print("Running Streaming Attention on Question")
-response=adapter.process_request(request)
-response_text=response.responses
+response = adapter.process_request(request)
+response_text = response.responses
print("Streaming Attention Response: ", response_text)
-
-
diff --git a/tutorials/py_examples/03_hashattention_demo.py b/tutorials/py_examples/03_hashattention_demo.py
index 33bec796..0985c213 100644
--- a/tutorials/py_examples/03_hashattention_demo.py
+++ b/tutorials/py_examples/03_hashattention_demo.py
@@ -12,96 +12,117 @@
from typing import Dict, List
from transformers import AutoTokenizer, AutoModelForCausalLM
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+from sparse_attention_hub.sparse_attention.research_attention import (
+ ResearchAttentionConfig,
+)
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
- LocalMaskerConfig, SinkMaskerConfig, HashAttentionTopKMaskerConfig
+ LocalMaskerConfig,
+ SinkMaskerConfig,
+ HashAttentionTopKMaskerConfig,
)
from sparse_attention_hub.adapters import ModelAdapterHF
from sparse_attention_hub.adapters import Request
+
def convert_usa_weights_to_hash_attention(
- usa_checkpoint_path: str,
- num_layers: int = 32,
+ usa_checkpoint_path: str,
+ num_layers: int = 32,
num_heads: int = 32,
- device: str = 'cpu'
+ device: str = "cpu",
) -> Dict[int, Dict[str, List[torch.Tensor]]]:
"""Convert USA module weights to HashAttentionTopKMasker format."""
-
+
print(f"Loading USA weights from {usa_checkpoint_path}")
usa_state_dict = torch.load(usa_checkpoint_path, map_location=device)
-
+
hat_weights = {}
-
+
for layer_idx in range(num_layers):
layer_weights = {
"query_matrix": [],
"query_bias": [],
"key_matrix": [],
- "key_bias": []
+ "key_bias": [],
}
-
+
# Collect weights for all heads in this layer
query_matrices_per_layer = [[] for _ in range(3)] # 3 linear layers
query_biases_per_layer = [[] for _ in range(3)]
key_matrices_per_layer = [[] for _ in range(3)]
key_biases_per_layer = [[] for _ in range(3)]
-
+
for head_idx in range(num_heads):
query_prefix = f"{layer_idx}.learning_to_hash_transformation_q.{head_idx}"
key_prefix = f"{layer_idx}.learning_to_hash_transformation_k.{head_idx}"
-
+
# Extract weights from 3-layer MLP (linear layers at indices 0, 2, 4)
for i, linear_idx in enumerate([0, 2, 4]):
# Query weights
weight_key = f"{query_prefix}.{linear_idx}.weight"
bias_key = f"{query_prefix}.{linear_idx}.bias"
-
+
if weight_key in usa_state_dict:
- weight = usa_state_dict[weight_key].t() # Transpose to (in_features, out_features)
+ weight = usa_state_dict[
+ weight_key
+ ].t() # Transpose to (in_features, out_features)
query_matrices_per_layer[i].append(weight)
-
+
if bias_key in usa_state_dict:
query_biases_per_layer[i].append(usa_state_dict[bias_key])
else:
- query_biases_per_layer[i].append(torch.zeros(usa_state_dict[weight_key].shape[0]))
-
+ query_biases_per_layer[i].append(
+ torch.zeros(usa_state_dict[weight_key].shape[0])
+ )
+
# Key weights
weight_key = f"{key_prefix}.{linear_idx}.weight"
bias_key = f"{key_prefix}.{linear_idx}.bias"
-
+
if weight_key in usa_state_dict:
- weight = usa_state_dict[weight_key].t() # Transpose to (in_features, out_features)
+ weight = usa_state_dict[
+ weight_key
+ ].t() # Transpose to (in_features, out_features)
key_matrices_per_layer[i].append(weight)
-
+
if bias_key in usa_state_dict:
key_biases_per_layer[i].append(usa_state_dict[bias_key])
else:
- key_biases_per_layer[i].append(torch.zeros(usa_state_dict[weight_key].shape[0]))
-
+ key_biases_per_layer[i].append(
+ torch.zeros(usa_state_dict[weight_key].shape[0])
+ )
+
# Stack all heads for each layer
for i in range(3):
if query_matrices_per_layer[i]:
- layer_weights["query_matrix"].append(torch.stack(query_matrices_per_layer[i]))
- layer_weights["query_bias"].append(torch.stack(query_biases_per_layer[i]))
- layer_weights["key_matrix"].append(torch.stack(key_matrices_per_layer[i]))
+ layer_weights["query_matrix"].append(
+ torch.stack(query_matrices_per_layer[i])
+ )
+ layer_weights["query_bias"].append(
+ torch.stack(query_biases_per_layer[i])
+ )
+ layer_weights["key_matrix"].append(
+ torch.stack(key_matrices_per_layer[i])
+ )
layer_weights["key_bias"].append(torch.stack(key_biases_per_layer[i]))
-
+
hat_weights[layer_idx] = layer_weights
-
+
print(f"โ
Converted weights for {num_layers} layers, {num_heads} heads")
return hat_weights
-def create_dummy_weights(num_layers: int = 32, num_heads: int = 32) -> Dict[int, Dict[str, List[torch.Tensor]]]:
+def create_dummy_weights(
+ num_layers: int = 32, num_heads: int = 32
+) -> Dict[int, Dict[str, List[torch.Tensor]]]:
"""Create dummy weights for demonstration purposes."""
-
+
hat_weights = {}
for layer_idx in range(num_layers):
hat_weights[layer_idx] = {
"query_matrix": [
torch.randn(num_heads, 128, 128), # First linear layer
torch.randn(num_heads, 128, 128), # Second linear layer
- torch.randn(num_heads, 128, 32), # Third linear layer
+ torch.randn(num_heads, 128, 32), # Third linear layer
],
"query_bias": [
torch.randn(num_heads, 128),
@@ -124,20 +145,24 @@ def create_dummy_weights(num_layers: int = 32, num_heads: int = 32) -> Dict[int,
def main():
"""Main function demonstrating HashAttention."""
-
+
print("๐ HashAttention Example")
print("=" * 50)
-
+
# Setup
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")
-
+
# Load weights
- usa_checkpoint_path = "/home/apd10/HashAttention-1.0/artifacts/llama3.1-8b-patch.64K.v1.pt"
-
+ usa_checkpoint_path = (
+ "/home/apd10/HashAttention-1.0/artifacts/llama3.1-8b-patch.64K.v1.pt"
+ )
+
if os.path.exists(usa_checkpoint_path):
try:
- hat_weights = convert_usa_weights_to_hash_attention(usa_checkpoint_path, device=device)
+ hat_weights = convert_usa_weights_to_hash_attention(
+ usa_checkpoint_path, device=device
+ )
print("โ
Successfully loaded USA weights")
except Exception as e:
print(f"โ Error loading USA weights: {e}")
@@ -147,7 +172,7 @@ def main():
print(f"โ USA checkpoint not found at {usa_checkpoint_path}")
print("Creating dummy weights for demonstration...")
hat_weights = create_dummy_weights()
-
+
# Configure HashAttention
local_config = LocalMaskerConfig(window_size=4)
sink_config = SinkMaskerConfig(sink_size=4)
@@ -157,22 +182,24 @@ def main():
hat_mlp_layers=3,
hat_mlp_hidden_size=128,
hat_mlp_activation="silu",
- hat_weights=hat_weights
+ hat_weights=hat_weights,
)
-
+
research_config = ResearchAttentionConfig(
masker_configs=[local_config, sink_config, hash_config]
)
-
+
print("โ
HashAttention config: Local(4) + Sink(4) + Hash(12 bits, 32 heavy)")
-
+
# Load model
model_name = "meta-llama/Llama-3.1-8B-Instruct"
-
- adapter = ModelAdapterHF(model_name=model_name,
- sparse_attention_config=research_config,
- model_kwargs={"torch_dtype": torch.bfloat16, "device_map": "cuda"},
- device="cuda")
+
+ adapter = ModelAdapterHF(
+ model_name=model_name,
+ sparse_attention_config=research_config,
+ model_kwargs={"torch_dtype": torch.bfloat16, "device_map": "cuda"},
+ device="cuda",
+ )
# Prepare test input
test_context = """
The concept of attention mechanisms has revolutionized natural language processing and machine learning.
@@ -182,8 +209,8 @@ def main():
This approach maintains performance while reducing computational costs for long sequences.
"""
test_questions = [
- "Summarize the above in a single title with less than 10 words. Given only the title.",
- "What are other attention mechanisms that are used in the field of LLMs?"
+ "Summarize the above in a single title with less than 10 words. Given only the title.",
+ "What are other attention mechanisms that are used in the field of LLMs?",
]
request = Request(
@@ -192,12 +219,12 @@ def main():
)
print("Running Hash Attention on Question")
- response=adapter.process_request(request, generation_kwargs={"max_new_tokens": 50}, request_kwargs={})
- response_text=response.responses
+ response = adapter.process_request(
+ request, generation_kwargs={"max_new_tokens": 50}, request_kwargs={}
+ )
+ response_text = response.responses
print("Hash Attention Response: ", response_text)
-
-
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/tutorials/py_examples/04_simple_benchmark_example.py b/tutorials/py_examples/04_simple_benchmark_example.py
index 30b55c07..c0a2d28c 100644
--- a/tutorials/py_examples/04_simple_benchmark_example.py
+++ b/tutorials/py_examples/04_simple_benchmark_example.py
@@ -22,12 +22,16 @@
# Ensure we're in the correct directory and add to Python path
import sys
-os.chdir('/data/apdesai/code/sparse-attention-hub')
-sys.path.insert(0, '/data/apdesai/code/sparse-attention-hub')
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+os.chdir("/data/apdesai/code/sparse-attention-hub")
+sys.path.insert(0, "/data/apdesai/code/sparse-attention-hub")
+
+from sparse_attention_hub.sparse_attention.research_attention import (
+ ResearchAttentionConfig,
+)
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
- LocalMaskerConfig, SinkMaskerConfig
+ LocalMaskerConfig,
+ SinkMaskerConfig,
)
from sparse_attention_hub.adapters.huggingface import ModelAdapterHF
from benchmark import MockBenchmark
@@ -35,85 +39,88 @@
def main():
"""Run a simple benchmark comparison between dense and sparse attention."""
-
+
print("๐ฏ Simple Benchmark Example")
print("=" * 40)
-
+
# Configuration
model_name = "microsoft/Phi-4-mini-instruct" # Small model for quick testing
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
print(f"Model: {model_name}")
print(f"Device: {device}")
-
+
# Create StreamingLLM configuration: Sink (4 tokens) + Local (32 tokens)
sink_config = SinkMaskerConfig(sink_size=4)
local_config = LocalMaskerConfig(window_size=32)
- streaming_config = ResearchAttentionConfig(masker_configs=[sink_config, local_config])
-
+ streaming_config = ResearchAttentionConfig(
+ masker_configs=[sink_config, local_config]
+ )
+
print("StreamingLLM Config: Sink(4) + Local(32)")
-
+
# Common model arguments
model_kwargs = {
"model_kwargs": {"torch_dtype": torch.bfloat16},
- "device": str(device)
+ "device": str(device),
}
-
+
# Initialize adapters
print("\n๐ง Loading models...")
-
+
# Sparse adapter (StreamingLLM)
print(" โ Loading sparse attention model...")
sparse_adapter = ModelAdapterHF(
- model_name=model_name,
- sparse_attention_config=streaming_config,
- **model_kwargs
+ model_name=model_name, sparse_attention_config=streaming_config, **model_kwargs
)
-
+
# Create mock benchmark - fast and simple
mock_benchmark = MockBenchmark()
-
+
print(f"\n๐ Running MockBenchmark:")
print(f" - 5 simple reading comprehension samples")
print(f" - 3 different contexts (science, history, geography)")
print(f" - Context sharing to test grouping efficiency")
-
+
# Create result directories
result_dir = Path("./simple_benchmark_results")
result_dir.mkdir(exist_ok=True)
-
+
sparse_dir = result_dir / "sparse"
sparse_dir.mkdir(exist_ok=True)
-
+
print("\n๐งช Running Sparse Attention Benchmark...")
start_time = time.time()
-
+
# Show dataset overview
dataset_df = mock_benchmark._load_datasets()
- print(f" Processing {len(dataset_df)} samples across {dataset_df['context'].nunique()} contexts...")
-
+ print(
+ f" Processing {len(dataset_df)} samples across {dataset_df['context'].nunique()} contexts..."
+ )
+
# Run sparse benchmark
with sparse_adapter.enable_sparse_mode():
sparse_metrics = mock_benchmark.run_benchmark(
- adapter=sparse_adapter,
- result_dir=str(sparse_dir)
+ adapter=sparse_adapter, result_dir=str(sparse_dir)
)
-
+
sparse_time = time.time() - start_time
print(f" โ
Sparse completed in {sparse_time:.2f}s")
print(f" Accuracy: {sparse_metrics.get('accuracy', 'N/A')}")
- print(f" Correct predictions: {sparse_metrics.get('correct_predictions', 'N/A')}/{sparse_metrics.get('total_samples', 'N/A')}")
-
+ print(
+ f" Correct predictions: {sparse_metrics.get('correct_predictions', 'N/A')}/{sparse_metrics.get('total_samples', 'N/A')}"
+ )
+
print("\n๐ Results Summary:")
print(f" โข Benchmark: {sparse_metrics['summary']['benchmark']}")
print(f" โข Task: {sparse_metrics['summary']['task']}")
print(f" โข Unique contexts: {sparse_metrics['summary']['unique_contexts']}")
print(f" โข Evaluation method: {sparse_metrics['summary']['evaluation_method']}")
-
+
print(f"\n๐พ Results saved to: {sparse_dir}")
print(" - raw_results.csv: Detailed predictions for each sample")
print(" - metrics.json: Evaluation metrics and summary")
if __name__ == "__main__":
- main()
+ main()
diff --git a/tutorials/py_examples/04_streaming_llm_benchmark_demo.py b/tutorials/py_examples/04_streaming_llm_benchmark_demo.py
index e5dd22cb..ba249b1a 100644
--- a/tutorials/py_examples/04_streaming_llm_benchmark_demo.py
+++ b/tutorials/py_examples/04_streaming_llm_benchmark_demo.py
@@ -32,12 +32,16 @@
# Ensure we're in the correct directory and add to Python path
import sys
-os.chdir('/data/apdesai/code/sparse-attention-hub')
-sys.path.insert(0, '/data/apdesai/code/sparse-attention-hub')
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+os.chdir("/data/apdesai/code/sparse-attention-hub")
+sys.path.insert(0, "/data/apdesai/code/sparse-attention-hub")
+
+from sparse_attention_hub.sparse_attention.research_attention import (
+ ResearchAttentionConfig,
+)
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
- LocalMaskerConfig, SinkMaskerConfig
+ LocalMaskerConfig,
+ SinkMaskerConfig,
)
from sparse_attention_hub.adapters.huggingface import ModelAdapterHF
from benchmark import Benchmark, LongBench, MockBenchmark
@@ -45,18 +49,18 @@
class BenchmarkRunner:
"""Comprehensive benchmark runner for streaming attention evaluation."""
-
+
def __init__(
- self,
+ self,
model_name: str = "microsoft/Phi-4-mini-instruct",
device: str = "cuda",
sink_size: int = 4,
local_window: int = 64,
result_dir: str = "./benchmark_results",
- benchmark_type: str = "mock"
+ benchmark_type: str = "mock",
):
"""Initialize the benchmark runner.
-
+
Args:
model_name: HuggingFace model identifier
device: Device to run on ('cuda' or 'cpu')
@@ -66,73 +70,73 @@ def __init__(
benchmark_type: Type of benchmark to use ('mock' or 'longbench')
"""
self.model_name = model_name
- self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
+ self.device = torch.device(device if torch.cuda.is_available() else "cpu")
self.sink_size = sink_size
self.local_window = local_window
self.result_dir = Path(result_dir)
self.result_dir.mkdir(exist_ok=True)
self.benchmark_type = benchmark_type
-
+
print(f"๐ Initializing Benchmark Runner")
print(f" Model: {model_name}")
print(f" Device: {self.device}")
print(f" Benchmark: {benchmark_type.upper()}")
print(f" StreamingLLM Config: Sink({sink_size}) + Local({local_window})")
print(f" Results will be saved to: {self.result_dir}")
-
+
# Create StreamingLLM configuration
self.streaming_config = self._create_streaming_config()
-
+
# Initialize adapters
self.dense_adapter = None
self.sparse_adapter = None
-
+
def _create_streaming_config(self) -> ResearchAttentionConfig:
"""Create StreamingLLM attention configuration."""
sink_config = SinkMaskerConfig(sink_size=self.sink_size)
local_config = LocalMaskerConfig(window_size=self.local_window)
return ResearchAttentionConfig(masker_configs=[sink_config, local_config])
-
+
def initialize_adapters(self) -> None:
"""Initialize both dense and sparse attention adapters."""
print("๐ง Initializing model adapters...")
-
+
# Common model arguments
common_kwargs = {
"model_kwargs": {"torch_dtype": torch.bfloat16},
- "device": str(self.device)
+ "device": str(self.device),
}
-
+
# Dense adapter (no sparse attention)
print(" โ Loading dense attention adapter...")
self.dense_adapter = ModelAdapterHF(
model_name=self.model_name,
sparse_attention_config=None, # No sparse attention = dense mode
- **common_kwargs
+ **common_kwargs,
)
-
+
# Sparse adapter (StreamingLLM)
print(" โ Loading sparse attention adapter...")
self.sparse_adapter = ModelAdapterHF(
model_name=self.model_name,
sparse_attention_config=self.streaming_config,
- **common_kwargs
+ **common_kwargs,
)
-
+
print("โ
Adapters initialized successfully!")
-
+
def run_benchmark_comparison(
- self,
+ self,
benchmark_subsets: Optional[List[str]] = None,
- max_samples: Optional[int] = 10
+ max_samples: Optional[int] = 10,
) -> Dict[str, Dict]:
"""Run benchmark comparison between dense and sparse attention.
-
+
Args:
benchmark_subsets: List of benchmark subsets to run. For MockBenchmark, this parameter is ignored.
For LongBench, if None, uses a small default set.
max_samples: Maximum number of samples per subset for quick testing (ignored for MockBenchmark)
-
+
Returns:
Dictionary containing results for both dense and sparse runs
"""
@@ -147,103 +151,109 @@ def run_benchmark_comparison(
if benchmark_subsets is None:
# Use a small subset for demonstration
benchmark_subsets = ["narrativeqa", "qasper", "samsum"]
-
+
print(f"๐งช Running LongBench comparison on subsets: {benchmark_subsets}")
benchmark = LongBench(subsets_to_run=benchmark_subsets)
-
+
results = {}
-
+
# Run dense attention benchmark
print("\n๐ Running Dense Attention Benchmark...")
if self.dense_adapter is None:
- raise RuntimeError("Dense adapter not initialized. Call initialize_adapters() first.")
- results['dense'] = self._run_single_benchmark(
+ raise RuntimeError(
+ "Dense adapter not initialized. Call initialize_adapters() first."
+ )
+ results["dense"] = self._run_single_benchmark(
adapter=self.dense_adapter,
benchmark=benchmark,
mode_name="dense",
- max_samples=max_samples
+ max_samples=max_samples,
)
-
+
# Run sparse attention benchmark
print("\nโก Running Sparse Attention Benchmark...")
if self.sparse_adapter is None:
- raise RuntimeError("Sparse adapter not initialized. Call initialize_adapters() first.")
- results['sparse'] = self._run_single_benchmark(
+ raise RuntimeError(
+ "Sparse adapter not initialized. Call initialize_adapters() first."
+ )
+ results["sparse"] = self._run_single_benchmark(
adapter=self.sparse_adapter,
benchmark=benchmark,
mode_name="sparse",
- max_samples=max_samples
+ max_samples=max_samples,
)
-
+
return results
-
+
def _run_single_benchmark(
- self,
- adapter: ModelAdapterHF,
- benchmark: Benchmark,
+ self,
+ adapter: ModelAdapterHF,
+ benchmark: Benchmark,
mode_name: str,
- max_samples: Optional[int] = None
+ max_samples: Optional[int] = None,
) -> Dict:
"""Run benchmark with a single adapter and track performance metrics.
-
+
Args:
adapter: The model adapter to use
benchmark: The benchmark instance
mode_name: Name for this benchmark run ('dense' or 'sparse')
max_samples: Maximum samples to process (for quick testing)
-
+
Returns:
Dictionary containing benchmark results and performance metrics
"""
# Start memory tracking
tracemalloc.start()
start_time = time.time()
-
+
# Create timestamped result directory
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
mode_result_dir = self.result_dir / f"{mode_name}_{timestamp}"
mode_result_dir.mkdir(exist_ok=True)
-
+
try:
# Load and optionally limit dataset size
print(f" ๐ Loading {mode_name} benchmark data...")
dataset_df = benchmark._load_datasets()
-
+
if max_samples:
- print(f" โ๏ธ Limiting to {max_samples} samples per task for quick testing")
- dataset_df = dataset_df.groupby('task').head(max_samples).reset_index(drop=True)
-
+ print(
+ f" โ๏ธ Limiting to {max_samples} samples per task for quick testing"
+ )
+ dataset_df = (
+ dataset_df.groupby("task").head(max_samples).reset_index(drop=True)
+ )
+
print(f" ๐ Processing {len(dataset_df)} samples...")
-
+
# Run the benchmark
if mode_name == "sparse":
# Use sparse attention mode
with adapter.enable_sparse_mode():
metrics = benchmark.run_benchmark(
- adapter=adapter,
- result_dir=str(mode_result_dir)
+ adapter=adapter, result_dir=str(mode_result_dir)
)
else:
- # Use dense attention mode
+ # Use dense attention mode
metrics = benchmark.run_benchmark(
- adapter=adapter,
- result_dir=str(mode_result_dir)
+ adapter=adapter, result_dir=str(mode_result_dir)
)
-
+
# Track performance metrics
end_time = time.time()
current, peak = tracemalloc.get_traced_memory()
tracemalloc.stop()
-
+
# Compile results
performance_metrics = {
"execution_time_seconds": end_time - start_time,
"peak_memory_mb": peak / 1024 / 1024,
"current_memory_mb": current / 1024 / 1024,
"samples_processed": len(dataset_df),
- "avg_time_per_sample": (end_time - start_time) / len(dataset_df)
+ "avg_time_per_sample": (end_time - start_time) / len(dataset_df),
}
-
+
result = {
"benchmark_metrics": metrics,
"performance_metrics": performance_metrics,
@@ -251,129 +261,157 @@ def _run_single_benchmark(
"model_config": {
"model_name": self.model_name,
"mode": mode_name,
- "streaming_config": {
- "sink_size": self.sink_size,
- "local_window": self.local_window
- } if mode_name == "sparse" else None
- }
+ "streaming_config": (
+ {"sink_size": self.sink_size, "local_window": self.local_window}
+ if mode_name == "sparse"
+ else None
+ ),
+ },
}
-
+
# Save detailed results
with open(mode_result_dir / "detailed_results.json", "w") as f:
json.dump(result, f, indent=2, default=str)
-
+
print(f" โ
{mode_name.title()} benchmark completed!")
print(f" Time: {performance_metrics['execution_time_seconds']:.2f}s")
print(f" Peak Memory: {performance_metrics['peak_memory_mb']:.1f}MB")
print(f" Avg Score: {metrics.get('average_score', 'N/A')}")
-
+
return result
-
+
except Exception as e:
print(f" โ Error in {mode_name} benchmark: {str(e)}")
tracemalloc.stop()
raise e
-
+
def analyze_results(self, results: Dict[str, Dict]) -> None:
"""Analyze and visualize benchmark comparison results.
-
+
Args:
results: Results dictionary from run_benchmark_comparison
"""
print("\n๐ Analyzing Results...")
-
+
# Extract metrics for comparison
- dense_metrics = results['dense']['performance_metrics']
- sparse_metrics = results['sparse']['performance_metrics']
-
- dense_benchmark = results['dense']['benchmark_metrics']
- sparse_benchmark = results['sparse']['benchmark_metrics']
-
+ dense_metrics = results["dense"]["performance_metrics"]
+ sparse_metrics = results["sparse"]["performance_metrics"]
+
+ dense_benchmark = results["dense"]["benchmark_metrics"]
+ sparse_benchmark = results["sparse"]["benchmark_metrics"]
+
# Performance comparison
print("\n๐ Performance Comparison:")
print("โ" * 50)
print(f"{'Metric':<25} {'Dense':<15} {'Sparse':<15} {'Speedup':<10}")
print("โ" * 50)
-
- time_speedup = dense_metrics['execution_time_seconds'] / sparse_metrics['execution_time_seconds']
- memory_reduction = (dense_metrics['peak_memory_mb'] - sparse_metrics['peak_memory_mb']) / dense_metrics['peak_memory_mb'] * 100
-
- print(f"{'Execution Time (s)':<25} {dense_metrics['execution_time_seconds']:<15.2f} {sparse_metrics['execution_time_seconds']:<15.2f} {time_speedup:<10.2f}x")
- print(f"{'Peak Memory (MB)':<25} {dense_metrics['peak_memory_mb']:<15.1f} {sparse_metrics['peak_memory_mb']:<15.1f} {memory_reduction:<10.1f}%")
- print(f"{'Avg Time/Sample (s)':<25} {dense_metrics['avg_time_per_sample']:<15.3f} {sparse_metrics['avg_time_per_sample']:<15.3f}")
-
+
+ time_speedup = (
+ dense_metrics["execution_time_seconds"]
+ / sparse_metrics["execution_time_seconds"]
+ )
+ memory_reduction = (
+ (dense_metrics["peak_memory_mb"] - sparse_metrics["peak_memory_mb"])
+ / dense_metrics["peak_memory_mb"]
+ * 100
+ )
+
+ print(
+ f"{'Execution Time (s)':<25} {dense_metrics['execution_time_seconds']:<15.2f} {sparse_metrics['execution_time_seconds']:<15.2f} {time_speedup:<10.2f}x"
+ )
+ print(
+ f"{'Peak Memory (MB)':<25} {dense_metrics['peak_memory_mb']:<15.1f} {sparse_metrics['peak_memory_mb']:<15.1f} {memory_reduction:<10.1f}%"
+ )
+ print(
+ f"{'Avg Time/Sample (s)':<25} {dense_metrics['avg_time_per_sample']:<15.3f} {sparse_metrics['avg_time_per_sample']:<15.3f}"
+ )
+
# Accuracy comparison
print("\n๐ฏ Accuracy Comparison:")
print("โ" * 40)
- dense_score = dense_benchmark.get('average_score', 0)
- sparse_score = sparse_benchmark.get('average_score', 0)
- accuracy_retention = (sparse_score / dense_score * 100) if dense_score > 0 else 0
-
+ dense_score = dense_benchmark.get("average_score", 0)
+ sparse_score = sparse_benchmark.get("average_score", 0)
+ accuracy_retention = (
+ (sparse_score / dense_score * 100) if dense_score > 0 else 0
+ )
+
print(f"Dense Attention Score: {dense_score:.3f}")
print(f"Sparse Attention Score: {sparse_score:.3f}")
print(f"Accuracy Retention: {accuracy_retention:.1f}%")
-
+
# Create visualization
self._create_visualization(results)
-
+
# Summary
print(f"\n๐ Summary:")
print(f" StreamingLLM achieves {time_speedup:.2f}x speedup")
print(f" Reduces memory usage by {memory_reduction:.1f}%")
print(f" Retains {accuracy_retention:.1f}% of original accuracy")
-
+
def _create_visualization(self, results: Dict[str, Dict]) -> None:
"""Create visualization comparing dense vs sparse performance."""
try:
import matplotlib.pyplot as plt
-
+
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 10))
- fig.suptitle('Dense vs Sparse Attention Benchmark Comparison', fontsize=16)
-
+ fig.suptitle("Dense vs Sparse Attention Benchmark Comparison", fontsize=16)
+
# Extract data
- dense_perf = results['dense']['performance_metrics']
- sparse_perf = results['sparse']['performance_metrics']
- dense_bench = results['dense']['benchmark_metrics']
- sparse_bench = results['sparse']['benchmark_metrics']
-
+ dense_perf = results["dense"]["performance_metrics"]
+ sparse_perf = results["sparse"]["performance_metrics"]
+ dense_bench = results["dense"]["benchmark_metrics"]
+ sparse_bench = results["sparse"]["benchmark_metrics"]
+
# Execution time comparison
- times = [dense_perf['execution_time_seconds'], sparse_perf['execution_time_seconds']]
- ax1.bar(['Dense', 'Sparse'], times, color=['#ff7f0e', '#2ca02c'])
- ax1.set_ylabel('Execution Time (seconds)')
- ax1.set_title('Execution Time Comparison')
-
+ times = [
+ dense_perf["execution_time_seconds"],
+ sparse_perf["execution_time_seconds"],
+ ]
+ ax1.bar(["Dense", "Sparse"], times, color=["#ff7f0e", "#2ca02c"])
+ ax1.set_ylabel("Execution Time (seconds)")
+ ax1.set_title("Execution Time Comparison")
+
# Memory usage comparison
- memories = [dense_perf['peak_memory_mb'], sparse_perf['peak_memory_mb']]
- ax2.bar(['Dense', 'Sparse'], memories, color=['#ff7f0e', '#2ca02c'])
- ax2.set_ylabel('Peak Memory (MB)')
- ax2.set_title('Memory Usage Comparison')
-
+ memories = [dense_perf["peak_memory_mb"], sparse_perf["peak_memory_mb"]]
+ ax2.bar(["Dense", "Sparse"], memories, color=["#ff7f0e", "#2ca02c"])
+ ax2.set_ylabel("Peak Memory (MB)")
+ ax2.set_title("Memory Usage Comparison")
+
# Accuracy comparison
- scores = [dense_bench.get('average_score', 0), sparse_bench.get('average_score', 0)]
- ax3.bar(['Dense', 'Sparse'], scores, color=['#ff7f0e', '#2ca02c'])
- ax3.set_ylabel('Average Score')
- ax3.set_title('Accuracy Comparison')
+ scores = [
+ dense_bench.get("average_score", 0),
+ sparse_bench.get("average_score", 0),
+ ]
+ ax3.bar(["Dense", "Sparse"], scores, color=["#ff7f0e", "#2ca02c"])
+ ax3.set_ylabel("Average Score")
+ ax3.set_title("Accuracy Comparison")
ax3.set_ylim(0, 1)
-
+
# Per-sample time comparison
- per_sample_times = [dense_perf['avg_time_per_sample'], sparse_perf['avg_time_per_sample']]
- ax4.bar(['Dense', 'Sparse'], per_sample_times, color=['#ff7f0e', '#2ca02c'])
- ax4.set_ylabel('Time per Sample (seconds)')
- ax4.set_title('Per-Sample Processing Time')
-
+ per_sample_times = [
+ dense_perf["avg_time_per_sample"],
+ sparse_perf["avg_time_per_sample"],
+ ]
+ ax4.bar(["Dense", "Sparse"], per_sample_times, color=["#ff7f0e", "#2ca02c"])
+ ax4.set_ylabel("Time per Sample (seconds)")
+ ax4.set_title("Per-Sample Processing Time")
+
plt.tight_layout()
-
+
# Save plot
- plot_path = self.result_dir / f"benchmark_comparison_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
- plt.savefig(plot_path, dpi=300, bbox_inches='tight')
+ plot_path = (
+ self.result_dir
+ / f"benchmark_comparison_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
+ )
+ plt.savefig(plot_path, dpi=300, bbox_inches="tight")
print(f"๐ Visualization saved to: {plot_path}")
-
+
# Also try to show if in interactive environment
try:
plt.show()
except:
pass
-
+
except ImportError:
print("๐ Matplotlib not available for visualization")
except Exception as e:
@@ -383,31 +421,47 @@ def _create_visualization(self, results: Dict[str, Dict]) -> None:
def main():
"""Main function to run the streaming LLM benchmark demo."""
parser = argparse.ArgumentParser(description="Streaming LLM Benchmark Demo")
- parser.add_argument("--model", default="microsoft/Phi-4-mini-instruct",
- help="HuggingFace model name")
- parser.add_argument("--benchmark", default="mock", choices=["mock", "longbench"],
- help="Benchmark type: 'mock' for quick testing (5 samples) or 'longbench' for comprehensive evaluation")
- parser.add_argument("--subsets", default="narrativeqa,qasper,samsum",
- help="Comma-separated list of LongBench subsets (ignored for MockBenchmark)")
- parser.add_argument("--max-samples", type=int, default=5,
- help="Maximum samples per subset for quick testing (ignored for MockBenchmark)")
- parser.add_argument("--sink-size", type=int, default=4,
- help="Number of sink tokens")
- parser.add_argument("--local-window", type=int, default=64,
- help="Local attention window size")
- parser.add_argument("--device", default="cuda",
- help="Device to use (cuda/cpu)")
- parser.add_argument("--result-dir", default="./benchmark_results",
- help="Directory to save results")
-
+ parser.add_argument(
+ "--model",
+ default="microsoft/Phi-4-mini-instruct",
+ help="HuggingFace model name",
+ )
+ parser.add_argument(
+ "--benchmark",
+ default="mock",
+ choices=["mock", "longbench"],
+ help="Benchmark type: 'mock' for quick testing (5 samples) or 'longbench' for comprehensive evaluation",
+ )
+ parser.add_argument(
+ "--subsets",
+ default="narrativeqa,qasper,samsum",
+ help="Comma-separated list of LongBench subsets (ignored for MockBenchmark)",
+ )
+ parser.add_argument(
+ "--max-samples",
+ type=int,
+ default=5,
+ help="Maximum samples per subset for quick testing (ignored for MockBenchmark)",
+ )
+ parser.add_argument(
+ "--sink-size", type=int, default=4, help="Number of sink tokens"
+ )
+ parser.add_argument(
+ "--local-window", type=int, default=64, help="Local attention window size"
+ )
+ parser.add_argument("--device", default="cuda", help="Device to use (cuda/cpu)")
+ parser.add_argument(
+ "--result-dir", default="./benchmark_results", help="Directory to save results"
+ )
+
args = parser.parse_args()
-
+
print("๐ฏ Streaming LLM Benchmark Demo")
print("=" * 50)
-
+
# Parse subsets
subsets = [s.strip() for s in args.subsets.split(",")]
-
+
# Initialize benchmark runner
runner = BenchmarkRunner(
model_name=args.model,
@@ -415,24 +469,23 @@ def main():
sink_size=args.sink_size,
local_window=args.local_window,
result_dir=args.result_dir,
- benchmark_type=args.benchmark
+ benchmark_type=args.benchmark,
)
-
+
# Initialize adapters
runner.initialize_adapters()
-
+
# Run benchmark comparison
results = runner.run_benchmark_comparison(
- benchmark_subsets=subsets,
- max_samples=args.max_samples
+ benchmark_subsets=subsets, max_samples=args.max_samples
)
-
+
# Analyze results
runner.analyze_results(results)
-
+
print("\nโ
Benchmark demo completed successfully!")
print(f"๐ Results saved to: {runner.result_dir}")
if __name__ == "__main__":
- main()
+ main()
diff --git a/tutorials/py_examples/05_local_sink_oracle_adaptive_demo.py b/tutorials/py_examples/05_local_sink_oracle_adaptive_demo.py
index ba2954ee..f666e62a 100644
--- a/tutorials/py_examples/05_local_sink_oracle_adaptive_demo.py
+++ b/tutorials/py_examples/05_local_sink_oracle_adaptive_demo.py
@@ -23,15 +23,20 @@
# Ensure we're in the correct directory and add to Python path
import sys
-os.chdir('/data/apdesai/code/sparse-attention-hub')
-sys.path.insert(0, '/data/apdesai/code/sparse-attention-hub')
-from sparse_attention_hub.sparse_attention.research_attention import ResearchAttentionConfig
+os.chdir("/data/apdesai/code/sparse-attention-hub")
+sys.path.insert(0, "/data/apdesai/code/sparse-attention-hub")
+
+from sparse_attention_hub.sparse_attention.research_attention import (
+ ResearchAttentionConfig,
+)
from sparse_attention_hub.sparse_attention.research_attention.maskers.fixed.implementations import (
- LocalMaskerConfig, SinkMaskerConfig, OracleTopKConfig
+ LocalMaskerConfig,
+ SinkMaskerConfig,
+ OracleTopKConfig,
)
from sparse_attention_hub.sparse_attention.research_attention.maskers.sampling.implementations import (
- AdaptiveSamplingMaskerConfig
+ AdaptiveSamplingMaskerConfig,
)
from sparse_attention_hub.adapters.huggingface import ModelAdapterHF
from sparse_attention_hub.adapters import Request
@@ -39,79 +44,80 @@
def main():
"""Run a demo with combined maskers: Local + Sink + Oracle-TopK + Adaptive Sampling."""
-
+
print("๐ฏ Local + Sink + Oracle-TopK + Adaptive Sampling Demo")
print("=" * 60)
-
+
# Configuration
model_name = "microsoft/Phi-4-mini-instruct" # Small model for quick testing
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
print(f"Model: {model_name}")
print(f"Device: {device}")
-
+
# Create combined masker configuration
print("\n๐ง Creating combined masker configuration...")
-
+
# 1. Local masker: 4-token window for local attention
local_config = LocalMaskerConfig(window_size=4)
print(" โ LocalMasker: window_size=4")
-
+
# 2. Sink masker: 4 sink tokens for global information
sink_config = SinkMaskerConfig(sink_size=4)
print(" โ SinkMasker: sink_size=4")
-
+
# 3. Oracle-TopK masker: 4 tokens based on oracle attention scores
oracle_config = OracleTopKConfig(heavy_size=4)
print(" โ OracleTopKMasker: heavy_size=4")
-
+
# 4. Adaptive sampling masker: 10% base rate with statistical guarantees
adaptive_config = AdaptiveSamplingMaskerConfig(
base_rate_sampling=0.1, # 10% base sampling rate
- epsilon=0.1, # Error bound
- delta=0.05, # Confidence bound
- init_offset=4, # Start from beginning
- local_offset=4 # End at sequence end
+ epsilon=0.1, # Error bound
+ delta=0.05, # Confidence bound
+ init_offset=4, # Start from beginning
+ local_offset=4, # End at sequence end
)
print(" โ AdaptiveSamplingMasker: base_rate=0.1, epsilon=0.1, delta=0.05")
-
+
# Combine all maskers in order of application
combined_config = ResearchAttentionConfig(
- masker_configs=[local_config,
- sink_config,
- oracle_config,
- adaptive_config,
+ masker_configs=[
+ local_config,
+ sink_config,
+ oracle_config,
+ adaptive_config,
]
)
-
+
print("\n๐ Combined Configuration:")
print(" โข Local(4) + Sink(4) + Oracle-TopK(4) + Adaptive(0.1)")
print(" โข Total maskers: 4")
print(" โข Expected sparsity: High (multiple sparse patterns combined)")
-
+
# Common model arguments
model_kwargs = {
"model_kwargs": {"torch_dtype": torch.bfloat16},
- "device": str(device)
+ "device": str(device),
}
-
+
# Initialize adapter
print("\n๐ง Loading model with combined maskers...")
-
+
try:
adapter = ModelAdapterHF(
model_name=model_name,
sparse_attention_config=combined_config,
- **model_kwargs
+ **model_kwargs,
)
print(" โ
Successfully loaded model with combined maskers")
except Exception as e:
print(f" โ Error loading model: {e}")
return
-
+
# Prepare test input
print("\n๐ Preparing test input...")
-
+
test_context = """
The sparse attention mechanism combines multiple attention patterns to achieve
both computational efficiency and performance. This approach uses:
@@ -124,50 +130,55 @@ def main():
This combination allows the model to maintain high performance while significantly
reducing computational complexity for long sequences.
"""
-
+
test_questions = [
"What are the four attention patterns used in this sparse attention mechanism?",
"How does adaptive sampling contribute to the overall efficiency?",
- "Explain the difference between local and sink attention patterns."
+ "Explain the difference between local and sink attention patterns.",
]
-
+
request = Request(
context=test_context,
questions=test_questions,
)
-
+
print(f" โ Context length: {len(test_context.split())} words")
print(f" โ Number of questions: {len(test_questions)}")
-
+
# Run inference
print("\n๐งช Running inference with combined maskers...")
start_time = time.time()
-
+
try:
- response = adapter.process_request(request, generation_kwargs={"max_new_tokens": 100}, request_kwargs={"max_context": 1024})
+ response = adapter.process_request(
+ request,
+ generation_kwargs={"max_new_tokens": 100},
+ request_kwargs={"max_context": 1024},
+ )
response_text = response.responses
-
+
inference_time = time.time() - start_time
print(f" โ
Inference completed in {inference_time:.2f}s")
-
+
# Display results
print("\n๐ Results:")
print("-" * 40)
-
+
for i, (question, answer) in enumerate(zip(test_questions, response_text), 1):
print(f"\nQ{i}: {question}")
print(f"A{i}: {answer}")
-
+
except Exception as e:
print(f" โ Error during inference: {e}")
import traceback
+
traceback.print_exc()
return
-
+
# Performance analysis
print("\n๐ Performance Analysis:")
print("-" * 40)
-
+
# Calculate expected sparsity
# This is a rough estimate based on the masker configurations
print(" โข Local attention: ~25% sparsity (4-token window)")
@@ -175,18 +186,20 @@ def main():
print(" โข Oracle-TopK: ~25% sparsity (4 top tokens)")
print(" โข Adaptive sampling: ~10% base rate + adaptive budget")
print(" โข Combined effect: High sparsity with maintained performance")
-
+
print("\n๐ก Key Benefits:")
print(" โข Local attention captures immediate context")
print(" โข Sink attention preserves global information")
print(" โข Oracle attention selects most relevant tokens")
print(" โข Adaptive sampling provides statistical guarantees")
print(" โข Combined approach balances efficiency and performance")
-
+
print(f"\nโ
Demo completed successfully!")
print(" The combined masker approach demonstrates how different attention")
- print(" patterns can be layered to create sophisticated sparse attention mechanisms.")
+ print(
+ " patterns can be layered to create sophisticated sparse attention mechanisms."
+ )
if __name__ == "__main__":
- main()
+ main()
+