Update m5.py

m5.py

@@ -4,9 +4,21 @@ from collections import defaultdict
 from accelerate import Accelerator
 from transformers import AutoTokenizer, AutoModel
 from sklearn.metrics.pairwise import cosine_similarity
+import termcolor
+
+# Set the cache directory path
+cache_dir = '/app/cache'
+
+# Create the directory if it doesn't exist
+if not os.path.exists(cache_dir):
+    os.makedirs(cache_dir)
+
+# Set the environment variable
+os.environ['TRANSFORMERS_CACHE'] = cache_dir
+
+# Verify the environment variable is set
+print(f"TRANSFORMERS_CACHE is set to: {os.environ['TRANSFORMERS_CACHE']}")
 
-# Set the TRANSFORMERS_CACHE environment variable to a writable directory
-os.environ['TRANSFORMERS_CACHE'] = '/app/cache'
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 class DM(nn.Module):