|
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 ctransformers import AutoModelForCausalLM |
|
|
|
import os |
|
os.environ['TRANSFORMERS_CACHE'] = '/code/model/cache/' |
|
|
|
model_kwargs = {'trust_remote_code': True} |
|
|
|
|
|
|
|
|
|
|
|
|
|
db = Chroma( |
|
persist_directory="./chroma_db", |
|
|
|
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() |
|
) |
|
|
|
|
|
|
|
|
|
|