fix nlp device kwargs (#2632)

* fix_device_stuff * reqs
deepchecks · Jul 16, 2023 · 6017cb9 · 6017cb9
1 parent 78eebd2
commit 6017cb9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/deepchecks/nlp/utils/text_properties.py b/deepchecks/nlp/utils/text_properties.py
@@ -189,14 +189,13 @@ def get_transformer_model(
         model_path = models_storage / 'onnx' / model_name
 
         if model_path.exists():
-            return onnx.ORTModelForSequenceClassification.from_pretrained(model_path, device_map=device)
+            return onnx.ORTModelForSequenceClassification.from_pretrained(model_path).to(device or -1)
 
         model = onnx.ORTModelForSequenceClassification.from_pretrained(
             model_name,
             export=True,
             cache_dir=models_storage,
-            device_map=device
-        )
+        ).to(device or -1)
         # NOTE:
         # 'optimum', after exporting/converting a model to the ONNX format,
         # does not store it onto disk we need to save it now to not reconvert
@@ -207,7 +206,7 @@ def get_transformer_model(
     model_path = models_storage / 'onnx' / 'quantized' / model_name
 
     if model_path.exists():
-        return onnx.ORTModelForSequenceClassification.from_pretrained(model_path, device_map=device)
+        return onnx.ORTModelForSequenceClassification.from_pretrained(model_path).to(device or -1)
 
     not_quantized_model = get_transformer_model(
         property_name,
@@ -217,7 +216,7 @@ def get_transformer_model(
         models_storage=models_storage
     )
 
-    quantizer = onnx.ORTQuantizer.from_pretrained(not_quantized_model, device_map=device)
+    quantizer = onnx.ORTQuantizer.from_pretrained(not_quantized_model).to(device or -1)
 
     quantizer.quantize(
         save_dir=model_path,
@@ -227,7 +226,7 @@ def get_transformer_model(
             per_channel=False
         )
     )
-    return onnx.ORTModelForSequenceClassification.from_pretrained(model_path, device_map=device)
+    return onnx.ORTModelForSequenceClassification.from_pretrained(model_path).to(device or -1)
 
 
 def get_transformer_pipeline(

diff --git a/requirements/nlp-prop-requirements.txt b/requirements/nlp-prop-requirements.txt
@@ -1,2 +1,2 @@
-optimum[onnxruntime]>=1.7.0
+optimum[onnxruntime]>=1.8.8
 fasttext>=0.8.0