gpt2.py 4.88 KB
Newer Older
Jason Phang's avatar
gpt3  
Jason Phang committed
1
import torch
Xingjian Shi's avatar
Xingjian Shi committed
2
import transformers
3
from typing import Optional
4
from lm_eval.base import BaseLM
Jason Phang's avatar
gpt3  
Jason Phang committed
5

6
class HFLM(BaseLM):
Fabrizio Milo's avatar
Fabrizio Milo committed
7
8
9
10
11
    def __init__(
        self,
        device="cuda",
        pretrained="gpt2",
        revision="main",
Xingjian Shi's avatar
Xingjian Shi committed
12
        low_cpu_mem_usage=None,
Fabrizio Milo's avatar
Fabrizio Milo committed
13
14
15
        subfolder=None,
        tokenizer=None,
        batch_size=1,
16
17
        load_in_8bit: Optional[bool] = False,
        trust_remote_code: Optional[bool] = False,
Fabrizio Milo's avatar
Fabrizio Milo committed
18
    ):
19
20
21
22
        super().__init__()

        assert isinstance(device, str)
        assert isinstance(pretrained, str)
23
        assert isinstance(batch_size, (int,str))
24

25
26
        device_list = set(["cuda", "cpu"] + [f'cuda:{i}' for i in range(torch.cuda.device_count())])
        if device and device in device_list:
27
            self._device = torch.device(device)
28
            print(f"Using device '{device}'")
29
        else:
Fabrizio Milo's avatar
Fabrizio Milo committed
30
            print("Device not specified")
31
            print(f"Cuda Available? {torch.cuda.is_available()}")
Fabrizio Milo's avatar
Fabrizio Milo committed
32
33
34
35
36
            self._device = (
                torch.device("cuda")
                if torch.cuda.is_available()
                else torch.device("cpu")
            )
37
38

        # TODO: update this to be less of a hack once subfolder is fixed in HF
39
40
        revision = revision + ("/" + subfolder if subfolder is not None else "")

41
        self.gpt2 = transformers.AutoModelForCausalLM.from_pretrained(
42
43
44
45
46
            pretrained,
            load_in_8bit=load_in_8bit,
            low_cpu_mem_usage=low_cpu_mem_usage,
            revision=revision,
            trust_remote_code=trust_remote_code,
47
        ).to(self.device)
48
49
        self.gpt2.eval()

50
        self.tokenizer = transformers.AutoTokenizer.from_pretrained(
kabbi159's avatar
kabbi159 committed
51
            pretrained if tokenizer is None else tokenizer,
cardy20's avatar
cardy20 committed
52
<<<<<<< HEAD
kabbi159's avatar
kabbi159 committed
53
            revision=revision + ("/" + subfolder if subfolder is not None else ""))
54

soqeue1's avatar
soqeue1 committed
55
56
57
58
        # assert isinstance(self.tokenizer, (
        #     transformers.GPT2Tokenizer, transformers.GPT2TokenizerFast,
        #     transformers.T5Tokenizer, transformers.T5TokenizerFast,
        # )), "this tokenizer has not been checked for compatibility yet!"
59
60
61

        self.vocab_size = self.tokenizer.vocab_size

soqeue1's avatar
soqeue1 committed
62
63
64
        # if isinstance(self.tokenizer, (transformers.GPT2Tokenizer, transformers.GPT2TokenizerFast)):
        #     assert self.tokenizer.encode('hello\n\nhello') == [31373, 198, 198, 31373], \
        #         self.tokenizer.encode('hello\n\nhello')
65
66

        # multithreading and batching
67
        self.batch_size_per_gpu = batch_size  # todo: adaptive batch size
68
69

        # TODO: fix multi-gpu
70
        # gpus = torch.cuda.device_count()
71
72
        # if gpus > 1:
        #     self.gpt2 = nn.DataParallel(self.gpt2)
cardy20's avatar
cardy20 committed
73
=======
74
            revision=revision,
75
            trust_remote_code=trust_remote_code,
Fabrizio Milo's avatar
Fabrizio Milo committed
76
        )
77

78
        self.vocab_size = self.tokenizer.vocab_size
79

Fabrizio Milo's avatar
Fabrizio Milo committed
80
81
82
83
84
85
86
87
88
        if isinstance(
            self.tokenizer, (transformers.GPT2Tokenizer, transformers.GPT2TokenizerFast)
        ):
            assert self.tokenizer.encode("hello\n\nhello") == [
                31373,
                198,
                198,
                31373,
            ], self.tokenizer.encode("hello\n\nhello")
Leo Gao's avatar
Leo Gao committed
89

90
91
92
93
94
        # setup for automatic batch size detection
        if batch_size == 'auto': 
            self.batch_size_per_gpu = batch_size
        else:
            self.batch_size_per_gpu = int(batch_size) 
cardy20's avatar
cardy20 committed
95
>>>>>>> 0542d35d5e56768dd9041ef9b88b90256970d843
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

    @property
    def eot_token_id(self):
        # we use EOT because end of *text* is more accurate for what we're doing than end of *sentence*
        return self.tokenizer.eos_token_id

    @property
    def max_length(self):
        try:
            return self.gpt2.config.n_ctx
        except AttributeError:
            # gptneoconfig doesn't have n_ctx apparently
            return self.gpt2.config.max_position_embeddings

    @property
    def max_gen_toks(self):
        return 256

    @property
    def batch_size(self):
        # TODO: fix multi-gpu
        return self.batch_size_per_gpu  # * gpus

    @property
    def device(self):
        # TODO: fix multi-gpu
        return self._device

124
125
    def tok_encode(self, string: str):
        return self.tokenizer.encode(string, add_special_tokens=False)
soqeue1's avatar
soqeue1 committed
126

127
128
129
130
131
132
133
134
135
    def tok_decode(self, tokens):
        return self.tokenizer.decode(tokens)

    def _model_call(self, inps):
        """
        inps: a torch tensor of shape [batch, sequence]
        the size of sequence may vary from call to call

        returns: a torch tensor of shape [batch, sequence, vocab] with the
136
        logits returned from the model
137
138
        """
        with torch.no_grad():
Jiwung Hyun's avatar
Jiwung Hyun committed
139
            return self.gpt2(inps)[0]
soqeue1's avatar
soqeue1 committed
140

141
    def _model_generate(self, context, max_length, eos_token_id):
142
143
144
145
        generation_kwargs = {'do_sample': False, 'max_length': max_length}
        if eos_token_id is not None:
            generation_kwargs['eos_token_id'] = eos_token_id
        return self.gpt2.generate(context, **generation_kwargs)
146
147


148
149
# for backwards compatibility
GPT2LM = HFLM