Update app_v4.py

app_v4.py

@@ -3,6 +3,7 @@ from transformers import AutoTokenizer
 from auto_gptq import AutoGPTQForCausalLM
 import torch
 import subprocess
+import traceback
 
 # Function to get memory info
 def get_gpu_memory():
@@ -38,11 +39,11 @@ try:
         model_basename="Jackson2-4bit-128g-GPTQ",
         use_safetensors=True,
         device=device,
-        max_memory={0: "10GIB"}
+        max_memory={0: "15GIB"}
     )
 except RuntimeError as e:
     if 'CUDA out of memory' in str(e):
-        st.error("CUDA out of memory. Try reducing the model size or input length.")
+        st.error("CUDA out of memory before toking generation. Try reducing the model size or input length.")
         st.stop()
     else:
         raise e
@@ -60,5 +61,11 @@ if st.button("Generate the prompt"):
     except RuntimeError as e:
         if 'CUDA out of memory' in str(e):
             st.error("CUDA out of memory during generation. Try reducing the input length.")
+            # Log the detailed error message
+            with open('error_log.txt', 'a') as f:
+                f.write(traceback.format_exc())
         else:
+            # Log the error and re-raise it
+            with open('error_log.txt', 'a') as f:
+                f.write(traceback.format_exc())
             raise e