[update] run.py test

run.py

@@ -6,19 +6,18 @@ from langchain_community.vectorstores import Chroma
 from transformers import AutoModelForCausalLM
 
 import os
-os.environ['TRANSFORMERS_CACHE'] = '/code/model/cache/'
 
+# Ensure the TRANSFORMERS_CACHE directory exists and is writable
+cache_dir = '/code/model/cache/'
+if not os.path.exists(cache_dir):
+    os.makedirs(cache_dir)
 
-model_kwargs = {'trust_remote_code': True}
+os.environ['TRANSFORMERS_CACHE'] = cache_dir
 
-# embedding = HuggingFaceEmbeddings(
-#     model_name="nomic-ai/nomic-embed-text-v1.5",
-#     model_kwargs=model_kwargs
-# )
+model_kwargs = {'trust_remote_code': True}
 
 db = Chroma(
     persist_directory="./chroma_db",
-    # embedding_function=embedding,
     collection_name='CVE'        
 )
 
@@ -32,11 +31,15 @@ Question: {question}
 
 prompt = ChatPromptTemplate.from_template(template)
 
-model = AutoModelForCausalLM.from_pretrained(
-    "zephyr-7b-beta.Q4_K_S.gguf",
-    model_type='mistral',
-    threads=3,
-)
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        "zephyr-7b-beta.Q4_K_S.gguf",
+        model_type='mistral',
+        threads=3,
+    )
+except Exception as e:
+    print(f"Failed to load model: {e}")
+    exit(1)
 
 chain = (
     {"context": retriever, "question": RunnablePassthrough()}