[update] run.py

run.py

@@ -6,15 +6,18 @@ from langchain_community.vectorstores import Chroma
 from transformers import AutoModelForCausalLM
 
 import os
-
-# Ensure the TRANSFORMERS_CACHE directory exists and is writable
-
-cache_dir = os.environ['TRANSFORMERS_CACHE'] 
+os.environ['TRANSFORMERS_CACHE'] = '/code/model/cache/'
 
 model_kwargs = {'trust_remote_code': True}
 
+# embedding = HuggingFaceEmbeddings(
+#     model_name="nomic-ai/nomic-embed-text-v1.5",
+#     model_kwargs=model_kwargs
+# )
+
 db = Chroma(
     persist_directory="./chroma_db",
+    # embedding_function=embedding,
     collection_name='CVE'        
 )
 
@@ -28,15 +31,11 @@ Question: {question}
 
 prompt = ChatPromptTemplate.from_template(template)
 
-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)
+model = AutoModelForCausalLM.from_pretrained(
+    "zephyr-7b-beta.Q4_K_S.gguf",
+    model_type='mistral',
+    threads=3,
+)
 
 chain = (
     {"context": retriever, "question": RunnablePassthrough()}