Skip to content

Fix config initialization, serialization, pipeline compatibility, and Apple MPS GPU compatibility #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gokulrejith wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix config initialization, serialization, pipeline compatibility, and Apple MPS GPU compatibility #29

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 24 additions & 10 deletions fastfit/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ def forward(self, features, labels=None, mask=None, sim=None):
Returns:
A loss scalar.
"""
device = torch.device("cuda") if features.is_cuda else torch.device("cpu")
device = features.device

if len(features.shape) < 3:
raise ValueError(
Expand Down Expand Up @@ -256,11 +256,14 @@ def __init__(
self.length_norm_at_inference = length_norm_at_inference

assert inference_direction in ["query", "doc", "both"]
assert "encoder" in kwargs, "Config has to be initialized with encoder config"
encoder_config = kwargs.pop("encoder")
encoder_model_type = encoder_config.pop("model_type")
self.encoder = AutoConfig.for_model(encoder_model_type, **encoder_config)
self.hidden_size = self.encoder.hidden_size
if "encoder" in kwargs:
encoder_config = kwargs.pop("encoder")
encoder_model_type = encoder_config.pop("model_type")
self.encoder = AutoConfig.for_model(encoder_model_type, **encoder_config)
self.hidden_size = self.encoder.hidden_size
else:
self.encoder = None
self.hidden_size = None

@classmethod
def from_encoder_config(
Expand All @@ -270,7 +273,8 @@ def from_encoder_config(

def to_dict(self):
output = copy.deepcopy(self.__dict__)
output["encoder"] = self.encoder.to_dict()
if self.encoder is not None:
output["encoder"] = self.encoder.to_dict()
output["model_type"] = self.__class__.model_type
return output

Expand Down Expand Up @@ -544,12 +548,22 @@ def reduce(self, encodings, input_ids, do_mask_zeros=False):

def forward(
self,
query_input_ids,
query_attention_mask,
input_ids=None,
attention_mask=None,
query_input_ids=None,
query_attention_mask=None,
doc_input_ids=None,
doc_attention_mask=None,
labels=None,
):
if query_input_ids is None and input_ids is not None:
query_input_ids = input_ids
if query_attention_mask is None and attention_mask is not None:
query_attention_mask = attention_mask
if labels is None:
scores = self.inference_forward(query_input_ids, query_attention_mask)
return SequenceClassifierOutput(logits=scores)

scores = None
if not self.training:
scores = self.inference_forward(query_input_ids, query_attention_mask)
Expand Down Expand Up @@ -625,7 +639,7 @@ def sim_loss(
a, b = Q[:bs, :, :], Q[bs:, :, :]

else:
lbls = Variable(labels.type(torch.DoubleTensor), requires_grad=True)
lbls = Variable(labels.type(torch.FloatTensor), requires_grad=True).to(self.device)
if self.config.rep_tokens == "all":
sim_mat = self.query_doc_similarity_matrix(
Q, D, query_attention_mask, doc_attention_mask,
Expand Down