Upload 4 files

llm/llm.py

@@ -0,0 +1,10 @@
+from langchain.llms import CTransformers
+
+def setup_llm():  
+    llm = CTransformers(
+        model="models/mistral-7b-instruct-v0.1.Q8_0.gguf",
+        model_type="mistral",
+        config={"max_new_tokens": 2048, "context_length": 4096, "temperature": 0},
+    )
+
+    return llm

llm/prompts.py

@@ -0,0 +1,14 @@
+# Note: Precise formatting of spacing and indentation of the prompt template is important,
+# as it is highly sensitive to whitespace changes. For example, it could have problems generating
+# a summary from the pieces of context if the spacing is not done correctly
+
+qa_template = """Your role is financial controller. 
+You are working on invoice documents.
+Your main work is to extract data from the invoice.
+I would like to extract the following data from the invoices: date, number, sender, final amount and short description of what was purchased. 
+
+Context: {context}
+Question: {question}
+
+Helpful answer:
+"""

llm/wrapper.py

@@ -0,0 +1,38 @@
+import box
+import yaml
+
+from langchain.prompts import PromptTemplate
+from langchain.chains import RetrievalQA
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from llm.prompts import qa_template
+from llm.llm import setup_llm
+
+# Import config vars
+with open('config.yml', 'r', encoding='utf8') as ymlfile:
+    cfg = box.Box(yaml.safe_load(ymlfile))
+
+
+def set_qa_prompt():
+    prompt = PromptTemplate(template=qa_template, input_variables=['context', 'question'])
+    return prompt
+
+
+def build_retrieval_qa_chain(llm, prompt, vectordb):
+    qa_chain = RetrievalQA.from_chain_type(llm=llm,
+                                       chain_type='stuff',
+                                       retriever=vectordb.as_retriever(search_kwargs={'k': cfg.VECTOR_COUNT}),
+                                       return_source_documents=cfg.RETURN_SOURCE_DOCUMENTS,
+                                       chain_type_kwargs={'prompt': prompt}
+                                       )
+    return qa_chain
+
+
+def setup_qa_chain():
+    embeddings = HuggingFaceEmbeddings(model_name=cfg.EMBEDDINGS,model_kwargs={'device': 'cpu'})
+    vectordb = FAISS.load_local(cfg.DB_FAISS_PATH, embeddings)
+    llm = setup_llm()
+    qa_prompt = set_qa_prompt()
+    qa_chain = build_retrieval_qa_chain(llm, qa_prompt, vectordb)
+
+    return qa_chain