Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow fixing sequence length to the model's context length #2

Merged
merged 2 commits into from
Dec 18, 2024
Merged

Allow fixing sequence length to the model's context length #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
75d8818
Add option to fix context length
epwalsh Dec 18, 2024
762dbea
clean up
epwalsh Dec 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

## Unreleased

### Added

- Allowed passing additional kwargs to the task through `build_task()`.
- Added the option to fix the context length of every batch to the model's context length.

## [v0.2.0](https://github.com/allenai/OLMo-in-loop-evals/releases/tag/v0.2.0) - 2024-10-29

### Added
Expand Down
62 changes: 58 additions & 4 deletions src/olmo_eval/tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import abc
import logging
import re
from typing import Any, Dict, List, Optional, Sequence, Union
from typing import Any, Dict, List, Optional, Sequence, Type, Union, cast

import datasets
import torch
Expand All @@ -26,6 +26,7 @@ def __init__(
dataset_path: str,
dataset_name: Union[str, Sequence[str], None] = None,
model_ctx_len: int = 2048,
fixed_ctx_len: bool = False,
split="validation",
metric_type=None, # Override default metric type
prompts: Optional[List[Optional[str]]] = None, # List of prompt variants to use
Expand All @@ -36,6 +37,7 @@ def __init__(
self.dataset_path = dataset_path
self.dataset_name = dataset_name
self.model_ctx_len = model_ctx_len
self.fixed_ctx_len = fixed_ctx_len
self.prompts = prompts or [None]
self.current_prompt: Optional[str] = None
if metric_type is not None:
Expand Down Expand Up @@ -213,7 +215,12 @@ def collate_fn(self, data):
cont_byte_lens.append(sample["cont_byte_len"])

queries.append(
torch.LongTensor(self.pad_tokens_until_max(sample["query"], max_len=max_query_len))
torch.LongTensor(
self.pad_tokens_until_max(
sample["query"],
max_len=self.model_ctx_len if self.fixed_ctx_len else max_query_len,
)
)
)
dc_queries.append(
torch.LongTensor(
Expand Down Expand Up @@ -298,11 +305,13 @@ def __init__(
tokenizer,
dataset_path="piqa",
dataset_name="plain_text",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -341,11 +350,13 @@ def __init__(
tokenizer,
dataset_path="hellaswag",
dataset_name=None,
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

@classmethod
Expand Down Expand Up @@ -406,12 +417,14 @@ def __init__(
tokenizer,
dataset_path="winogrande",
dataset_name="winogrande_xl",
**kwargs,
):
# all winogrande datasets have same val set
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def prep_examples(self):
Expand Down Expand Up @@ -508,11 +521,13 @@ def __init__(
tokenizer,
dataset_path="openbookqa",
dataset_name="main",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -548,11 +563,13 @@ def __init__(
tokenizer,
dataset_path="boolq",
dataset_name=None,
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -598,11 +615,13 @@ def __init__(
tokenizer,
dataset_path="sciq",
dataset_name=None,
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -645,11 +664,13 @@ def __init__(
tokenizer,
dataset_path: str = "ai2_arc",
dataset_name: Optional[str] = "ARC-Easy",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -685,11 +706,13 @@ def __init__(
tokenizer,
dataset_path="ai2_arc",
dataset_name="ARC-Challenge",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)


Expand Down Expand Up @@ -726,11 +749,13 @@ def __init__(
tokenizer,
dataset_path="allenai/basic_arithmetic",
dataset_name: Optional[str] = None,
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)


Expand All @@ -752,11 +777,13 @@ def __init__(
tokenizer,
dataset_path="tau/commonsense_qa",
dataset_name=None,
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)


Expand All @@ -777,11 +804,13 @@ def __init__(
tokenizer,
dataset_path="social_i_qa",
dataset_name=None,
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -830,11 +859,13 @@ def __init__(
tokenizer,
dataset_path="super_glue",
dataset_name="copa",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -877,11 +908,13 @@ def __init__(
tokenizer,
dataset_path="glue",
dataset_name="rte",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -921,11 +954,13 @@ def __init__(
tokenizer,
dataset_path="super_glue",
dataset_name="cb",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -968,11 +1003,13 @@ def __init__(
tokenizer,
dataset_path="glue",
dataset_name="mrpc",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

@classmethod
Expand Down Expand Up @@ -1038,11 +1075,13 @@ def __init__(
tokenizer,
dataset_path="glue",
dataset_name="sst2",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

@classmethod
Expand Down Expand Up @@ -1168,6 +1207,7 @@ def __init__(
prompt_variations=None,
mc_labels=False,
metric_type=None,
**kwargs,
):
dataset_names = []
# Collect the relevant categories
Expand Down Expand Up @@ -1203,6 +1243,7 @@ def __init__(
split=split,
prompts=prompts,
metric_type=metric_type,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -1288,11 +1329,13 @@ def __init__(
tokenizer,
dataset_path="trivia_qa",
dataset_name="rc.wikipedia.nocontext",
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand Down Expand Up @@ -1326,11 +1369,13 @@ def __init__(
tokenizer,
dataset_path="nq_open",
dataset_name=None,
**kwargs,
):
super().__init__(
tokenizer=tokenizer,
dataset_path=dataset_path,
dataset_name=dataset_name,
**kwargs,
)

def doc_to_text(self, doc):
Expand All @@ -1357,6 +1402,7 @@ def __init__(
dataset_path: str,
dataset_name: Union[str, Sequence[str], None] = None,
model_ctx_len: int = 2048,
fixed_ctx_len: bool = False,
split=None,
metric_type=None,
prompts: Optional[List[Optional[str]]] = None, # List of prompt variants to use
Expand All @@ -1367,6 +1413,7 @@ def __init__(
self.dataset_path = dataset_path
self.dataset_name = dataset_name
self.model_ctx_len = model_ctx_len
self.fixed_ctx_len = fixed_ctx_len
self.log_instances = 0 # Set to > 0 to log the first few instances as a sanity check

self.samples: List[Dict[str, Any]] = []
Expand Down Expand Up @@ -1847,9 +1894,16 @@ def list_tasks() -> List[str]:
return list(label_to_task_map.keys())


def build_task(label: str, tokenizer: Tokenizer) -> ICLMultiChoiceTaskDataset:
def build_task(
label: str,
tokenizer: Tokenizer,
model_ctx_len: int = 2048,
fixed_ctx_len: bool = False,
) -> ICLMultiChoiceTaskDataset:
task_class = label_to_task_map[label]
task_kwargs = {}
if isinstance(task_class, tuple):
task_class, task_kwargs = task_class
return task_class(tokenizer=tokenizer, **task_kwargs) # type: ignore
return cast(Type[ICLMultiChoiceTaskDataset], task_class)(
tokenizer=tokenizer, model_ctx_len=model_ctx_len, fixed_ctx_len=fixed_ctx_len, **task_kwargs
)
Loading