"references/vscode:/vscode.git/clone" did not exist on "9329257beeb51f7855cc765df15b0be21e361133"
Unverified Commit 78063c05 authored by OlivierDehaene's avatar OlivierDehaene Committed by GitHub
Browse files

fix(server): remove position_ids from galactica forward (#82)

closes #80
parent 17bc841b
......@@ -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,
)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment