Update run.py

run.py

@@ -1,44 +1,47 @@
-from langchain.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnablePassthrough
-from langchain_community.vectorstores import Chroma
-from transformers import AutoModelForCausalLM
-
-embedding = AutoModelForCausalLM.from_pretrained(
-    "deployllm/nomic-embed-text-v1.5-GGUF",
-    model_type='llama',
-    threads=3,
-)
-
-db = Chroma(
-    persist_directory="./chroma_db",
-    embedding_function=embedding,
-    collection_name='CVE'
-)
-
-retriever = db.as_retriever()
-
-template = """Answer the question based only on the following context:
-{context}
-Do not tell the source of the data
-Question: {question}
-"""
-
-prompt = ChatPromptTemplate.from_template(template)
-
-model = AutoModelForCausalLM.from_pretrained(
-    "zephyr-7b-beta.Q4_K_S.gguf",
-    model_type='mistral',
-    threads=3,
-)
-
-chain = (
-    {"context": retriever, "question": RunnablePassthrough()}
-    | prompt
-    | model
-    | StrOutputParser()
-)
-
-# Uncomment and use the following for testing
-# for chunk in chain.stream("Your question here"):
-#     print(chunk, end="", flush=True)
+from langchain.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
+from transformers import AutoModelForCausalLM
+
+model_kwargs = {'trust_remote_code': True}
+
+embedding = AutoModelForCausalLM.from_pretrained(
+    "nomic-embed-text-v1.5.Q4_k_S.gguf",
+    model_type='llama',
+    threads=3,
+)
+
+db = Chroma(
+    persist_directory="./chroma_db",
+    embedding_function=embedding,
+    collection_name='CVE'        
+)
+
+retriever = db.as_retriever()
+
+template = """Answer the question based only on the following context:
+{context}
+Do not tell the source of the data
+Question: {question}
+"""
+
+prompt = ChatPromptTemplate.from_template(template)
+
+model = AutoModelForCausalLM.from_pretrained(
+    "zephyr-7b-beta.Q4_K_S.gguf",
+    model_type='mistral',
+    threads=3,
+)
+
+chain = (
+    {"context": retriever, "question": RunnablePassthrough()}
+    | prompt
+    | model
+    | StrOutputParser()
+)
+
+# Uncomment and use the following for testing
+# for chunk in chain.stream("Your question here"):
+#     print(chunk, end="", flush=True)