[Fix #16] Set transformers package version and add actions

.github/workflows/docker-image.yml

@@ -24,6 +24,14 @@ jobs:
         run: |
           docker run --rm onxxim-test echo "Docker build successful!"
 
+      - name: Test generating onnx file - GPT2
+        run: |
+          docker run --rm onxxim-test python3 /workspace/ONNXim/scripts/generate_transformer_onnx.py --model gpt2
+
+      - name: Test generating onnx file - BERT
+        run: |
+          docker run --rm onxxim-test python3 /workspace/ONNXim/scripts/generate_transformer_onnx.py --model bert
+
       - name: Cache Docker layers
         uses: actions/cache@v3
         with:

Dockerfile

@@ -25,7 +25,7 @@ RUN wget https://github.com/Kitware/CMake/releases/download/v3.22.0/cmake-3.22.0
     make install
 
 # Install specific Python packages with pip
-RUN pip3 install conan==1.57.0 onnx onnxruntime torch torchvision optimum
+RUN pip3 install conan==1.57.0 transformers==4.40.1 onnx onnxruntime torch torchvision optimum
 
 # Copy your project files into the image
 COPY ./ ONNXim

scripts/generate_transformer_onnx.py

@@ -37,6 +37,13 @@
 pathlib.Path(f'{HOME}/models/{args.model}/').mkdir(parents=True, exist_ok=True)
 pathlib.Path(f"{HOME}/model_lists").mkdir(parents=True, exist_ok=True)
 
+# Check attention node
+nodes = optimized_model.graph().node
+node_types = [node.op_type for node in nodes]
+if "Attention" not in node_types:
+    print("Opimizing model failed...")
+    exit(1)
+
 # Save optimized onnx file
 optimized_model.save_model_to_file(f'{HOME}/models/{args.model}/{args.model}.onnx', use_external_data_format=True)