From 70cd8ecaf67b8908f37569edabcb9eb697a2aada Mon Sep 17 00:00:00 2001 From: Vy Hong Date: Mon, 1 Apr 2024 17:37:07 -0700 Subject: [PATCH] Declare fields in PatchscopeBase Also moves the Context classes to patchscope base to avoid circular dependency --- obvs/patchscope.py | 96 +---------------------------------- obvs/patchscope_base.py | 107 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 108 insertions(+), 95 deletions(-) diff --git a/obvs/patchscope.py b/obvs/patchscope.py index 51a803d..890f4bc 100644 --- a/obvs/patchscope.py +++ b/obvs/patchscope.py @@ -32,106 +32,14 @@ from __future__ import annotations -from collections.abc import Callable, Sequence -from dataclasses import dataclass, field +from collections.abc import Sequence import torch from nnsight import LanguageModel from tqdm import tqdm from obvs.logging import logger -from obvs.patchscope_base import PatchscopeBase - - -@dataclass -class SourceContext: - """ - Source context for the patchscope - """ - _prompt: str | torch.Tensor = field(init=False, repr=False, default="<|endoftext|>") - _text_prompt: str = field(init=False, repr=False) - _soft_prompt: torch.Tensor | None = field(init=False, repr=False) - - prompt: str | torch.Tensor - position: Sequence[int] | None = None - layer: int = -1 - model_name: str = "gpt2" - device: str = "cuda" if torch.cuda.is_available() else "cpu" - - # This overrides the `prompt` field - # See https://florimond.dev/en/posts/2018/10/reconciling-dataclasses-and-properties-in-python - @property - def prompt(self) -> str | torch.Tensor: - """ - The prompt - """ - return self._prompt - - @prompt.setter - def prompt(self, value: str | torch.Tensor | None): - if isinstance(value, property): - # initial value not specified, use default - value = SourceContext._prompt - - if value is None: - value = "<|endoftext|>" - - if isinstance(value, torch.Tensor) and value.dim() != 2: - raise ValueError(f"Soft prompt must have shape [pos, dmodel]. prompt.shape = {value.shape}") - - self._prompt = value - if isinstance(value, torch.Tensor): - self._text_prompt = " ".join("_" * value.shape[0]) - self._soft_prompt = value - else: - self._text_prompt = value - self._soft_prompt = None - - @property - def text_prompt(self) -> str: - """ - The text prompt input or generated from soft prompt - """ - return self._text_prompt - - @property - def soft_prompt(self) -> torch.Tensor | None: - """ - The soft prompt input or None - """ - return self._soft_prompt - - -@dataclass -class TargetContext(SourceContext): - """ - Target context for the patchscope - Parameters identical to the source context, with the addition of - a mapping function and max_new_tokens to control generation length - """ - - mapping_function: Callable[[torch.Tensor], torch.Tensor] = lambda x: x - max_new_tokens: int = 10 - - @staticmethod - def from_source( - source: SourceContext, - mapping_function: Callable[[torch.Tensor], torch.Tensor] | None = None, - max_new_tokens: int = 10, - ) -> TargetContext: - """ - Construct a target context from the source context - """ - - return TargetContext( - prompt=source.prompt, - position=source.position, - model_name=source.model_name, - layer=source.layer, - mapping_function=mapping_function or (lambda x: x), - max_new_tokens=max_new_tokens, - device=source.device, - ) +from obvs.patchscope_base import PatchscopeBase, SourceContext, TargetContext class ModelLoader: diff --git a/obvs/patchscope_base.py b/obvs/patchscope_base.py index 84698b8..fabf17e 100644 --- a/obvs/patchscope_base.py +++ b/obvs/patchscope_base.py @@ -1,9 +1,103 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import Sequence +from dataclasses import dataclass, field +from typing import Callable, Sequence import torch +from nnsight import LanguageModel +from transformers import PreTrainedTokenizer + + +@dataclass +class SourceContext: + """ + Source context for the patchscope + """ + _prompt: str | torch.Tensor = field(init=False, repr=False, default="<|endoftext|>") + _text_prompt: str = field(init=False, repr=False) + _soft_prompt: torch.Tensor | None = field(init=False, repr=False) + + prompt: str | torch.Tensor + position: Sequence[int] | None = None + layer: int = -1 + model_name: str = "gpt2" + device: str = "cuda" if torch.cuda.is_available() else "cpu" + + # This overrides the `prompt` field + # See https://florimond.dev/en/posts/2018/10/reconciling-dataclasses-and-properties-in-python + @property + def prompt(self) -> str | torch.Tensor: + """ + The prompt + """ + return self._prompt + + @prompt.setter + def prompt(self, value: str | torch.Tensor | None): + if isinstance(value, property): + # initial value not specified, use default + value = SourceContext._prompt + + if value is None: + value = "<|endoftext|>" + + if isinstance(value, torch.Tensor) and value.dim() != 2: + raise ValueError(f"Soft prompt must have shape [pos, dmodel]. prompt.shape = {value.shape}") + + self._prompt = value + if isinstance(value, torch.Tensor): + self._text_prompt = " ".join("_" * value.shape[0]) + self._soft_prompt = value + else: + self._text_prompt = value + self._soft_prompt = None + + @property + def text_prompt(self) -> str: + """ + The text prompt input or generated from soft prompt + """ + return self._text_prompt + + @property + def soft_prompt(self) -> torch.Tensor | None: + """ + The soft prompt input or None + """ + return self._soft_prompt + + +@dataclass +class TargetContext(SourceContext): + """ + Target context for the patchscope + Parameters identical to the source context, with the addition of + a mapping function and max_new_tokens to control generation length + """ + + mapping_function: Callable[[torch.Tensor], torch.Tensor] = lambda x: x + max_new_tokens: int = 10 + + @staticmethod + def from_source( + source: SourceContext, + mapping_function: Callable[[torch.Tensor], torch.Tensor] | None = None, + max_new_tokens: int = 10, + ) -> TargetContext: + """ + Construct a target context from the source context + """ + + return TargetContext( + prompt=source.prompt, + device=source.device, + position=source.position, + model_name=source.model_name, + layer=source.layer, + mapping_function=mapping_function or (lambda x: x), + max_new_tokens=max_new_tokens, + ) class PatchscopeBase(ABC): @@ -11,6 +105,17 @@ class PatchscopeBase(ABC): A base class with lots of helper functions """ + source: SourceContext + target: TargetContext + source_model: LanguageModel + target_model: LanguageModel + tokenizer: PreTrainedTokenizer + + MODEL_TARGET: str + LAYER_TARGET: str + + _target_outputs: list[torch.Tensor] + def get_model_specifics(self, model_name): """ Get the model specific attributes.