diff --git a/config/experiments.yaml b/config/experiments.yaml index 7bae119..4dc61c4 100644 --- a/config/experiments.yaml +++ b/config/experiments.yaml @@ -341,3 +341,204 @@ experiments: - rouge2 - rougeL + + # QMSum - DeepSeek V4 Flash baseline + - id: qmsum_baseline_deepseekv4flash + tags: [qmsum] + model: + name: openrouter:deepseek/deepseek-v4-flash + mode: baseline + generation: + max_tokens: 256 + temperature: 1.2 + top_p: 1.0 + dataset: + name: qmsum + split: test + num_examples: 281 + seed: 40 + task: + name: summarization + metrics: + - rouge1 + - rouge2 + - rougeL + + # QMSum - DeepSeek V4 Pro baseline + - id: qmsum_baseline_deepseekv4pro + tags: [qmsum] + model: + name: openrouter:deepseek/deepseek-v4-pro + mode: baseline + generation: + max_tokens: 256 + temperature: 1.2 + top_p: 1.0 + dataset: + name: qmsum + split: test + num_examples: 281 + seed: 40 + task: + name: summarization + metrics: + - rouge1 + - rouge2 + - rougeL + + # CLINC OOS - DeepSeek V4 Flash baseline (Classify task) + # NOTE: num_examples/seed here are placeholders — confirm the sample size + # and seed used to produce the other models' Classify numbers (not found + # anywhere in this repo) before treating this as comparable. + - id: classify_clinc_baseline_deepseekv4flash + tags: [clinc_oos] + model: + name: openrouter:deepseek/deepseek-v4-flash + mode: baseline + generation: + max_tokens: 16 + temperature: 0.0 + dataset: + name: clinc_oos + split: test + num_examples: 400 + seed: 42 + task: + name: classification + metrics: + - exact_match + + # CLINC OOS - DeepSeek V4 Pro baseline (Classify task) + - id: classify_clinc_baseline_deepseekv4pro + tags: [clinc_oos] + model: + name: openrouter:deepseek/deepseek-v4-pro + mode: baseline + generation: + max_tokens: 16 + temperature: 0.0 + dataset: + name: clinc_oos + split: test + num_examples: 400 + seed: 42 + task: + name: classification + metrics: + - exact_match + + # FinanceBench - DeepSeek V4 Flash baseline (Compress task) + # NOTE: uses evidence_text_full_page fallback unless FINANCEBENCH_PDF_DIR is + # set to a local clone of patronus-ai/financebench PDFs. Confirm this matches + # the context source used for the other models' Compress numbers. + - id: financebench_baseline_deepseekv4flash + model: + name: openrouter:deepseek/deepseek-v4-flash + mode: baseline + dataset: + name: financebench + split: train + num_examples: 150 + seed: 42 + task: + name: rag_qa + metrics: + - exact_match + - f1 + + # FinanceBench - DeepSeek V4 Pro baseline (Compress task) + - id: financebench_baseline_deepseekv4pro + model: + name: openrouter:deepseek/deepseek-v4-pro + mode: baseline + dataset: + name: financebench + split: train + num_examples: 150 + seed: 42 + task: + name: rag_qa + metrics: + - exact_match + - f1 + + # LongExtractBench - DeepSeek V4 Flash baseline (Extract task) + # num_examples: 50 = the full LongExtractBench-50 set (no sampling needed). + - id: extract_longextractbench_deepseekv4flash + tags: [longextractbench] + model: + name: openrouter:deepseek/deepseek-v4-flash + mode: baseline + generation: + max_tokens: 12000 + temperature: 0.0 + dataset: + name: longextractbench + split: test + num_examples: 50 + seed: 42 + task: + name: extraction + metrics: + - field_precision + - field_recall + - field_f1 + - parse_rate + + # LongExtractBench - DeepSeek V4 Pro baseline (Extract task) + - id: extract_longextractbench_deepseekv4pro + tags: [longextractbench] + model: + name: openrouter:deepseek/deepseek-v4-pro + mode: baseline + generation: + max_tokens: 12000 + temperature: 0.0 + dataset: + name: longextractbench + split: test + num_examples: 50 + seed: 42 + task: + name: extraction + metrics: + - field_precision + - field_recall + - field_f1 + - parse_rate + + # CUAD - DeepSeek V4 Flash baseline (Extract task, matches methodology used + # for the other benchmark pages' Extract number, instead of LongExtractBench + # — CUAD contexts are ~11K chars vs LongExtractBench's ~250K-token documents, + # so this is far cheaper/faster and avoids the output-truncation problem + # LongExtractBench hit at 12K max_tokens). + - id: extract_cuad_baseline_deepseekv4flash + model: + name: openrouter:deepseek/deepseek-v4-flash + mode: baseline + dataset: + name: cuad + split: test + num_examples: 150 + seed: 42 + task: + name: rag_qa + metrics: + - exact_match + - f1 + + # CUAD - DeepSeek V4 Pro baseline (Extract task) + - id: extract_cuad_baseline_deepseekv4pro + model: + name: openrouter:deepseek/deepseek-v4-pro + mode: baseline + dataset: + name: cuad + split: test + num_examples: 150 + seed: 42 + task: + name: rag_qa + metrics: + - exact_match + - f1 diff --git a/src/core/registry.py b/src/core/registry.py index 35e4f6d..0f3a440 100644 --- a/src/core/registry.py +++ b/src/core/registry.py @@ -5,6 +5,8 @@ # Import model clients from models.openai_client import OpenAIClient from models.gemini_client import GeminiClient +from models.deepseek_client import DeepSeekClient +from models.openrouter_client import OpenRouterClient from models.scaledown_client import ScaleDownClient from models.scaledown_summarize_client import ScaleDownSummarizeClient from models.base import ModelClient, PricingInfo, DEFAULT_MODEL_PRICING @@ -16,12 +18,17 @@ from dataset.msmarco import MSMARCODataset from dataset.financebench import FinanceBenchDataset from dataset.qmsum import QMSumDataset +from dataset.clinc_oos import ClincOosDataset +from dataset.longextractbench import LongExtractBenchDataset +from dataset.cuad import CuadDataset from dataset.base import Dataset # Import tasks from tasks.rag_task import RAGTask from tasks.retrieval_task import RetrievalTask from tasks.summarization_task import SummarizationTask +from tasks.classification_task import ClassificationTask +from tasks.extraction_task import ExtractionTask from tasks.base import Task # Import metrics @@ -38,6 +45,12 @@ compute_set_recall, compute_set_f1, ) +from metrics.field_extraction import ( + compute_field_precision, + compute_field_recall, + compute_field_f1, + compute_parse_rate, +) from metrics.retrieval import ( compute_retrieval_precision, compute_retrieval_recall, @@ -52,6 +65,8 @@ MODEL_REGISTRY: Dict[str, type] = { "openai": OpenAIClient, "gemini": GeminiClient, + "deepseek": DeepSeekClient, + "openrouter": OpenRouterClient, } DATASET_REGISTRY: Dict[str, type] = { @@ -61,12 +76,17 @@ "msmarco": MSMARCODataset, "financebench": FinanceBenchDataset, "qmsum": QMSumDataset, + "clinc_oos": ClincOosDataset, + "longextractbench": LongExtractBenchDataset, + "cuad": CuadDataset, } TASK_REGISTRY: Dict[str, type] = { "rag_qa": RAGTask, "retrieval_task": RetrievalTask, "summarization": SummarizationTask, + "classification": ClassificationTask, + "extraction": ExtractionTask, } METRIC_REGISTRY: Dict[str, Callable] = { @@ -83,6 +103,10 @@ "retrieval_precision": compute_retrieval_precision, "retrieval_recall": compute_retrieval_recall, "retrieval_f1": compute_retrieval_f1, + "field_precision": compute_field_precision, + "field_recall": compute_field_recall, + "field_f1": compute_field_f1, + "parse_rate": compute_parse_rate, } diff --git a/src/dataset/clinc_oos.py b/src/dataset/clinc_oos.py new file mode 100644 index 0000000..88c8ef9 --- /dev/null +++ b/src/dataset/clinc_oos.py @@ -0,0 +1,61 @@ +"""CLINC OOS dataset loader. + +CLINC OOS (Larson et al., 2019) is an intent classification benchmark with +150 in-scope intents across 10 domains plus an out-of-scope class. Used here +for the "Classify" task on the ScaleDown benchmark pages. + +Loaded from HuggingFace clinc_oos ("plus" config, which includes out-of-scope +examples in addition to the 150 in-domain intents). +""" +import random +from typing import Iterator, Optional + +import datasets as hf_datasets + +from dataset.base import Dataset, Example + + +class ClincOosDataset(Dataset): + """CLINC OOS single-turn intent classification dataset.""" + + def __init__(self, name: str = "clinc_oos"): + super().__init__(name) + + def load_examples( + self, + split: str = "test", + limit: Optional[int] = None, + seed: int = 42, + ) -> Iterator[Example]: + """Load and yield CLINC OOS examples. + + Args: + split: Dataset split (train/validation/test). + limit: Maximum number of examples to load (None = all). + seed: Random seed for deterministic sampling. + + Yields: + Example dicts with {id, context ("", unused), question (query text), + answer ([intent label]), labels (full label list, for prompt building)}. + """ + hf_dataset = hf_datasets.load_dataset("clinc_oos", "plus", split=split) + label_names = hf_dataset.features["intent"].names + + all_examples = list(hf_dataset) + total_size = len(all_examples) + if limit is not None and limit < total_size: + random.seed(seed) + indices = random.sample(range(total_size), limit) + indices.sort() + else: + indices = range(total_size) + + for idx in indices: + row = all_examples[idx] + yield Example( + id=str(idx), + context="", + question=row["text"], + answer=[label_names[row["intent"]]], + labels=label_names, + ) diff --git a/src/dataset/cuad.py b/src/dataset/cuad.py new file mode 100644 index 0000000..3d8626e --- /dev/null +++ b/src/dataset/cuad.py @@ -0,0 +1,63 @@ +"""CUAD dataset loader. + +CUAD (Contract Understanding Atticus Dataset) is a legal-contract span-extraction +QA benchmark — the "Extract" task on the existing ScaleDown benchmark pages +(Grok/OpenAI/Gemini) was run against CUAD with the F1 metric. + +Loaded from chenghao/cuad_qa (a SQuAD-format mirror of the original CUAD +release; the official theatticusproject/cuad-qa loader script is no longer +loadable under current `datasets` versions, which dropped script-based +loaders — this mirror has the same fields as a plain Parquet dataset). +""" +import random +from typing import Iterator, Optional + +from datasets import load_dataset as hf_load_dataset + +from dataset.base import Dataset, Example + + +class CuadDataset(Dataset): + """CUAD contract clause extraction dataset (SQuAD-format span QA).""" + + def __init__(self, name: str = "cuad"): + super().__init__(name) + + def load_examples( + self, + split: str = "test", + limit: Optional[int] = None, + seed: int = 42, + ) -> Iterator[Example]: + """Load and yield CUAD examples. + + Args: + split: Dataset split (train/test). + limit: Maximum number of examples to load (None = all). + seed: Random seed for deterministic sampling. + + Yields: + Example dicts with {id, context (contract text), question (clause + prompt), answer (list of gold answer strings — empty list if the + clause type isn't present in this contract)}. + """ + hf_dataset = hf_load_dataset("chenghao/cuad_qa", split=split) + all_examples = list(hf_dataset) + + total_size = len(all_examples) + if limit is not None and limit < total_size: + random.seed(seed) + indices = random.sample(range(total_size), limit) + indices.sort() + else: + indices = range(total_size) + + for idx in indices: + row = all_examples[idx] + answer_texts = row["answers"]["text"] if row["answers"]["text"] else [""] + yield Example( + id=row["id"], + context=row["context"], + question=row["question"], + answer=answer_texts, + ) diff --git a/src/dataset/financebench.py b/src/dataset/financebench.py index d807ed6..5cdb2fb 100644 --- a/src/dataset/financebench.py +++ b/src/dataset/financebench.py @@ -1,4 +1,5 @@ """FinanceBench dataset loader.""" +import os import random from pathlib import Path from typing import Iterator, Optional @@ -7,9 +8,19 @@ from dataset.base import Dataset, Example -# Default path to cloned patronus-ai/financebench PDF directory +# Default path to cloned patronus-ai/financebench PDF directory. +# Override with the FINANCEBENCH_PDF_DIR env var, or pass pdf_dir=... to +# FinanceBenchDataset(). If the directory doesn't exist, load_examples() +# falls back to evidence_text_full_page (see below) rather than failing — +# but that means a run without the PDFs is scoring different (shorter, +# pre-extracted) context than a run with them. Confirm which one produced +# any FinanceBench numbers you're comparing against before trusting a +# cross-model comparison. DEFAULT_PDF_DIR = Path( - "/Users/soham/Desktop/scaledown-project/financebench-data/pdfs" + os.environ.get( + "FINANCEBENCH_PDF_DIR", + str(Path(__file__).resolve().parents[2] / "data" / "financebench-pdfs"), + ) ) diff --git a/src/dataset/longextractbench.py b/src/dataset/longextractbench.py new file mode 100644 index 0000000..4d5331f --- /dev/null +++ b/src/dataset/longextractbench.py @@ -0,0 +1,109 @@ +"""LongExtractBench dataset loader. + +LongExtractBench-50 (micro1-inc/longextract-bench-50 on HuggingFace) is a +50-document benchmark for schema-driven structured extraction from long, +table-heavy PDFs (financial filings, government statistics, healthcare +reports, etc). Each example is a folder containing: + document.pdf - the source document + schema.json - target JSON Schema for extraction + ground_truth.json - human-reconciled gold extraction + +Unlike the other datasets in this repo, LongExtractBench is not a +`datasets`-loadable parquet dataset (it's a raw file repo, one folder per +example), so this loader uses `huggingface_hub.snapshot_download` to pull +the whole repo locally and reads each folder directly. PDF text extraction +mirrors dataset/financebench.py's approach (PyMuPDF). +""" +import json +import random +from pathlib import Path +from typing import Iterator, Optional + +from huggingface_hub import snapshot_download + +from dataset.base import Dataset, Example + +REPO_ID = "micro1-inc/longextract-bench-50" + + +def _load_pdf_text(pdf_path: Path) -> str: + """Extract full text from a PDF using PyMuPDF. + + Args: + pdf_path: Path to the PDF file. + + Returns: + Full document text with pages joined by double newlines. + """ + import fitz # PyMuPDF + + doc = fitz.open(str(pdf_path)) + pages = [page.get_text() for page in doc] + doc.close() + return "\n\n".join(pages) + + +class LongExtractBenchDataset(Dataset): + """LongExtractBench-50 schema-driven structured extraction dataset.""" + + def __init__(self, name: str = "longextractbench", cache_dir: Optional[str] = None): + """Initialize LongExtractBench dataset. + + Args: + name: Dataset name. + cache_dir: Optional huggingface_hub cache directory override. + """ + super().__init__(name) + self.cache_dir = cache_dir + + def load_examples( + self, + split: str = "test", + limit: Optional[int] = None, + seed: int = 42, + ) -> Iterator[Example]: + """Load and yield LongExtractBench examples. + + Args: + split: Unused — LongExtractBench-50 has a single flat set of examples. + limit: Maximum number of examples to load (None = all 50). + seed: Random seed for deterministic sampling. + + Yields: + Example dicts with {id, context (document text), question (schema as + JSON string, for logging), answer ([ground truth as JSON string]), + schema (parsed dict), ground_truth (parsed dict)}. + """ + repo_dir = Path(snapshot_download( + repo_id=REPO_ID, repo_type="dataset", cache_dir=self.cache_dir, + )) + + example_dirs = sorted( + d for d in repo_dir.iterdir() + if d.is_dir() and (d / "document.pdf").exists() + and (d / "schema.json").exists() + and (d / "ground_truth.json").exists() + ) + + total_size = len(example_dirs) + if limit is not None and limit < total_size: + random.seed(seed) + indices = random.sample(range(total_size), limit) + indices.sort() + else: + indices = range(total_size) + + for idx in indices: + d = example_dirs[idx] + schema = json.loads((d / "schema.json").read_text()) + ground_truth = json.loads((d / "ground_truth.json").read_text()) + context = _load_pdf_text(d / "document.pdf") + + yield Example( + id=d.name, + context=context, + question=json.dumps(schema), + answer=[json.dumps(ground_truth)], + schema=schema, + ground_truth=ground_truth, + ) diff --git a/src/metrics/field_extraction.py b/src/metrics/field_extraction.py new file mode 100644 index 0000000..8385a53 --- /dev/null +++ b/src/metrics/field_extraction.py @@ -0,0 +1,105 @@ +"""Field-level extraction metrics for schema-driven structured extraction (e.g. LongExtractBench). + +Scores a predicted JSON object against a gold JSON object by flattening both +into dotted-path -> stringified-value pairs and computing set-based +precision/recall/F1 over those pairs. This is a first-pass scoring method: +it does exact string match per field after light normalization (strip, +lowercase, collapse whitespace) — no fuzzy numeric tolerance, no date +canonicalization, no schema validation. Good enough to get a directional +signal; revisit before publishing a precise number. + +Follows the one-metric-name -> one-function -> one-key convention used +elsewhere in this repo (see metrics/retrieval.py), so each of precision/ +recall/F1/parse-rate is its own registrable metric rather than one function +returning several keys. +""" +import json +from typing import List, Dict, Callable, Optional, Any + + +def _flatten(obj: Any, prefix: str = "") -> Dict[str, str]: + """Flatten a nested JSON-like structure into {dotted.path: stringified_value}. + + Args: + obj: Parsed JSON value (dict, list, or scalar). + prefix: Dotted-path prefix accumulated from parent keys. + + Returns: + Flat dict mapping dotted paths to normalized string values. + """ + out: Dict[str, str] = {} + if isinstance(obj, dict): + for k, v in obj.items(): + path = f"{prefix}.{k}" if prefix else str(k) + out.update(_flatten(v, path)) + elif isinstance(obj, list): + for i, v in enumerate(obj): + out.update(_flatten(v, f"{prefix}[{i}]")) + else: + out[prefix] = str(obj).strip().lower() + return out + + +def _safe_parse(text: str) -> Optional[dict]: + """Parse a JSON object from text, tolerant of stray whitespace. Returns None on failure.""" + try: + parsed = json.loads(text.strip()) + return parsed if isinstance(parsed, dict) else None + except (json.JSONDecodeError, AttributeError): + return None + + +def _per_example_scores(predictions: List[str], references: List[List[str]]) -> List[Dict[str, float]]: + """Shared per-example precision/recall/F1/parsed-ok computation.""" + scores = [] + for pred_text, refs in zip(predictions, references): + gold_text = refs[0] if refs else "{}" + gold = _flatten(_safe_parse(gold_text) or {}) + pred_parsed = _safe_parse(pred_text) + pred = _flatten(pred_parsed or {}) + + if not gold and not pred: + scores.append({"precision": 1.0, "recall": 1.0, "f1": 1.0, "parsed_ok": 1.0 if pred_parsed is not None else 0.0}) + continue + if not pred or not gold: + scores.append({"precision": 0.0, "recall": 0.0, "f1": 0.0, "parsed_ok": 1.0 if pred_parsed is not None else 0.0}) + continue + + correct = sum(1 for k, v in pred.items() if k in gold and gold[k] == v) + precision = correct / len(pred) + recall = correct / len(gold) + f1 = 0.0 if precision + recall == 0 else 2 * precision * recall / (precision + recall) + scores.append({"precision": precision, "recall": recall, "f1": f1, "parsed_ok": 1.0 if pred_parsed is not None else 0.0}) + return scores + + +def compute_field_precision( + predictions: List[str], references: List[List[str]], normalize_fn: Optional[Callable[[str], str]] = None, +) -> Dict[str, float]: + scores = _per_example_scores(predictions, references) + n = len(scores) or 1 + return {"field_precision": sum(s["precision"] for s in scores) / n} + + +def compute_field_recall( + predictions: List[str], references: List[List[str]], normalize_fn: Optional[Callable[[str], str]] = None, +) -> Dict[str, float]: + scores = _per_example_scores(predictions, references) + n = len(scores) or 1 + return {"field_recall": sum(s["recall"] for s in scores) / n} + + +def compute_field_f1( + predictions: List[str], references: List[List[str]], normalize_fn: Optional[Callable[[str], str]] = None, +) -> Dict[str, float]: + scores = _per_example_scores(predictions, references) + n = len(scores) or 1 + return {"field_f1": sum(s["f1"] for s in scores) / n} + + +def compute_parse_rate( + predictions: List[str], references: List[List[str]], normalize_fn: Optional[Callable[[str], str]] = None, +) -> Dict[str, float]: + scores = _per_example_scores(predictions, references) + n = len(scores) or 1 + return {"parse_rate": sum(s["parsed_ok"] for s in scores) / n} diff --git a/src/models/base.py b/src/models/base.py index 0d0f4b2..8564e45 100644 --- a/src/models/base.py +++ b/src/models/base.py @@ -45,6 +45,16 @@ def calculate_cost(self, input_tokens: int, output_tokens: int) -> Optional[floa "gemini-2.5-flash": PricingInfo(input_per_1m_tokens=0.30, output_per_1m_tokens=2.50), "gemini-2.5-flash-lite": PricingInfo(input_per_1m_tokens=0.10, output_per_1m_tokens=0.40), "gemini-2.5-pro": PricingInfo(input_per_1m_tokens=1.25, output_per_1m_tokens=10.00), + + # DeepSeek models (cache-miss input rate; see deepseek.ai/pricing, Aug 2026) + "deepseek-v4-flash": PricingInfo(input_per_1m_tokens=0.14, output_per_1m_tokens=0.28), + "deepseek-v4-pro": PricingInfo(input_per_1m_tokens=0.435, output_per_1m_tokens=0.87), + + # DeepSeek via OpenRouter (live from openrouter.ai/api/v1/models, Aug 2026). + # OpenRouterClient reports actual per-call cost from the API response when + # available, so these are a fallback only. + "deepseek/deepseek-v4-flash": PricingInfo(input_per_1m_tokens=0.14, output_per_1m_tokens=0.28), + "deepseek/deepseek-v4-pro": PricingInfo(input_per_1m_tokens=0.63168, output_per_1m_tokens=1.26336), } diff --git a/src/models/deepseek_client.py b/src/models/deepseek_client.py new file mode 100644 index 0000000..fd8378f --- /dev/null +++ b/src/models/deepseek_client.py @@ -0,0 +1,161 @@ +"""DeepSeek model client. + +DeepSeek exposes an OpenAI-compatible Chat Completions API, so this client +reuses the `openai` SDK pointed at DeepSeek's base URL rather than a bespoke +HTTP layer. Token usage for cost calculation comes straight from the API +response (DeepSeek has no public tokenizer package for local counting). + +Environment Variables: + DEEPSEEK_API_KEY: Required API key for DeepSeek. + +Model names (verify against your account before relying on these — DeepSeek +was mid-rollout on V4 Pro as of Aug 2026, see PR description): + deepseek-v4-flash + deepseek-v4-pro +""" +import logging +import time +from typing import Optional + +from openai import OpenAI, RateLimitError, APIError, APIConnectionError + +from models.base import ModelClient, GenerationInput, ModelOutput, PricingInfo + +logger = logging.getLogger(__name__) + +DEEPSEEK_BASE_URL = "https://api.deepseek.com" + + +class DeepSeekClient(ModelClient): + """Client for DeepSeek models via the OpenAI-compatible Chat Completions API.""" + + def __init__( + self, + model_name: str, + api_key: Optional[str] = None, + pricing: Optional[PricingInfo] = None, + max_retries: int = 3, + retry_delay: float = 2.0, + ): + """Initialize DeepSeek client. + + Args: + model_name: DeepSeek model name (e.g., "deepseek-v4-flash", "deepseek-v4-pro"). + api_key: DeepSeek API key. If None, uses DEEPSEEK_API_KEY env var. + pricing: Optional pricing information for cost calculation. + max_retries: Maximum number of retries for failed requests. + retry_delay: Initial delay in seconds for exponential backoff. + """ + super().__init__(model_name, api_key, pricing) + self.provider = "deepseek" + self.max_retries = max_retries + self.retry_delay = retry_delay + + # DEEPSEEK_API_KEY is read by the OpenAI SDK only if api_key is passed + # explicitly here; there is no implicit env var fallback like OPENAI_API_KEY. + import os + resolved_key = api_key or os.environ.get("DEEPSEEK_API_KEY") + if not resolved_key: + raise ValueError( + "DeepSeek API key required: pass api_key or set DEEPSEEK_API_KEY." + ) + self.client = OpenAI(api_key=resolved_key, base_url=DEEPSEEK_BASE_URL) + + def generate(self, inp: GenerationInput) -> ModelOutput: + """Generate text using DeepSeek's Chat Completions API. + + Args: + inp: GenerationInput with prompt and configuration. + + Returns: + ModelOutput with generated text and metadata. + """ + messages = [] + if inp.system_prompt: + messages.append({"role": "system", "content": inp.system_prompt}) + + user_content = inp.user_prompt + if inp.context: + user_content = f"{inp.context}\n\n{inp.user_prompt}" + messages.append({"role": "user", "content": user_content}) + + params = { + "model": self.model_name, + "messages": messages, + } + if inp.temperature is not None: + params["temperature"] = inp.temperature + if inp.max_output_tokens is not None: + params["max_tokens"] = inp.max_output_tokens + if inp.top_p is not None: + params["top_p"] = inp.top_p + if inp.stop is not None: + params["stop"] = inp.stop + if inp.response_schema is not None: + # DeepSeek supports {"type": "json_object"} but not strict json_schema + # mode as of Aug 2026 — fall back to json_object and rely on the + # system prompt to describe the schema. + params["response_format"] = {"type": "json_object"} + + last_exception = None + for attempt in range(self.max_retries): + try: + start_time = time.time() + response = self.client.chat.completions.create(**params) + latency_ms = (time.time() - start_time) * 1000 + + text = response.choices[0].message.content or "" + input_tokens = response.usage.prompt_tokens if response.usage else 0 + output_tokens = response.usage.completion_tokens if response.usage else 0 + cost_usd = self._calculate_cost(input_tokens, output_tokens) + + return ModelOutput( + text=text, + latency_ms=latency_ms, + input_tokens=input_tokens, + output_tokens=output_tokens, + cost_usd=cost_usd, + input_question=inp.user_prompt, + input_context=inp.context, + input_system_prompt=inp.system_prompt, + ) + + except RateLimitError as e: + last_exception = e + if attempt < self.max_retries - 1: + delay = self.retry_delay * (2 ** attempt) + logger.info(f"Rate limit hit, retrying in {delay}s (attempt {attempt + 1}/{self.max_retries})") + time.sleep(delay) + else: + logger.error(f"Rate limit exceeded after {self.max_retries} retries") + + except (APIError, APIConnectionError) as e: + last_exception = e + if attempt < self.max_retries - 1: + delay = self.retry_delay * (2 ** attempt) + logger.info(f"API error, retrying in {delay}s (attempt {attempt + 1}/{self.max_retries}): {e}") + time.sleep(delay) + else: + logger.error(f"API error after {self.max_retries} retries: {e}") + except Exception as e: + logger.error(f"Unexpected error in DeepSeekClient.generate: {e}", exc_info=True) + raise + if last_exception is not None: + raise last_exception + else: + raise RuntimeError("DeepSeekClient.generate failed with an unexpected error and no exception was captured.") + + def count_tokens(self, text: str) -> int: + """Approximate token count for text not yet sent to the API. + + DeepSeek has no public local tokenizer, so this is a rough heuristic + (chars / 4) used only for pre-flight estimates. Actual cost accounting + always uses the `usage` block returned by the API in generate(). + + Args: + text: Input text to count tokens for. + + Returns: + Approximate token count. + """ + return max(1, len(text) // 4) diff --git a/src/models/openrouter_client.py b/src/models/openrouter_client.py new file mode 100644 index 0000000..0cbcb2c --- /dev/null +++ b/src/models/openrouter_client.py @@ -0,0 +1,170 @@ +"""OpenRouter model client. + +OpenRouter exposes an OpenAI-compatible Chat Completions API that proxies to +many upstream providers, including DeepSeek. Model names use OpenRouter's +"provider/model" convention (e.g. "deepseek/deepseek-v4-flash"). + +Reasoning is disabled by default. DeepSeek V4 Flash/Pro on OpenRouter default +to an extended-thinking mode that spends completion tokens on a hidden +`reasoning` field before (or instead of, if max_tokens runs out first) the +visible `content` — confirmed live: a 10-token budget with reasoning enabled +returned `content: null` and `finish_reason: "length"`, all 10 tokens burned +on reasoning. That silently truncates low-max_tokens tasks (e.g. Classify) +and isn't apples-to-apples against non-reasoning competitor models on the +other benchmark pages (e.g. GPT-5.4 Mini). Pass reasoning_enabled=True to +this client only if you deliberately want to benchmark the reasoning variant. + +Environment Variables: + OPENROUTER_API_KEY: Required API key for OpenRouter. +""" +import logging +import time +from typing import Optional + +from openai import OpenAI, RateLimitError, APIError, APIConnectionError + +from models.base import ModelClient, GenerationInput, ModelOutput, PricingInfo + +logger = logging.getLogger(__name__) + +OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1" + + +class OpenRouterClient(ModelClient): + """Client for OpenRouter-routed models via the OpenAI-compatible Chat Completions API.""" + + def __init__( + self, + model_name: str, + api_key: Optional[str] = None, + pricing: Optional[PricingInfo] = None, + max_retries: int = 3, + retry_delay: float = 2.0, + reasoning_enabled: bool = False, + ): + """Initialize OpenRouter client. + + Args: + model_name: OpenRouter model id (e.g., "deepseek/deepseek-v4-flash"). + api_key: OpenRouter API key. If None, uses OPENROUTER_API_KEY env var. + pricing: Optional pricing information for cost calculation. + max_retries: Maximum number of retries for failed requests. + retry_delay: Initial delay in seconds for exponential backoff. + reasoning_enabled: Whether to allow the upstream model's extended + thinking/reasoning mode. Default False — see module docstring. + """ + super().__init__(model_name, api_key, pricing) + self.provider = "openrouter" + self.max_retries = max_retries + self.retry_delay = retry_delay + self.reasoning_enabled = reasoning_enabled + + import os + resolved_key = api_key or os.environ.get("OPENROUTER_API_KEY") + if not resolved_key: + raise ValueError( + "OpenRouter API key required: pass api_key or set OPENROUTER_API_KEY." + ) + self.client = OpenAI(api_key=resolved_key, base_url=OPENROUTER_BASE_URL) + + def generate(self, inp: GenerationInput) -> ModelOutput: + """Generate text using OpenRouter's Chat Completions API. + + Args: + inp: GenerationInput with prompt and configuration. + + Returns: + ModelOutput with generated text and metadata. + """ + messages = [] + if inp.system_prompt: + messages.append({"role": "system", "content": inp.system_prompt}) + + user_content = inp.user_prompt + if inp.context: + user_content = f"{inp.context}\n\n{inp.user_prompt}" + messages.append({"role": "user", "content": user_content}) + + params = { + "model": self.model_name, + "messages": messages, + "extra_body": {"reasoning": {"enabled": self.reasoning_enabled}}, + } + if inp.temperature is not None: + params["temperature"] = inp.temperature + if inp.max_output_tokens is not None: + params["max_tokens"] = inp.max_output_tokens + if inp.top_p is not None: + params["top_p"] = inp.top_p + if inp.stop is not None: + params["stop"] = inp.stop + if inp.response_schema is not None: + params["response_format"] = {"type": "json_object"} + + last_exception = None + for attempt in range(self.max_retries): + try: + start_time = time.time() + response = self.client.chat.completions.create(**params) + latency_ms = (time.time() - start_time) * 1000 + + if not response.choices: + raise RuntimeError( + f"OpenRouter returned no choices (response id={getattr(response, 'id', '?')}); " + "seen on at least one LongExtractBench example with a very large document — " + "root cause not yet identified (context window is 1M tokens, well above what " + "was sent, so it isn't a length issue). Not retried indefinitely by design: " + "the outer retry loop still gets 3 attempts, but if it fails all 3 it's treated " + "as a real per-example failure rather than transient." + ) + text = response.choices[0].message.content or "" + input_tokens = response.usage.prompt_tokens if response.usage else 0 + output_tokens = response.usage.completion_tokens if response.usage else 0 + + # Prefer OpenRouter's actual reported cost (it varies by + # upstream provider routing) over our static PricingInfo when + # available. + actual_cost = getattr(response.usage, "cost", None) if response.usage else None + cost_usd = actual_cost if actual_cost is not None else self._calculate_cost(input_tokens, output_tokens) + + return ModelOutput( + text=text, + latency_ms=latency_ms, + input_tokens=input_tokens, + output_tokens=output_tokens, + cost_usd=cost_usd, + input_question=inp.user_prompt, + input_context=inp.context, + input_system_prompt=inp.system_prompt, + ) + + except RateLimitError as e: + last_exception = e + if attempt < self.max_retries - 1: + delay = self.retry_delay * (2 ** attempt) + logger.info(f"Rate limit hit, retrying in {delay}s (attempt {attempt + 1}/{self.max_retries})") + time.sleep(delay) + else: + logger.error(f"Rate limit exceeded after {self.max_retries} retries") + + except (APIError, APIConnectionError) as e: + last_exception = e + if attempt < self.max_retries - 1: + delay = self.retry_delay * (2 ** attempt) + logger.info(f"API error, retrying in {delay}s (attempt {attempt + 1}/{self.max_retries}): {e}") + time.sleep(delay) + else: + logger.error(f"API error after {self.max_retries} retries: {e}") + except Exception as e: + logger.error(f"Unexpected error in OpenRouterClient.generate: {e}", exc_info=True) + raise + if last_exception is not None: + raise last_exception + else: + raise RuntimeError("OpenRouterClient.generate failed with an unexpected error and no exception was captured.") + + def count_tokens(self, text: str) -> int: + """Approximate token count (chars / 4) — OpenRouter has no single tokenizer + since it proxies many upstream models. Actual cost accounting always uses + the `usage` block returned by the API in generate().""" + return max(1, len(text) // 4) diff --git a/src/tasks/classification_task.py b/src/tasks/classification_task.py new file mode 100644 index 0000000..7caa20b --- /dev/null +++ b/src/tasks/classification_task.py @@ -0,0 +1,62 @@ +"""Intent classification task (e.g. CLINC OOS).""" +import re +from typing import Dict, Any, Optional + +from models.base import GenerationInput +from tasks.base import Task + +_SYSTEM_PROMPT_TEMPLATE = ( + "Classify the user's message into exactly one of the following intents:\n" + "{labels}\n\n" + "Respond with only the intent label, exactly as written above, and nothing else." +) + + +class ClassificationTask(Task): + """Single-label intent classification task. + + Designed for CLINC OOS and similar closed-label classification benchmarks. + Expects each example to carry a `labels` field (the full label set) so the + prompt can enumerate valid intents. + """ + + def __init__(self, name: str = "classification", dataset_name: Optional[str] = None): + super().__init__(name) + self.dataset_name = dataset_name + + def build_prompt(self, example: Dict[str, Any], **kwargs) -> GenerationInput: + """Build GenerationInput for intent classification. + + Args: + example: Example with 'question' (utterance) and 'labels' (label list). + **kwargs: Remaining generation parameters forwarded to GenerationInput. + + Returns: + GenerationInput configured for the target client. + """ + # See tasks/rag_task.py for why this pop is needed. + kwargs.pop("model_mode", None) + + labels = example.get("labels", []) + system_prompt = _SYSTEM_PROMPT_TEMPLATE.format(labels=", ".join(labels)) + return GenerationInput( + system_prompt=system_prompt, + context=None, + user_prompt=example.get("question", ""), + **kwargs, + ) + + def postprocess(self, output: str) -> str: + """Take the first non-empty line as the predicted label.""" + for line in output.strip().split("\n"): + line = line.strip() + if line: + return line + return output.strip() + + def normalize(self, text: str) -> str: + """Normalize label text for comparison: lowercase, spaces/hyphens -> underscore.""" + text = text.strip().lower() + text = re.sub(r"[\s\-]+", "_", text) + text = re.sub(r"[^a-z0-9_]", "", text) + return text diff --git a/src/tasks/extraction_task.py b/src/tasks/extraction_task.py new file mode 100644 index 0000000..61c7c16 --- /dev/null +++ b/src/tasks/extraction_task.py @@ -0,0 +1,62 @@ +"""Schema-driven structured extraction task (e.g. LongExtractBench).""" +import json +from typing import Dict, Any, Optional + +from models.base import GenerationInput +from tasks.base import Task + +_SYSTEM_PROMPT_TEMPLATE = ( + "Extract the requested fields from the document below and return ONLY a " + "single valid JSON object matching this schema. Do not include commentary, " + "explanations, or markdown code fences — output raw JSON only.\n\n" + "Schema:\n{schema}" +) + + +class ExtractionTask(Task): + """Structured extraction task scored via field-level F1 on parsed JSON. + + Designed for LongExtractBench and similar schema-driven extraction + benchmarks. Expects each example to carry a `schema` field (a JSON Schema + dict) so the prompt can describe the target structure. + """ + + def __init__(self, name: str = "extraction", dataset_name: Optional[str] = None): + super().__init__(name) + self.dataset_name = dataset_name + + def build_prompt(self, example: Dict[str, Any], **kwargs) -> GenerationInput: + """Build GenerationInput for schema-driven extraction. + + Args: + example: Example with 'context' (document text) and 'schema' (JSON Schema dict). + **kwargs: Remaining generation parameters forwarded to GenerationInput. + + Returns: + GenerationInput configured for the target client. + """ + # See tasks/rag_task.py for why this pop is needed. + kwargs.pop("model_mode", None) + + schema = example.get("schema", {}) + system_prompt = _SYSTEM_PROMPT_TEMPLATE.format(schema=json.dumps(schema)) + return GenerationInput( + system_prompt=system_prompt, + context=f"Document:\n{example.get('context', '')}", + user_prompt="Return the extracted JSON object now.", + response_schema={"name": "extraction", "schema": schema}, + **kwargs, + ) + + def postprocess(self, output: str) -> str: + """Strip markdown code fences if the model wrapped the JSON anyway.""" + text = output.strip() + if text.startswith("```"): + text = text.strip("`").strip() + if text.lower().startswith("json"): + text = text[4:].strip() + return text + + def normalize(self, text: str) -> str: + """No-op — the field_f1 metric parses and compares JSON directly.""" + return text diff --git a/src/tasks/rag_task.py b/src/tasks/rag_task.py index 7226b4e..6f59148 100644 --- a/src/tasks/rag_task.py +++ b/src/tasks/rag_task.py @@ -94,7 +94,13 @@ def build_prompt(self, example: Dict[str, Any], **kwargs) -> "GenerationInput": GenerationInput object. """ from models.base import GenerationInput - + + # model_mode is passed by evaluation/runner.py for every task (only + # SummarizationTask's ScaleDown-endpoint branch actually needs it); + # pop it here so it doesn't reach GenerationInput(), which has no + # such field and raises a TypeError on an unexpected kwarg. + kwargs.pop("model_mode", None) + question = example.get("question", "") context = example.get("context", "") diff --git a/src/tasks/retrieval_task.py b/src/tasks/retrieval_task.py index e447342..94e00a4 100644 --- a/src/tasks/retrieval_task.py +++ b/src/tasks/retrieval_task.py @@ -60,7 +60,10 @@ def build_prompt(self, example: Dict[str, Any], **kwargs) -> "GenerationInput": GenerationInput object with response_schema set. """ from models.base import GenerationInput - + + # See tasks/rag_task.py for why this pop is needed. + kwargs.pop("model_mode", None) + question = example.get("question", "") context = example.get("context", "")