wehub-resource-sync
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
"""Model configurations for SimpleQA evaluation.
|
|
|
|
This module provides model configurations to keep run_naive_simpleqa.py clean.
|
|
"""
|
|
|
|
import os
|
|
from typing import Dict, Optional
|
|
|
|
|
|
def get_model_config(model_name: str) -> Dict[str, Optional[str]]:
|
|
"""
|
|
Get model configuration based on model name.
|
|
|
|
Args:
|
|
model_name: Name of the model (e.g., 'Qwen/Qwen3-VL-4B-Instruct', 'gemini-3-pro-preview')
|
|
|
|
Returns:
|
|
Dictionary with 'api_base', 'api_key', and 'model' keys.
|
|
"""
|
|
model_lower = model_name.lower()
|
|
|
|
# Gemini models
|
|
if "gemini" in model_lower:
|
|
# Check for Vertex AI first
|
|
vertex_api_key = os.getenv("GEMINI_API_KEY")
|
|
use_vertex = os.getenv("GOOGLE_GENAI_USE_VERTEXAI", "").lower() == "true"
|
|
|
|
if use_vertex and vertex_api_key:
|
|
# Using Vertex AI - don't pass api_key, use environment variable instead
|
|
api_key = None # Vertex AI uses environment variable, not api_key parameter
|
|
else:
|
|
# Using standard Gemini API
|
|
api_key = os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
|
|
if not api_key:
|
|
raise ValueError(
|
|
"GOOGLE_API_KEY or GEMINI_API_KEY environment variable is required for Gemini models. "
|
|
"Set it with: export GOOGLE_API_KEY='your-api-key' or export GEMINI_API_KEY='your-api-key' and GOOGLE_GENAI_USE_VERTEXAI=true"
|
|
)
|
|
|
|
# For Gemini models, we use Google's Generative AI SDK directly
|
|
# The api_base is not used for Gemini (SDK handles it internally)
|
|
# But we set a placeholder for compatibility
|
|
api_base = None # Not used for Gemini SDK
|
|
|
|
return {
|
|
"api_base": api_base,
|
|
"api_key": api_key,
|
|
"model": model_name, # Use the model name as-is
|
|
}
|
|
|
|
# Default: assume OpenAI-compatible API (vLLM, etc.)
|
|
return {
|
|
"api_base": os.getenv("API_BASE", "http://localhost:8000/v1"),
|
|
"api_key": os.getenv("API_KEY", "dummy"),
|
|
"model": model_name,
|
|
}
|
|
|
|
|
|
def get_output_filename(
|
|
output_dir: str,
|
|
model_name: str,
|
|
mode: str = "naive",
|
|
num_examples: int = 1000,
|
|
url_screenshot: bool = False,
|
|
task: str = "simpleqa",
|
|
) -> str:
|
|
"""
|
|
Generate output filename with model name and task included.
|
|
|
|
Args:
|
|
output_dir: Base output directory (e.g., 'eval_output/naive_qa')
|
|
model_name: Model name (e.g., 'Qwen/Qwen3-VL-4B-Instruct')
|
|
mode: Evaluation mode ('naive', 'screenshot', 'retrieval')
|
|
num_examples: Number of examples
|
|
url_screenshot: Whether URL screenshot mode is enabled
|
|
task: Task/benchmark name (e.g., 'simpleqa', 'encyclopedic_vqa', 'worldvqa')
|
|
|
|
Returns:
|
|
Full output file path
|
|
"""
|
|
# Clean model name for filename (replace special chars)
|
|
model_safe = (
|
|
model_name.replace("/", "_").replace(":", "_").replace("-", "_").lower()
|
|
)
|
|
|
|
# Build filename components (task first for easy distinction)
|
|
parts = [task]
|
|
if url_screenshot:
|
|
parts.append("urlscreenshot")
|
|
parts.append(mode)
|
|
parts.append(model_safe)
|
|
parts.append(str(num_examples))
|
|
|
|
filename = "_".join(parts) + ".jsonl"
|
|
return os.path.join(output_dir, filename)
|