wehub-resource-sync
76d991c447
Auto Update PR / update-prs (push) Has been cancelled
CI / format-check (push) Has been cancelled
CI / test (3.10) (push) Has been cancelled
CI / test (3.11) (push) Has been cancelled
CI / test (3.12) (push) Has been cancelled
CI / live-api-tests (push) Has been cancelled
CI / plugin-integration-test (push) Has been cancelled
CI / ollama-integration-test (push) Has been cancelled
CI / test-fork-pr (push) Has been cancelled
439 lines
13 KiB
Python
439 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright 2025 Google LLC.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""LangExtract benchmark suite for performance and quality testing.
|
|
|
|
Measures tokenization speed and extraction quality across multiple languages
|
|
and text types. Automatically downloads test texts from Project Gutenberg
|
|
and generates comparative visualizations.
|
|
|
|
Usage:
|
|
# Run diverse text type benchmark (default)
|
|
python benchmarks/benchmark.py
|
|
|
|
# Test with specific model
|
|
python benchmarks/benchmark.py --model gemini-2.5-flash
|
|
python benchmarks/benchmark.py --model gemma2:2b # Local model via Ollama
|
|
|
|
# Generate comparison plots from existing results
|
|
python benchmarks/benchmark.py --compare
|
|
|
|
Requirements:
|
|
- Set GEMINI_API_KEY for cloud models
|
|
- Install Ollama for local model testing
|
|
- Results saved to benchmark_results/
|
|
"""
|
|
|
|
import argparse
|
|
from datetime import datetime
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import time
|
|
from typing import Any
|
|
import urllib.error
|
|
|
|
import dotenv
|
|
|
|
from benchmarks import config
|
|
from benchmarks import plotting
|
|
from benchmarks import utils
|
|
import langextract
|
|
from langextract import core
|
|
from langextract import data
|
|
from langextract import visualize
|
|
import langextract.io as lio
|
|
|
|
# Load API key from environment
|
|
dotenv.load_dotenv(override=True)
|
|
GEMINI_API_KEY = os.environ.get(
|
|
"GEMINI_API_KEY", os.environ.get("LANGEXTRACT_API_KEY")
|
|
)
|
|
|
|
|
|
class BenchmarkRunner:
|
|
"""Orchestrates benchmark execution and result collection."""
|
|
|
|
def __init__(self):
|
|
"""Initialize runner with timestamp and git metadata."""
|
|
self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
self.git_info = utils.get_git_info()
|
|
self.tokenizer = core.tokenizer.RegexTokenizer()
|
|
|
|
def set_tokenizer(self, tokenizer_type: str):
|
|
"""Set the tokenizer to use."""
|
|
if tokenizer_type.lower() == "unicode":
|
|
self.tokenizer = core.tokenizer.UnicodeTokenizer()
|
|
print("Using UnicodeTokenizer")
|
|
else:
|
|
self.tokenizer = core.tokenizer.RegexTokenizer()
|
|
print("Using RegexTokenizer (default)")
|
|
|
|
def print_header(self):
|
|
"""Print benchmark header."""
|
|
print("=" * config.DISPLAY.separator_width)
|
|
print("LANGEXTRACT BENCHMARK")
|
|
print("=" * config.DISPLAY.separator_width)
|
|
print(
|
|
f"Branch: {self.git_info['branch']} | Commit: {self.git_info['commit']}"
|
|
)
|
|
print("-" * config.DISPLAY.separator_width)
|
|
|
|
def benchmark_tokenization(self) -> list[dict[str, Any]]:
|
|
"""Measure tokenization throughput at different text sizes.
|
|
|
|
Returns:
|
|
List of dicts with words, tokens, timing, and throughput metrics.
|
|
"""
|
|
print("\nTokenization Performance")
|
|
print("-" * config.DISPLAY.subseparator_width)
|
|
|
|
results = []
|
|
|
|
for word_count in config.TOKENIZATION.default_text_sizes:
|
|
text = " ".join(["word"] * word_count)
|
|
|
|
_ = self.tokenizer.tokenize(text)
|
|
|
|
times = []
|
|
for _ in range(config.TOKENIZATION.benchmark_iterations):
|
|
start = time.perf_counter()
|
|
tokenized = self.tokenizer.tokenize(text)
|
|
elapsed = time.perf_counter() - start
|
|
times.append(elapsed)
|
|
|
|
avg_time = sum(times) / len(times)
|
|
avg_ms = avg_time * 1000
|
|
num_tokens = len(tokenized.tokens)
|
|
tokens_per_sec = num_tokens / avg_time if avg_time > 0 else 0
|
|
|
|
word_str = (
|
|
f"{word_count//1000:,}k" if word_count >= 1000 else f"{word_count:,}"
|
|
)
|
|
|
|
print(
|
|
f"{word_str:>6} words: {avg_ms:7.2f}ms "
|
|
f"({tokens_per_sec/1e6:.1f}M tokens/sec)"
|
|
)
|
|
|
|
results.append({
|
|
"words": word_count,
|
|
"tokens": num_tokens,
|
|
"avg_ms": avg_ms,
|
|
"tokens_per_sec": tokens_per_sec,
|
|
})
|
|
|
|
return results
|
|
|
|
def test_single_extraction(
|
|
self,
|
|
model_id: str = config.MODELS.default_model,
|
|
text_type: config.TextTypes = config.TextTypes.ENGLISH,
|
|
) -> dict[str, Any]:
|
|
"""Execute extraction test.
|
|
|
|
Args:
|
|
model_id: Model identifier (e.g., 'gemini-2.5-flash', 'gemma2:2b').
|
|
text_type: Language/text type to test.
|
|
|
|
Returns:
|
|
Dict with success status, timing, entity counts, and metrics.
|
|
"""
|
|
print("\nExtraction Test")
|
|
print("-" * config.DISPLAY.subseparator_width)
|
|
|
|
try:
|
|
# Get test text
|
|
test_text = utils.get_text_from_gutenberg(text_type)
|
|
test_text = utils.get_optimal_text_size(test_text, model_id)
|
|
|
|
print(f" Text: {len(test_text):,} characters ({text_type.value})")
|
|
print(f" Model: {model_id}")
|
|
|
|
# Analyze tokenization
|
|
tokenization_analysis = utils.analyze_tokenization(
|
|
test_text, self.tokenizer
|
|
)
|
|
print(
|
|
" Tokenization:"
|
|
f" {utils.format_tokenization_summary(tokenization_analysis)}"
|
|
)
|
|
|
|
# Get extraction config for text type
|
|
extraction_config = utils.get_extraction_example(text_type)
|
|
|
|
example = data.ExampleData(
|
|
text="MACBETH speaks to LADY MACBETH about Duncan.",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_text="Macbeth", extraction_class="Character"
|
|
),
|
|
data.Extraction(
|
|
extraction_text="Lady Macbeth", extraction_class="Character"
|
|
),
|
|
data.Extraction(
|
|
extraction_text="Duncan", extraction_class="Character"
|
|
),
|
|
],
|
|
)
|
|
|
|
max_retries = 5
|
|
retry_delay = 3.0
|
|
|
|
# Retry logic for transient network/API failures
|
|
for attempt in range(max_retries):
|
|
try:
|
|
start_time = time.time()
|
|
result = langextract.extract(
|
|
text_or_documents=test_text,
|
|
model_id=model_id,
|
|
api_key=GEMINI_API_KEY,
|
|
prompt_description=extraction_config["prompt"],
|
|
examples=[example],
|
|
max_workers=config.MODELS.default_max_workers,
|
|
temperature=config.MODELS.default_temperature,
|
|
extraction_passes=config.MODELS.default_extraction_passes,
|
|
tokenizer=self.tokenizer,
|
|
)
|
|
elapsed = time.time() - start_time
|
|
break
|
|
except (ConnectionError, TimeoutError):
|
|
if attempt < max_retries - 1:
|
|
print(f" Retrying in {retry_delay}s...")
|
|
time.sleep(retry_delay)
|
|
retry_delay *= 1.5
|
|
continue
|
|
raise
|
|
|
|
print(f"Extraction completed in {elapsed:.1f}s")
|
|
|
|
grounded_entities = []
|
|
ungrounded_entities = []
|
|
|
|
if result.extractions:
|
|
for extraction in result.extractions:
|
|
is_grounded = (
|
|
extraction.char_interval
|
|
and extraction.char_interval.start_pos is not None
|
|
and extraction.char_interval.end_pos is not None
|
|
)
|
|
|
|
entity_text = extraction.extraction_text
|
|
if entity_text:
|
|
if is_grounded:
|
|
grounded_entities.append(entity_text)
|
|
else:
|
|
ungrounded_entities.append(entity_text)
|
|
|
|
unique_grounded = list(set(grounded_entities))
|
|
unique_ungrounded = list(set(ungrounded_entities))
|
|
|
|
print(f"Found {len(unique_grounded)} grounded entities")
|
|
if unique_ungrounded:
|
|
print(f" ({len(unique_ungrounded)} ungrounded entities ignored)")
|
|
|
|
if unique_grounded:
|
|
sample = unique_grounded[:5]
|
|
sample_str = ", ".join(sample) + (
|
|
"..." if len(unique_grounded) > 5 else ""
|
|
)
|
|
print(f" Sample: {sample_str}")
|
|
|
|
return {
|
|
"success": True,
|
|
"model": model_id,
|
|
"text_type": text_type.value,
|
|
"time_seconds": elapsed,
|
|
"entity_count": len(unique_grounded),
|
|
"ungrounded_count": len(unique_ungrounded),
|
|
"sample_entities": unique_grounded[:10],
|
|
"tokenization": tokenization_analysis,
|
|
config.EXTRACTION_RESULT_KEY: result,
|
|
}
|
|
|
|
except (urllib.error.URLError, RuntimeError) as e:
|
|
# Handle expected text download failures.
|
|
print(f"Failed: {e}")
|
|
return {
|
|
"success": False,
|
|
"model": model_id,
|
|
"text_type": text_type.value,
|
|
"error": str(e),
|
|
}
|
|
|
|
def test_diverse_text_types(
|
|
self, models: list[str] | None = None
|
|
) -> list[dict[str, Any]]:
|
|
"""Test extraction with diverse text types."""
|
|
print("\n" + "=" * config.DISPLAY.separator_width)
|
|
print("DIVERSE TEXT TYPE MODE")
|
|
print("=" * config.DISPLAY.separator_width)
|
|
|
|
if models is None:
|
|
models = [config.MODELS.default_model]
|
|
|
|
results = []
|
|
test_count = 0
|
|
|
|
for model_id in models:
|
|
print(f"\nTesting {model_id}")
|
|
print("-" * 30)
|
|
|
|
for text_type in config.TextTypes:
|
|
print(f"\n Testing {text_type.value} text...")
|
|
result = self.test_single_extraction(model_id, text_type)
|
|
results.append(result)
|
|
|
|
if result.get("success"):
|
|
test_count += 1
|
|
if test_count % 3 == 0:
|
|
print(
|
|
" Rate limit delay"
|
|
f" ({config.MODELS.gemini_rate_limit_delay}s)..."
|
|
)
|
|
time.sleep(config.MODELS.gemini_rate_limit_delay)
|
|
|
|
print(f"\nCompleted {test_count} successful tests")
|
|
return results
|
|
|
|
def save_results(self, results: dict[str, Any]):
|
|
"""Save results and create plots."""
|
|
results["timestamp"] = self.timestamp
|
|
results["git"] = self.git_info
|
|
|
|
json_path = config.PATHS.get_result_path(self.timestamp, "").with_suffix(
|
|
".json"
|
|
)
|
|
|
|
viz_dir = json_path.parent / "visualizations" / self.timestamp
|
|
viz_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if config.RESULTS_KEY in results:
|
|
print(f"\nGenerating visualizations in: {viz_dir}")
|
|
for result in results[config.RESULTS_KEY]:
|
|
if result.get("success") and config.EXTRACTION_RESULT_KEY in result:
|
|
model_name = result["model"].replace("/", "_").replace(":", "_")
|
|
text_type = result["text_type"]
|
|
viz_name = f"{model_name}_{text_type}"
|
|
|
|
jsonl_path = viz_dir / f"{viz_name}.jsonl"
|
|
lio.save_annotated_documents(
|
|
[result[config.EXTRACTION_RESULT_KEY]],
|
|
output_name=jsonl_path.name,
|
|
output_dir=str(viz_dir),
|
|
)
|
|
|
|
html_content = visualize(str(jsonl_path))
|
|
html_path = viz_dir / f"{viz_name}.html"
|
|
with open(html_path, "w") as f:
|
|
f.write(getattr(html_content, "data", html_content))
|
|
|
|
# Remove extraction result objects before saving JSON
|
|
for result in results.get(config.RESULTS_KEY, []):
|
|
result.pop(config.EXTRACTION_RESULT_KEY, None)
|
|
|
|
with open(json_path, "w") as f:
|
|
json.dump(results, f, indent=2, default=str)
|
|
print(f"\nResults saved to: {json_path}")
|
|
|
|
plot_created = plotting.create_diverse_plots(results, json_path)
|
|
|
|
if plot_created:
|
|
print(f"Plot saved to: {json_path.with_suffix('.png')}")
|
|
else:
|
|
print(f"Warning: Failed to create plot for {json_path.name}")
|
|
|
|
def run_diverse_benchmark(self, models: list[str] | None = None):
|
|
"""Run benchmark."""
|
|
self.print_header()
|
|
|
|
tokenization_results = self.benchmark_tokenization()
|
|
diverse_results = self.test_diverse_text_types(models)
|
|
|
|
results = {
|
|
"tokenization": tokenization_results,
|
|
config.RESULTS_KEY: diverse_results,
|
|
}
|
|
|
|
self.save_results(results)
|
|
|
|
|
|
def main():
|
|
"""Main entry point."""
|
|
parser = argparse.ArgumentParser(description="LangExtract Benchmark Suite")
|
|
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default=None,
|
|
help=f"Model to use (default: {config.MODELS.default_model})",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--tokenizer",
|
|
type=str,
|
|
choices=["regex", "unicode"],
|
|
default="regex",
|
|
help="Tokenizer to use (default: regex)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--compare",
|
|
action="store_true",
|
|
help="Generate comparison plots from existing benchmark results",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Handle comparison mode
|
|
if args.compare:
|
|
results_dir = Path("benchmark_results")
|
|
json_files = sorted(results_dir.glob("benchmark_*.json"))
|
|
|
|
if len(json_files) < 2:
|
|
print(
|
|
"Need at least 2 benchmark results for comparison, found"
|
|
f" {len(json_files)}"
|
|
)
|
|
return
|
|
|
|
print(f"Found {len(json_files)} benchmark results to compare")
|
|
|
|
# Use last 10 results or all if less than 10
|
|
files_to_compare = json_files[-10:]
|
|
comparison_path = (
|
|
results_dir
|
|
/ f"comparison_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
|
|
)
|
|
|
|
plotting.create_comparison_plots(files_to_compare, comparison_path)
|
|
print(f"\nComparison plot saved to: {comparison_path}")
|
|
return
|
|
|
|
model_to_test = args.model or config.MODELS.default_model
|
|
if "gemini" in model_to_test.lower() and not GEMINI_API_KEY:
|
|
print(
|
|
f"Error: {model_to_test} requires GEMINI_API_KEY or LANGEXTRACT_API_KEY"
|
|
)
|
|
return
|
|
|
|
runner = BenchmarkRunner()
|
|
runner.set_tokenizer(args.tokenizer)
|
|
runner.run_diverse_benchmark([args.model] if args.model else None)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|