diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index ae2c1891..d2ba175c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -70,7 +70,7 @@ repos: hooks: - id: mypy-local name: Run mypy for local Python installation - entry: tools/mypy.sh 0 "local" + entry: tools/mypy.sh 1 "local" language: python types: [python] pass_filenames: false diff --git a/README.md b/README.md index dcd5565f..282b31d7 100644 --- a/README.md +++ b/README.md @@ -43,19 +43,21 @@ The current leaderboard is computed considering the accuracy and overall cost fo If you want your router on the leaderboard, please contact us via email at yifan.lu@rice.edu or jxing@rice.edu, or submit a GitHub issue. For fairness, we have withheld the ground truth answers for the full dataset. However, you can still test your router using the sub-sampled 10% dataset by following the steps below. -## Usage +## Setup -### Step 1: Install uv (if you don't have it) +### Step 1: Install uv and RouterArena ```bash curl -LsSf https://astral.sh/uv/install.sh | sh +cd RouterArena +uv sync ``` -### Step 2: Install RouterArena +### Step 2: Download Dataset +Run this command to download the dataset from the [HF dataset](https://huggingface.co/datasets/RouteWorks/RouterArena). ```bash -cd RouterArena -uv sync +uv run python ./scripts/process_datasets/prep_datasets.py ``` ### Step 3: Set Up API Keys @@ -64,24 +66,20 @@ This step is **required only if you plan to use our pipeline to make LLM inferen ```bash # Example .env file -OPENAI_API_KEY=your_openai_key_here -ANTHROPIC_API_KEY=your_anthropic_key_here -GOOGLE_API_KEY=your_google_key_here -MISTRAL_API_KEY=your_mistral_key_here -HF_TOKEN=your_huggingface_token_here -# ... add other keys as needed +OPENAI_API_KEY= +ANTHROPIC_API_KEY= +HF_TOKEN= +# ... ``` +#### Optional: See the `ModelInference` class in `RouterArena/llm_inference/model_inference.py` for the complete list of supported providers and required environment variables. You can extend that class to support additional models, or submit a GitHub issue to request support for new providers. -### Step 4: Download Dataset -Run this command to download the dataset from the [HF dataset](https://huggingface.co/datasets/RouteWorks/RouterArena). +## Usage -```bash -uv run python ./scripts/process_datasets/prep_datasets.py -``` +Follow the steps below to evaluate your router. You can start with the `sub_10` split (10% sub-sampled dataset) to test your setup and code. The `sub_10` split includes ground truth answers for local testing. Once ready, you can evaluate on the `full` dataset for official leaderboard submission. -### Step 5: Prepare Config File and Model Costs +### Step 1: Prepare Config File Create a config file in `./router_inference/config/.json`. We have created an example router for demonstration purposes: @@ -101,7 +99,7 @@ Create a config file in `./router_inference/config/.json`. We have *Note: The model name must be the same as the one used in `./universal_model_names.py` (see next step for details)* -**Important**: You also need to add cost information for each model (same model naming requirement as above) in `./model_cost/cost.json`. For each model in your config, add an entry with the pricing per million tokens: +**Important**: For each model in your config, add an entry with the pricing per million tokens in this format: ```json { @@ -109,143 +107,56 @@ Create a config file in `./router_inference/config/.json`. We have "input_token_price_per_million": 0.15, "output_token_price_per_million": 0.6 }, - "claude-3-haiku-20240307": { - "input_token_price_per_million": 0.25, - "output_token_price_per_million": 1.25 - }, - "gemini-2.0-flash-001": { - "input_token_price_per_million": 0.1, - "output_token_price_per_million": 0.4 - }, - "mistral-medium": { - "input_token_price_per_million": 2.7, - "output_token_price_per_million": 8.1 - } } ``` -### Step 6: -API providers may use different names for the same model (e.g., `gpt-4o`, `openai/gpt-4o`). We manage this via `./universal_model_names.py`: -- `universal_names`: canonical model names used in this project -- `mapping`: maps external provider names to our canonical names +### Step 2: Verify Model Names -Please make sure the model you used are listed here, or you have added it there (if you add a model, please make sure you add the API inference endpoint at `RouterArena/llm_inference/model_inference.py`). +Ensure all models in your config are listed in `./universal_model_names.py`. If you add a new model, you must also add the API inference endpoint in `RouterArena/llm_inference/model_inference.py`. -### Step 7: Generate Router's Prediction File +### Step 3: Generate Router's Prediction File -You need to create a prediction file that contains your router's model selections for each query. You can use the helper script to generate a template prediction file: +Generate a template prediction file: ```bash -uv run python ./router_inference/generate_prediction_file.py your-router 10 -``` - -This command generates a prediction file at `./router_inference/predictions/your-router.json` for the 10% split. Use `full` instead of `10` for the complete dataset. - -**Important**: The generated file uses a **placeholder router** that simply cycles through models in the config file sequentially. You **must replace the model choices** in the `prediction` field with your router's actual selections. The script is only meant to provide a template structure with all required fields populated. - -An example prediction file structure: - -```json -[ - { - "global index": "ArcMMLU_655", - "prompt": "Question text here...", - "prediction": "gpt-4o-mini", // Auto generated by the generate_prediction_file.py - "generated_result": null, // Will be filled after LLM inference - "cost": null, // Will be filled after evaluation - "accuracy": null // Will be filled after evaluation - } -] +uv run python ./router_inference/generate_prediction_file.py your-router sub_10 ``` -Alternatively, you can create the prediction file manually or integrate it into your router's inference pipeline. The `generated_result`, `cost`, and `accuracy` fields can be left as `null` initially—they will be populated by the LLM inference and evaluation in later steps. +Use `full` instead of `sub_10` for the complete dataset. **Important**: Replace the placeholder model choices in the `prediction` field with your router's actual selections. -### Step 8: Sanity Check for Config and Prediction Files +### Step 4: Validate Config and Prediction Files -Before proceeding with LLM inference, it's recommended to validate your config and prediction files using our validation script: +Validate your config and prediction files before proceeding: ```bash -uv run python ./router_inference/check_config_prediction_files.py your-router 10 +uv run python ./router_inference/check_config_prediction_files.py your-router sub_10 ``` -This script performs the following checks: +This script checks: (1) all model names are valid, (2) prediction file has correct size (809 for `sub_10`, 8400 for `full`), and (3) all entries have valid `global_index`, `prompt`, and `prediction` fields. -1. **Config Validation**: Verifies that all model names in your config file are valid and can be found in `ModelNameManager` -2. **Prediction File Size**: Ensures your prediction file has the correct number of entries (809 for 10% split, 8400 for full dataset) -3. **Field Validation**: Validates that each prediction entry: - - Has a `global_index` that exists in the dataset - - Has a `prompt` that exactly matches the dataset - - Has a `prediction` (model selection) that is one of the models listed in your config +## Run LLM Inference -If all checks pass, you'll see `✓ ALL CHECKS PASSED!` and can proceed to the next step. If there are errors, the script will list them so you can fix any issues before running LLM inference and evaluation. - -### Step 9: Run LLM Inference - -Once your prediction file is ready, run the LLM inference script to make API calls for each query using the selected models: +Run the inference script to make API calls for each query using the selected models: ```bash uv run python ./llm_inference/run.py your-router ``` -This script will: -1. **Load your prediction file** from `./router_inference/predictions/your-router.json` -2. **Make API calls** for each query using the model specified in the `prediction` field -3. **Use cached results** when available (if the same model has already processed the same query) -4. **Save results incrementally** back to the prediction file, updating the `generated_result` field with: - - `generated_answer`: The model's response - - `success`: Whether the API call succeeded - - `token_usage`: Input/output token counts - - `provider`: The API provider used - - `error`: Any error message (if failed) +The script loads your prediction file, makes API calls using the models specified in the `prediction` field, and saves results incrementally. It uses cached results when available and saves progress after each query, so you can safely interrupt and resume. Results are saved to `./cached_results/` for reuse across routers. -The script automatically saves progress after each query, so you can safely interrupt and resume later. Results are also saved to `./cached_results/` for reuse across different routers. +**Note**: Requires valid API keys (see Setup Step 3). The script skips entries that already have successful results. -**Note**: This step requires valid API keys (see Step 3) for the models you're using. The script will skip entries that already have successful results, making it safe to re-run. +## LLM Evaluation and Compute RouterArena Score -### Step 10: Run LLM Evaluation +**Important**: For the `sub_10` split (testing), you can run evaluation locally and get RouterArena scores. For the `full` dataset (official leaderboard), ground truth answers are not available locally. After running LLM inference on the `full` dataset, submit your prediction file via GitHub issue or contact us at yifan.lu@rice.edu or jxing@rice.edu for official evaluation. -After LLM inference is complete, evaluate the generated answers to compute accuracy and cost metrics: +For local evaluation on the `sub_10` split, run the evaluation script: ```bash uv run python ./llm_evaluation/run.py your-router sub_10 ``` -This script will: -1. **Load your prediction file** from `./router_inference/predictions/your-router.json` -2. **Determine the dataset** for each query based on its `global_index` (e.g., "AIME_112" → AIME dataset) -3. **Evaluate each generated answer** against the ground truth using dataset-specific metrics: - - Math problems (AIME, MATH, etc.) → `math_metric` - - Multiple-choice questions (MMLUPro, ArcMMLU, etc.) → `mcq_accuracy` - - Code problems (LiveCodeBench) → `code_accuracy` - - And other specialized metrics as needed -4. **Calculate inference cost** based on token usage and model pricing from `./model_cost/cost.json` -5. **Save results incrementally** to the prediction file, updating: - - `accuracy`: Evaluation score (0.0 to 1.0) - - `cost`: Inference cost in dollars - -The script uses the `sub_10` split for testing (with ground truth answers available locally). For the full dataset evaluation, use `full` instead, but note that ground truth answers are not available locally—you'll need to submit your predictions via GitHub Issue for official evaluation. - -The script automatically skips entries that are already evaluated, making it safe to re-run or resume after interruption. - -### Step 11: Compute RouterArena Score - -After evaluation is complete, compute your router's RouterArena score: - -```bash -uv run python ./router_evaluation/compute_scores.py your-router -``` - -This script calculates: -1. **Average Accuracy**: The mean accuracy across all evaluated queries -2. **Total Cost**: The sum of all inference costs -3. **Average Cost per 1K Queries**: Total cost normalized to 1000 queries -4. **RouterArena Score**: A composite score that balances accuracy and cost. It ranges from 0 to 1, with higher scores indicating better trade-offs between accuracy and cost efficiency. - -**Note**: Scores computed on the `sub_10` split are for testing purposes. To submit your router for the official leaderboard, you need to: -1. Generate predictions and run evaluation for the `full` dataset -2. Contact us at yifan.lu@rice.edu or jxing@rice.edu, or submit a GitHub issue with your results - -The leaderboard rankings are based on RouterArena scores computed on the full dataset. +The script evaluates generated answers against ground truth, calculates inference costs, and computes router-level metrics including the RouterArena score (ranging 0-1). It skips already-evaluated entries, making it safe to re-run or resume. ## Citation: If you find our project helpful, please give us a star and cite us by: diff --git a/llm_evaluation/eval_reasoning.py b/llm_evaluation/eval_reasoning.py index 73f3b312..330fd586 100644 --- a/llm_evaluation/eval_reasoning.py +++ b/llm_evaluation/eval_reasoning.py @@ -15,7 +15,7 @@ superglue_exact_match, superglue_clozetest, ) -from datasets import load_from_disk +from datasets import load_from_disk # type: ignore[import-not-found,import-untyped] # Dataset to metric mapping dataset2metric = { diff --git a/llm_evaluation/evaluate_models.py b/llm_evaluation/evaluate_models.py index a7eede3f..d40f2378 100644 --- a/llm_evaluation/evaluate_models.py +++ b/llm_evaluation/evaluate_models.py @@ -18,7 +18,7 @@ import glob from typing import Dict, List, Any, Optional import sys -from tqdm import tqdm +from tqdm import tqdm # type: ignore[import-untyped] # Add the current directory to Python path to import eval modules sys.path.append(os.path.dirname(os.path.abspath(__file__))) diff --git a/llm_evaluation/livecodebench_util.py b/llm_evaluation/livecodebench_util.py index 1a4ae9ea..009dbcd9 100644 --- a/llm_evaluation/livecodebench_util.py +++ b/llm_evaluation/livecodebench_util.py @@ -488,20 +488,44 @@ def reliability_guard(maximum_memory_bytes: Optional[int] = None): import builtins - from typing import Any + from typing import Any, cast + + builtins.exit = cast(Any, None) # type: ignore[assignment] + builtins.quit = cast(Any, None) # type: ignore[assignment] # Prepare Any-typed aliases to avoid mypy assignment errors - builtins_mod: Any = builtins os_mod: Any = os - shutil_mod: Any = shutil subprocess_mod: Any = subprocess - modules_any: Any = sys.modules os.environ["OMP_NUM_THREADS"] = "1" - # Disable selected builtins - setattr(builtins_mod, "exit", None) - setattr(builtins_mod, "quit", None) + os.kill = cast(Any, None) # type: ignore[assignment] + os.system = cast(Any, None) # type: ignore[assignment] + os.putenv = cast(Any, None) # type: ignore[assignment] + os.remove = cast(Any, None) # type: ignore[assignment] + os.removedirs = cast(Any, None) # type: ignore[assignment] + os.rmdir = cast(Any, None) # type: ignore[assignment] + os.fchdir = cast(Any, None) # type: ignore[assignment] + os.setuid = cast(Any, None) # type: ignore[assignment] + os.fork = cast(Any, None) # type: ignore[assignment] + os.forkpty = cast(Any, None) # type: ignore[assignment] + os.killpg = cast(Any, None) # type: ignore[assignment] + os.rename = cast(Any, None) # type: ignore[assignment] + os.renames = cast(Any, None) # type: ignore[assignment] + os.truncate = cast(Any, None) # type: ignore[assignment] + os.replace = cast(Any, None) # type: ignore[assignment] + os.unlink = cast(Any, None) # type: ignore[assignment] + os.fchmod = cast(Any, None) # type: ignore[assignment] + os.fchown = cast(Any, None) # type: ignore[assignment] + os.chmod = cast(Any, None) # type: ignore[assignment] + os.chown = cast(Any, None) # type: ignore[assignment] + os.chroot = cast(Any, None) # type: ignore[assignment] + os.fchdir = cast(Any, None) # type: ignore[assignment] + os.lchflags = cast(Any, None) # type: ignore[attr-defined,assignment] + os.lchmod = cast(Any, None) # type: ignore[attr-defined,assignment] + os.lchown = cast(Any, None) # type: ignore[assignment] + os.getcwd = cast(Any, None) # type: ignore[assignment] + os.chdir = cast(Any, None) # type: ignore[assignment] # Disable destructive os functions (guard where platform-specific) for name in [ @@ -537,19 +561,22 @@ def reliability_guard(maximum_memory_bytes: Optional[int] = None): except Exception: pass - # Disable dangerous shutil functions - for name in ["rmtree", "move", "chown"]: - try: - setattr(shutil_mod, name, None) - except Exception: - pass + shutil.rmtree = cast(Any, None) # type: ignore[assignment] + shutil.move = cast(Any, None) # type: ignore[assignment] + shutil.chown = cast(Any, None) # type: ignore[assignment] # Disable subprocess.Popen setattr(subprocess_mod, "Popen", None) - # Hide selected modules - for name in ["ipdb", "joblib", "resource", "psutil", "tkinter"]: - modules_any[name] = None + setattr(subprocess, "Popen", cast(Any, None)) # type: ignore[misc] + + # __builtins__["help"] = None # this line is commented out as it results into error + + sys.modules["ipdb"] = None # type: ignore[assignment] + sys.modules["joblib"] = None # type: ignore[assignment] + sys.modules["resource"] = None # type: ignore[assignment] + sys.modules["psutil"] = None # type: ignore[assignment] + sys.modules["tkinter"] = None # type: ignore[assignment] def save_original_references(): diff --git a/llm_evaluation/metric_utils.py b/llm_evaluation/metric_utils.py index 22e0784e..f54d822b 100644 --- a/llm_evaluation/metric_utils.py +++ b/llm_evaluation/metric_utils.py @@ -2,15 +2,15 @@ # SPDX-License-Identifier: Apache-2.0 import re -import regex +import regex # type: ignore[import-untyped] from math import isclose from typing import Any, Optional, List -from latex2sympy2 import latex2sympy -from sympy import N, simplify -from sympy.parsing.latex import parse_latex -from sympy.parsing.sympy_parser import parse_expr +from latex2sympy2 import latex2sympy # type: ignore[import-not-found,import-untyped] +from sympy import N, simplify # type: ignore[import-not-found,import-untyped] +from sympy.parsing.latex import parse_latex # type: ignore[import-not-found,import-untyped] +from sympy.parsing.sympy_parser import parse_expr # type: ignore[import-not-found,import-untyped] def choice_answer_clean(pred: str) -> str: diff --git a/llm_evaluation/metrics.py b/llm_evaluation/metrics.py index 014ddc76..f9dd8add 100644 --- a/llm_evaluation/metrics.py +++ b/llm_evaluation/metrics.py @@ -6,13 +6,13 @@ import json import copy -import jieba -from fuzzywuzzy import fuzz +import jieba # type: ignore[import-not-found,import-untyped] +from fuzzywuzzy import fuzz # type: ignore[import-not-found,import-untyped] import difflib from collections import Counter -from rouge import Rouge -import regex +from rouge import Rouge # type: ignore[import-not-found,import-untyped] +import regex # type: ignore[import-untyped] from metric_utils import ( choice_answer_clean, diff --git a/llm_evaluation/run.py b/llm_evaluation/run.py index 6a74abdf..13aacd3b 100644 --- a/llm_evaluation/run.py +++ b/llm_evaluation/run.py @@ -24,6 +24,7 @@ import sys import logging import datetime +import math from typing import Dict, Any, List, Optional # Add parent directory to path for imports @@ -41,6 +42,37 @@ logger = logging.getLogger(__name__) +def compute_arena_score(cost, accuracy, beta=0.1, c_max=200, c_min=0.0044): + """ + Compute the RouterArena score S_i,β for a given cost and accuracy. + + Parameters: + ----------- + cost : float + The cost c_i of the model or router (per 1000 queries). + accuracy : float + The accuracy A_i of the model or router. + beta : float, optional + Weighting factor between accuracy and cost (default = 0.1). + c_max : float, optional + Maximum cost (default = 200). + c_min : float, optional + Minimum cost (default = 0.0044). + + Returns: + -------- + float + The computed RouterArena score S_i,β. + """ + # Compute normalized cost C_i + C_i = (math.log2(c_max) - math.log2(cost)) / (math.log2(c_max) - math.log2(c_min)) + + # Compute score S_i,β + S = ((1 + beta) * accuracy * C_i) / (beta * accuracy + C_i) + + return S + + def load_predictions_file(router_name: str) -> List[Dict[str, Any]]: """ Load router predictions from JSON file. @@ -95,8 +127,8 @@ def load_ground_truth_dataset(split: str) -> Dict[str, Dict[str, Any]]: Returns: Dictionary mapping global_index to ground truth data """ - from datasets import load_from_disk - import pandas as pd + from datasets import load_from_disk # type: ignore[import-not-found,import-untyped] + import pandas as pd # type: ignore[import-untyped] if split not in ["sub_10", "full"]: raise ValueError(f"Invalid split: {split}. Must be 'sub_10' or 'full'") @@ -176,7 +208,7 @@ def get_livecodebench_ground_truth(global_index: str) -> Optional[Dict[str, Any] """ global _livecodebench_cache try: - from datasets import load_from_disk + from datasets import load_from_disk # type: ignore[import-not-found,import-untyped] # Load LiveCodeBench dataset (cache it if needed) if _livecodebench_cache is None: @@ -296,7 +328,7 @@ def evaluate_single_prediction( # which works because code_accuracy accepts dict as ground_truth score, metric_name = evaluator._evaluate_single_entry( generated_answer, - ground_truth, + ground_truth, # type: ignore[arg-type] scorer_func, dataset_name, ) @@ -415,42 +447,58 @@ def process_router_predictions( def compute_router_metrics(predictions: List[Dict[str, Any]], router_name: str) -> None: """ - Compute router-level metrics (accuracy, cost, etc.) and display them. + Compute router-level metrics (accuracy, cost, RouterArena score, etc.) and display them. Args: predictions: List of prediction dictionaries with evaluation results router_name: Name of the router """ - total_accuracy = 0.0 - total_cost = 0.0 - count = 0 - cost_count = 0 + accuracies = [] + costs = [] + valid_cost_count = 0 for prediction in predictions: accuracy = prediction.get("accuracy") if accuracy is not None: - total_accuracy += accuracy - count += 1 + accuracies.append(accuracy) cost = prediction.get("cost") if cost is not None and cost > 0: - total_cost += cost - cost_count += 1 + costs.append(cost) + valid_cost_count += 1 - avg_accuracy = total_accuracy / count if count > 0 else 0.0 - avg_cost = total_cost / cost_count if cost_count > 0 else 0.0 - avg_cost_per_1k = avg_cost * 1000 if avg_cost else 0.0 + # Compute average accuracy + avg_accuracy = sum(accuracies) / len(accuracies) if accuracies else 0.0 - logger.info("\n" + "=" * 60) - logger.info(f"Router Metrics: {router_name}") - logger.info("=" * 60) + # Compute total cost (sum of all costs) + total_cost = sum(costs) if costs else 0.0 + + # Compute average cost per 1000 queries for RouterArena score calculation + num_queries = len(predictions) + avg_cost_per_1000 = (total_cost / num_queries * 1000) if num_queries > 0 else 0.0 + + # Compute RouterArena score using average cost per 1000 queries and average accuracy + arena_score = compute_arena_score(avg_cost_per_1000, avg_accuracy) + + # Print results + logger.info("\n" + "=" * 80) + logger.info(f"Router: {router_name}") + logger.info("=" * 80) + logger.info(f"Total Queries: {num_queries}") + logger.info(f"Queries with Accuracy: {len(accuracies)}") + logger.info(f"Queries with Valid Cost: {valid_cost_count}") logger.info(f"Average Accuracy: {avg_accuracy:.4f}") - logger.info(f"Average Cost per Query: ${avg_cost:.6f}") - logger.info(f"Average Cost per 1K Queries: ${avg_cost_per_1k:.4f}") - logger.info(f"Total Queries: {len(predictions)}") - logger.info(f"Queries with Accuracy: {count}") - logger.info(f"Queries with Cost: {cost_count}") - logger.info("=" * 60 + "\n") + logger.info(f"Total Cost: ${total_cost:.6f}") + if num_queries > 0: + logger.info(f"Average Cost per Query: ${total_cost / num_queries:.6f}") + else: + logger.info("Average Cost per Query: $0.00") + logger.info(f"Average Cost per 1K Queries: ${avg_cost_per_1000:.4f}") + logger.info(f"RouterArena Score: {arena_score:.4f}") + logger.info( + "PLEASE NOTE: The sub_10 dataset is a subset of the full dataset and is used for testing purposes. It is generally easier than the full dataset." + ) + logger.info("=" * 80 + "\n") def main(): diff --git a/llm_evaluation/utils.py b/llm_evaluation/utils.py index 80614227..1d9c3a7e 100644 --- a/llm_evaluation/utils.py +++ b/llm_evaluation/utils.py @@ -1,8 +1,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the RouterArena project # SPDX-License-Identifier: Apache-2.0 -import pandas as pd -from datasets import load_from_disk +import pandas as pd # type: ignore[import-untyped] +from datasets import load_from_disk # type: ignore[import-not-found,import-untyped] def escape_format_braces(text): diff --git a/llm_inference/model_inference.py b/llm_inference/model_inference.py index 354c3574..765e3ea8 100644 --- a/llm_inference/model_inference.py +++ b/llm_inference/model_inference.py @@ -189,8 +189,8 @@ def _get_provider(self, model_name: str) -> str: def _call_xai(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call XAI API.""" - from xai_sdk import Client - from xai_sdk.chat import user, system + from xai_sdk import Client # type: ignore[import-untyped] + from xai_sdk.chat import user, system # type: ignore[import-untyped] client = Client( api_key=os.getenv("XAI_API_KEY"), @@ -217,7 +217,7 @@ def _call_xai(self, model_name: str, prompt: str) -> Dict[str, Any]: } def _call_zhipu(self, model_name: str, prompt: str) -> Dict[str, Any]: - from zhipuai import ZhipuAI + from zhipuai import ZhipuAI # type: ignore[import-untyped] client = ZhipuAI(api_key=os.getenv("ZHIPU_API_KEY")) @@ -242,7 +242,7 @@ def _call_zhipu(self, model_name: str, prompt: str) -> Dict[str, Any]: def _call_replicate(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call Replicate API.""" - import replicate + import replicate # type: ignore[import-not-found] client = replicate.Client(api_token=self.replicate_api_key) @@ -311,7 +311,7 @@ def _call_openrouter(self, model_name: str, prompt: str) -> Dict[str, Any]: def _call_openai(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call OpenAI API.""" - import openai + import openai # type: ignore[import-untyped] client = openai.OpenAI(api_key=self.openai_api_key) @@ -345,7 +345,7 @@ def _call_openai(self, model_name: str, prompt: str) -> Dict[str, Any]: def _call_together(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call Together AI API.""" - import together + import together # type: ignore[import-untyped] client = together.Together(api_key=self.together_api_key) @@ -383,7 +383,7 @@ def _call_together(self, model_name: str, prompt: str) -> Dict[str, Any]: def _call_anthropic(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call Anthropic API.""" - import anthropic + import anthropic # type: ignore[import-untyped] client = anthropic.Anthropic(api_key=self.anthropic_api_key) @@ -416,7 +416,7 @@ def _call_anthropic(self, model_name: str, prompt: str) -> Dict[str, Any]: def _call_google(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call Google AI API.""" - import google.generativeai as genai + import google.generativeai as genai # type: ignore[import-untyped] genai.configure(api_key=self.google_api_key) @@ -444,7 +444,7 @@ def _call_google(self, model_name: str, prompt: str) -> Dict[str, Any]: def _call_mistral(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call Mistral AI API.""" - from mistralai import Mistral + from mistralai import Mistral # type: ignore[import-untyped] client = Mistral(api_key=self.mistral_api_key) @@ -492,7 +492,7 @@ def _call_mistral(self, model_name: str, prompt: str) -> Dict[str, Any]: def _call_azure(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call Azure OpenAI API.""" - import openai + import openai # type: ignore[import-untyped] client = openai.AzureOpenAI( api_key=self.azure_api_key, @@ -617,8 +617,8 @@ def _call_perplexity(self, model_name: str, prompt: str) -> Dict[str, Any]: def _call_aws(self, model_name: str, prompt: str) -> Dict[str, Any]: """Call AWS Bedrock API.""" - import boto3 - from botocore.exceptions import ClientError + import boto3 # type: ignore[import-untyped] + from botocore.exceptions import ClientError # type: ignore[import-untyped] # Map model names to their inference profile ARNs model_arn_mapping = { diff --git a/llm_inference/shared_utils.py b/llm_inference/shared_utils.py index 3c74569b..4831f336 100644 --- a/llm_inference/shared_utils.py +++ b/llm_inference/shared_utils.py @@ -175,7 +175,7 @@ def build_chat(tokenizer, prompt, chat_template): messages, tokenize=False, add_generation_prompt=True ) elif "longchat" in chat_template or "vicuna" in chat_template: - from fastchat.model import get_conversation_template + from fastchat.model import get_conversation_template # type: ignore[import-not-found] conv = get_conversation_template("vicuna") conv.append_message(conv.roles[0], prompt) diff --git a/router_inference/check_config_prediction_files.py b/router_inference/check_config_prediction_files.py index 4b13cae5..5713d2c0 100644 --- a/router_inference/check_config_prediction_files.py +++ b/router_inference/check_config_prediction_files.py @@ -15,7 +15,7 @@ Usage: python router_inference/check_config_prediction_files.py - split: either "10" for 10% split or "full" for full dataset + split: either "sub_10" for 10% split or "full" for full dataset """ import argparse @@ -31,13 +31,13 @@ # Expected dataset sizes EXPECTED_SIZES = { - "10": 809, + "sub_10": 809, "full": 8400, } # Dataset file paths DATASET_PATHS = { - "10": "./dataset/router_data_10.json", + "sub_10": "./dataset/router_data_10.json", "full": "./dataset/router_data.json", } @@ -89,7 +89,7 @@ def load_dataset(split: str) -> List[Dict[str, Any]]: Load dataset file. Args: - split: Either "10" or "full" + split: Either "sub_10" or "full" Returns: List of dataset entries @@ -97,7 +97,7 @@ def load_dataset(split: str) -> List[Dict[str, Any]]: dataset_path = DATASET_PATHS.get(split) if not dataset_path: - raise ValueError(f"Invalid split: {split}. Must be '10' or 'full'") + raise ValueError(f"Invalid split: {split}. Must be 'sub_10' or 'full'") if not os.path.exists(dataset_path): raise FileNotFoundError(f"Dataset file not found: {dataset_path}") @@ -143,7 +143,7 @@ def check_prediction_size( Args: predictions: List of prediction dictionaries - split: Either "10" or "full" + split: Either "sub_10" or "full" Returns: Tuple of (is_valid, error_message) @@ -151,7 +151,7 @@ def check_prediction_size( expected_size = EXPECTED_SIZES.get(split) if expected_size is None: - return False, f"Invalid split: {split}. Must be '10' or 'full'" + return False, f"Invalid split: {split}. Must be 'sub_10' or 'full'" actual_size = len(predictions) @@ -286,8 +286,8 @@ def main(): parser.add_argument( "split", type=str, - choices=["10", "full"], - help="Dataset split: '10' for 10%% split (809 entries) or 'full' (8400 entries)", + choices=["sub_10", "full"], + help="Dataset split: 'sub_10' for 10%% split (809 entries) or 'full' (8400 entries)", ) args = parser.parse_args() diff --git a/router_inference/compare_router_accuracy.py b/router_inference/compare_router_accuracy.py index bf05f311..804a0602 100644 --- a/router_inference/compare_router_accuracy.py +++ b/router_inference/compare_router_accuracy.py @@ -199,7 +199,7 @@ def main(): evaluation_dict = build_complete_evaluation_dictionary() # Load RouterEvalBench dataset and create global_index to bloom_level mapping - from datasets import load_dataset + from datasets import load_dataset # type: ignore[import-untyped] # Load the routerevalbench dataset from local path dataset_path = "./dataset/routerevalbench" diff --git a/router_inference/generate_prediction_file.py b/router_inference/generate_prediction_file.py index 3d0f2002..e532cf02 100644 --- a/router_inference/generate_prediction_file.py +++ b/router_inference/generate_prediction_file.py @@ -11,7 +11,7 @@ Usage: python router_inference/generate_prediction_file.py - split: either "10" for 10% split (809 entries) or "full" (8400 entries) + split: either "sub_10" for 10% split (809 entries) or "full" (8400 entries) """ import argparse @@ -25,7 +25,7 @@ # Dataset file paths DATASET_PATHS = { - "10": "./dataset/router_data_10.json", + "sub_10": "./dataset/router_data_10.json", "full": "./dataset/router_data.json", } @@ -56,7 +56,7 @@ def load_dataset(split: str) -> List[Dict[str, Any]]: Load dataset file. Args: - split: Either "10" or "full" + split: Either "sub_10" or "full" Returns: List of dataset entries @@ -64,7 +64,7 @@ def load_dataset(split: str) -> List[Dict[str, Any]]: dataset_path = DATASET_PATHS.get(split) if not dataset_path: - raise ValueError(f"Invalid split: {split}. Must be '10' or 'full'") + raise ValueError(f"Invalid split: {split}. Must be 'sub_10' or 'full'") if not os.path.exists(dataset_path): raise FileNotFoundError(f"Dataset file not found: {dataset_path}") @@ -155,8 +155,8 @@ def main(): parser.add_argument( "split", type=str, - choices=["10", "full"], - help="Dataset split: '10' for 10%% split or 'full'", + choices=["sub_10", "full"], + help="Dataset split: 'sub_10' for 10%% split or 'full'", ) args = parser.parse_args() diff --git a/scripts/process_datasets/prep_datasets.py b/scripts/process_datasets/prep_datasets.py index 543869f2..36398888 100644 --- a/scripts/process_datasets/prep_datasets.py +++ b/scripts/process_datasets/prep_datasets.py @@ -9,7 +9,7 @@ import zlib sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../"))) -from datasets import load_dataset, load_from_disk +from datasets import load_dataset, load_from_disk # type: ignore[import-untyped] from typing import Dict, Any, List save_dir = "./dataset/"