Support ONNX export for causal LM sequence classifiers (#27450)

dwyatte · web-flow · commit 1394e08cf099 · 2023-11-16T18:56:34.000+09:00
support onnx for causal lm sequence classification
diff --git a/src/transformers/models/ctrl/modeling_ctrl.py b/src/transformers/models/ctrl/modeling_ctrl.py
@@ -796,7 +796,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/deprecated/open_llama/modeling_open_llama.py b/src/transformers/models/deprecated/open_llama/modeling_open_llama.py
@@ -924,7 +924,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/gpt2/modeling_gpt2.py b/src/transformers/models/gpt2/modeling_gpt2.py
@@ -1451,7 +1451,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py b/src/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py
@@ -1184,7 +1184,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/gpt_neo/modeling_gpt_neo.py b/src/transformers/models/gpt_neo/modeling_gpt_neo.py
@@ -1090,7 +1090,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/src/transformers/models/gpt_neox/modeling_gpt_neox.py
@@ -948,7 +948,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/gptj/modeling_gptj.py b/src/transformers/models/gptj/modeling_gptj.py
@@ -1001,7 +1001,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/llama/modeling_llama.py b/src/transformers/models/llama/modeling_llama.py
@@ -1204,7 +1204,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/mistral/modeling_mistral.py b/src/transformers/models/mistral/modeling_mistral.py
@@ -1174,7 +1174,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/openai/modeling_openai.py b/src/transformers/models/openai/modeling_openai.py
@@ -814,7 +814,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/opt/modeling_opt.py b/src/transformers/models/opt/modeling_opt.py
@@ -1030,7 +1030,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/persimmon/modeling_persimmon.py b/src/transformers/models/persimmon/modeling_persimmon.py
@@ -925,7 +925,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/phi/modeling_phi.py b/src/transformers/models/phi/modeling_phi.py
@@ -938,7 +938,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:
diff --git a/src/transformers/models/transfo_xl/modeling_transfo_xl.py b/src/transformers/models/transfo_xl/modeling_transfo_xl.py
@@ -1247,7 +1247,7 @@ def forward(
             sequence_lengths = -1
         else:
             if input_ids is not None:
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
                     logits.device
                 )
             else:

Original file line number	Diff line number	Diff line change
`@@ -796,7 +796,7 @@ def forward(`
`796`	`796`	`sequence_lengths = -1`
`797`	`797`	`else:`
`798`	`798`	`if input_ids is not None:`
`799`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`799`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`800`	`800`	`logits.device`
`801`	`801`	`)`
`802`	`802`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -924,7 +924,7 @@ def forward(`
`924`	`924`	`sequence_lengths = -1`
`925`	`925`	`else:`
`926`	`926`	`if input_ids is not None:`
`927`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`927`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`928`	`928`	`logits.device`
`929`	`929`	`)`
`930`	`930`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -1451,7 +1451,7 @@ def forward(`
`1451`	`1451`	`sequence_lengths = -1`
`1452`	`1452`	`else:`
`1453`	`1453`	`if input_ids is not None:`
`1454`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`1454`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`1455`	`1455`	`logits.device`
`1456`	`1456`	`)`
`1457`	`1457`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -1184,7 +1184,7 @@ def forward(`
`1184`	`1184`	`sequence_lengths = -1`
`1185`	`1185`	`else:`
`1186`	`1186`	`if input_ids is not None:`
`1187`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`1187`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`1188`	`1188`	`logits.device`
`1189`	`1189`	`)`
`1190`	`1190`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -1090,7 +1090,7 @@ def forward(`
`1090`	`1090`	`sequence_lengths = -1`
`1091`	`1091`	`else:`
`1092`	`1092`	`if input_ids is not None:`
`1093`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`1093`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`1094`	`1094`	`logits.device`
`1095`	`1095`	`)`
`1096`	`1096`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -948,7 +948,7 @@ def forward(`
`948`	`948`	`sequence_lengths = -1`
`949`	`949`	`else:`
`950`	`950`	`if input_ids is not None:`
`951`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`951`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`952`	`952`	`logits.device`
`953`	`953`	`)`
`954`	`954`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -1001,7 +1001,7 @@ def forward(`
`1001`	`1001`	`sequence_lengths = -1`
`1002`	`1002`	`else:`
`1003`	`1003`	`if input_ids is not None:`
`1004`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`1004`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`1005`	`1005`	`logits.device`
`1006`	`1006`	`)`
`1007`	`1007`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -1204,7 +1204,7 @@ def forward(`
`1204`	`1204`	`sequence_lengths = -1`
`1205`	`1205`	`else:`
`1206`	`1206`	`if input_ids is not None:`
`1207`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`1207`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`1208`	`1208`	`logits.device`
`1209`	`1209`	`)`
`1210`	`1210`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -1174,7 +1174,7 @@ def forward(`
`1174`	`1174`	`sequence_lengths = -1`
`1175`	`1175`	`else:`
`1176`	`1176`	`if input_ids is not None:`
`1177`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`1177`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`1178`	`1178`	`logits.device`
`1179`	`1179`	`)`
`1180`	`1180`	`else:`
Original file line number	Diff line number	Diff line change
`@@ -814,7 +814,7 @@ def forward(`
`814`	`814`	`sequence_lengths = -1`
`815`	`815`	`else:`
`816`	`816`	`if input_ids is not None:`
`817`		`- sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(`
	`817`	`+ sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(`
`818`	`818`	`logits.device`
`819`	`819`	`)`
`820`	`820`	`else:`