Update handler.py

handler.py

@@ -1,8 +1,20 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import subprocess
+import pkg_resources
 
-# # Install flash-attn
+# Verify and print the transformers version
+transformers_version = pkg_resources.get_distribution("transformers").version
+print(f"Transformers version: {transformers_version}")
+
+# Update transformers to the development version if necessary
+if transformers_version != "4.40.2":
+    subprocess.run('pip uninstall -y transformers', shell=True)
+    subprocess.run('pip install git+https://github.com/huggingface/transformers', shell=True)
+    transformers_version = pkg_resources.get_distribution("transformers").version
+    print(f"Updated Transformers version: {transformers_version}")
+
+# Install flash-attn if needed
 # subprocess.run('pip install flash-attn', shell=True)
 
 class CustomModelHandler: