huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-09 11:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(server): remove position_ids from galactica forward

Browse Source
This commit is contained in:
OlivierDehaene 2023-02-20 19:04:49 +01:00
parent 17bc841b1b
commit 0363b67d75
1 changed files with 15 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
server/text_generation/models/galactica.py
View File

@ -2,7 +2,7 @@ import re
import torch
import torch.distributed
from typing import List, Optional, Type
from typing import List, Optional, Type, Tuple
from accelerate import init_empty_weights
from safetensors import safe_open
@ -145,6 +145,20 @@ class Galactica(CausalLM):
generated_ids, skip_special_tokens=False, cleanup_tokenization_spaces=False
)
def forward(
self, input_ids, attention_mask, position_ids, past_key_values: Optional = None
) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]:
"""Overwrite forward to ignore position_ids"""
# Model Forward
outputs = self.model.forward(
input_ids=input_ids,
attention_mask=attention_mask,
past_key_values=past_key_values,
use_cache=True,
)
return outputs.logits, outputs.past_key_values
class GalacticaSharded(Galactica):
def __init__(
@ -322,7 +336,6 @@ class GalacticaSharded(Galactica):
outputs = self.model.forward(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
use_cache=True,
)