Update optimum_encoder.py

optimum_encoder.py

@@ -60,11 +60,11 @@ class OptimumEncoder(BaseEncoder):
             **self.tokenizer_kwargs,
         )
         
-        provider_options = {
-            "trt_engine_cache_enable": True,
-            "trt_engine_cache_path": os.getenv('HF_HOME'),
-            "trt_fp16_enable": True
-        }
+        #provider_options = {
+        #    "trt_engine_cache_enable": True,
+        #    "trt_engine_cache_path": os.getenv('HF_HOME'),
+        #    "trt_fp16_enable": True
+        #}
 
         session_options = ort.SessionOptions()
         session_options.log_severity_level = 0
@@ -73,33 +73,34 @@ class OptimumEncoder(BaseEncoder):
             model_id=self.name, 
             file_name='model_fp16.onnx',
             subfolder='onnx',
-            provider='TensorrtExecutionProvider',
-            provider_options=provider_options,
+            provider='CUDAExecutionProvider',
+            use_io_binding=True,
+            #provider_options=provider_options,
             session_options=session_options,
             **self.model_kwargs
         )
 
-        print("Building engine for a short sequence...")
-        short_text = ["short"]
-        short_encoded_input = tokenizer(
-            short_text, padding=True, truncation=True, return_tensors="pt"
-        ).to(self.device)
-        short_output = ort_model(**short_encoded_input)
+        # print("Building engine for a short sequence...")
+        # short_text = ["short"]
+        # short_encoded_input = tokenizer(
+        #     short_text, padding=True, truncation=True, return_tensors="pt"
+        # ).to(self.device)
+        # short_output = ort_model(**short_encoded_input)
         
-        print("Building engine for a long sequence...")
-        long_text = ["a very long input just for demo purpose, this is very long" * 10]
-        long_encoded_input = tokenizer(
-            long_text, padding=True, truncation=True, return_tensors="pt"
-        ).to(self.device)
-        long_output = ort_model(**long_encoded_input)
-
-        text = ["Replace me by any text you'd like."]
-        encoded_input = tokenizer(
-            text, padding=True, truncation=True, return_tensors="pt"
-        ).to(self.device)
+        # print("Building engine for a long sequence...")
+        # long_text = ["a very long input just for demo purpose, this is very long" * 10]
+        # long_encoded_input = tokenizer(
+        #     long_text, padding=True, truncation=True, return_tensors="pt"
+        # ).to(self.device)
+        # long_output = ort_model(**long_encoded_input)
+
+        # text = ["Replace me by any text you'd like."]
+        # encoded_input = tokenizer(
+        #     text, padding=True, truncation=True, return_tensors="pt"
+        # ).to(self.device)
         
-        for i in range(3):
-            output = ort_model(**encoded_input)
+        # for i in range(3):
+        #     output = ort_model(**encoded_input)
 
         return tokenizer, ort_model