Update app.py

app.py

@@ -10,6 +10,24 @@ from transformers import T5ForConditionalGeneration, T5Tokenizer
 import os
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+
+def install_apex():
+    # Install Apex in editable mode from the specified GitHub repository
+    cmd = [
+        'pip', 'install', '--no-cache-dir', '--no-build-isolation',
+        '--config-settings', 'build-option=--cpp_ext', '--config-settings',
+        'build-option=--cuda_ext', '-e', 'git+https://github.com/NVIDIA/apex.git'
+    ]
+    subprocess.run(cmd, check=True)
+
+try:
+    import apex
+except ModuleNotFoundError:
+    print("Apex not found, installing...")
+    install_apex()
+    # Try to import Apex again after installation
+    import apex
     
 
 def download_t5_model(model_id, save_directory):