Skip to content

Commit

Permalink
Pin library name to transformers for feature extraction (#804)
Browse files Browse the repository at this point in the history
  • Loading branch information
IlyasMoutawwakil authored and echarlaix committed Jul 8, 2024
1 parent a2870c8 commit 75357d3
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions optimum/intel/ipex/modeling_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,7 @@ def _from_pretrained(
model = TasksManager.get_model_from_task(
task,
model_id,
library_name="transformers",
trust_remote_code=trust_remote_code,
torch_dtype=torch_dtype,
_commit_hash=commit_hash,
Expand Down Expand Up @@ -267,6 +268,7 @@ def forward(
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
token_type_ids: torch.Tensor = None,
position_ids: torch.Tensor = None,
**kwargs,
):
inputs = {
Expand All @@ -277,6 +279,9 @@ def forward(
if "token_type_ids" in self.input_names:
inputs["token_type_ids"] = token_type_ids

if "position_ids" in self.input_names:
inputs["position_ids"] = position_ids

outputs = self._call_model(**inputs)
if isinstance(outputs, dict):
model_output = ModelOutput(**outputs)
Expand Down

0 comments on commit 75357d3

Please sign in to comment.