Update handler.py

handler.py

@@ -4,18 +4,22 @@ import subprocess
 import pkg_resources
 
 # Verify and print the transformers version
-transformers_version = pkg_resources.get_distribution("transformers").version
-print(f"Transformers version: {transformers_version}")
+try:
+    transformers_version = pkg_resources.get_distribution("transformers").version
+    print(f"Transformers version: {transformers_version}")
+except pkg_resources.DistributionNotFound:
+    transformers_version = None
+    print("Transformers not installed")
 
 # 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)
+    try:
+        subprocess.run('pip uninstall -y transformers', shell=True, check=True)
+        subprocess.run('pip install git+https://github.com/huggingface/transformers', shell=True, check=True)
+        transformers_version = pkg_resources.get_distribution("transformers").version
+        print(f"Updated Transformers version: {transformers_version}")
+    except subprocess.CalledProcessError as e:
+        print(f"Error occurred while updating transformers: {e}")
 
 class CustomModelHandler:
     def __init__(self, model_name_or_path: str):
@@ -30,8 +34,7 @@ class CustomModelHandler:
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_name_or_path, 
                 trust_remote_code=True,
-                torch_dtype="auto",
-                _attn_implementation="eager"
+                torch_dtype="auto"
             )
             self.model.to(self.device)
             print(f"Model loaded and moved to {self.device}")