mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Make it possible to evaluate exllama perplexity (#3138)

Browse source
This commit is contained in:
oobabooga 2023-07-16 01:52:55 -03:00 • committed by GitHub
parent b284f2407d
commit 94dfcec237
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 15 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

20
modules/exllama_hf.py
View file

@ -29,6 +29,7 @@ class ExllamaHF(PreTrainedModel):
super().__init__(PretrainedConfig()) super().__init__(PretrainedConfig())
self.ex_config = config self.ex_config = config
self.ex_model = ExLlama(self.ex_config) self.ex_model = ExLlama(self.ex_config)
self.ex_cache = ExLlamaCache(self.ex_model)
self.generation_config = GenerationConfig() self.generation_config = GenerationConfig()
self.lora = None self.lora = None
@ -52,11 +53,20 @@ class ExllamaHF(PreTrainedModel):
labels = kwargs.get('labels', None) labels = kwargs.get('labels', None)
seq = kwargs['input_ids'][0].tolist() seq = kwargs['input_ids'][0].tolist()
cache = kwargs['past_key_values'] if 'past_key_values' in kwargs else None cache = kwargs['past_key_values'] if 'past_key_values' in kwargs else None
if cache is None:
cache = ExLlamaCache(self.ex_model)
self.ex_model.forward(torch.tensor([seq[:-1]], dtype=torch.long), cache, preprocess_only=True, lora=self.lora)
logits = self.ex_model.forward(torch.tensor([seq[-1:]], dtype=torch.long), cache, lora=self.lora).to(kwargs['input_ids'].device) if labels is None:
if cache is None:
self.ex_cache.current_seq_len = 0
cache = self.ex_cache
self.ex_model.forward(torch.tensor([seq[:-1]], dtype=torch.long), cache, preprocess_only=True, lora=self.lora)
logits = self.ex_model.forward(torch.tensor([seq[-1:]], dtype=torch.long), cache, lora=self.lora).to(kwargs['input_ids'].device)
else:
if cache is None:
self.ex_cache.current_seq_len = 0
cache = self.ex_cache
logits = self.ex_model.forward(torch.tensor([seq], dtype=torch.long), cache, last_id_only=False, lora=self.lora)
loss = None loss = None
if labels is not None: if labels is not None:
@ -71,7 +81,7 @@ class ExllamaHF(PreTrainedModel):
shift_labels = shift_labels.to(shift_logits.device) shift_labels = shift_labels.to(shift_logits.device)
loss = loss_fct(shift_logits, shift_labels) loss = loss_fct(shift_logits, shift_labels)
return CausalLMOutputWithPast(logits=logits, past_key_values=cache if use_cache else None) return CausalLMOutputWithPast(logits=logits, past_key_values=cache if use_cache else None, loss=loss)
@classmethod @classmethod
def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs): def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):