196 lines
8.8 KiB
Plaintext
196 lines
8.8 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 1,
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from langchain.chains.question_answering import load_qa_chain\n",
|
||
"from langchain.prompts import PromptTemplate\n",
|
||
"from lib.embeds import MyEmbeddings\n",
|
||
"from lib.faiss import FAISSVS\n",
|
||
"from langchain.chains.combine_documents.stuff import StuffDocumentsChain\n",
|
||
"from langchain.chains.llm import LLMChain\n",
|
||
"from lib.chatglm_llm import ChatGLM, AlpacaGLM\n",
|
||
"from lib.config import *\n",
|
||
"from lib.utils import get_docs\n",
|
||
"\n",
|
||
"\n",
|
||
"class LocalDocQA:\n",
|
||
" def __init__(self, \n",
|
||
" embedding_model=EMBEDDING_MODEL, \n",
|
||
" embedding_device=EMBEDDING_DEVICE, \n",
|
||
" llm_model=LLM_MODEL, \n",
|
||
" llm_device=LLM_DEVICE, \n",
|
||
" llm_history_len=LLM_HISTORY_LEN, \n",
|
||
" top_k=VECTOR_SEARCH_TOP_K,\n",
|
||
" vs_name = VS_NAME\n",
|
||
" ) -> None:\n",
|
||
" \n",
|
||
" torch.cuda.empty_cache()\n",
|
||
" torch.cuda.empty_cache()\n",
|
||
"\n",
|
||
" self.embedding_model = embedding_model\n",
|
||
" self.llm_model = llm_model\n",
|
||
" self.embedding_device = embedding_device\n",
|
||
" self.llm_device = llm_device\n",
|
||
" self.llm_history_len = llm_history_len\n",
|
||
" self.top_k = top_k\n",
|
||
" self.vs_name = vs_name\n",
|
||
"\n",
|
||
" self.llm = AlpacaGLM()\n",
|
||
" self.llm.load_model(model_name_or_path=llm_model_dict[llm_model], llm_device=llm_device)\n",
|
||
"\n",
|
||
" self.embeddings = MyEmbeddings(model_name=embedding_model_dict[embedding_model])\n",
|
||
" self.load_vector_store(vs_name)\n",
|
||
"\n",
|
||
" self.prompt = PromptTemplate(\n",
|
||
" template=PROMPT_TEMPLATE,\n",
|
||
" input_variables=[\"context\", \"question\"]\n",
|
||
" )\n",
|
||
" self.search_params = {\n",
|
||
" \"engine\": \"bing\",\n",
|
||
" \"gl\": \"us\",\n",
|
||
" \"hl\": \"en\",\n",
|
||
" \"serpapi_api_key\": \"\"\n",
|
||
" }\n",
|
||
"\n",
|
||
" def init_knowledge_vector_store(self, vs_name: str):\n",
|
||
" \n",
|
||
" docs = get_docs(KNOWLEDGE_PATH)\n",
|
||
" vector_store = FAISSVS.from_documents(docs, self.embeddings)\n",
|
||
" vs_path = VECTORSTORE_PATH + vs_name\n",
|
||
" vector_store.save_local(vs_path)\n",
|
||
"\n",
|
||
" def add_knowledge_to_vector_store(self, vs_name: str):\n",
|
||
" docs = get_docs(ADD_KNOWLEDGE_PATH)\n",
|
||
" new_vector_store = FAISSVS.from_documents(docs, self.embeddings)\n",
|
||
" vector_store = FAISSVS.load_local(VECTORSTORE_PATH + vs_name, self.embeddings) \n",
|
||
" vector_store.merge_from(new_vector_store)\n",
|
||
" vector_store.save_local(VECTORSTORE_PATH + vs_name)\n",
|
||
"\n",
|
||
" def load_vector_store(self, vs_name: str):\n",
|
||
" self.vector_store = FAISSVS.load_local(VECTORSTORE_PATH + vs_name, self.embeddings)\n",
|
||
"\n",
|
||
" # def get_search_based_answer(self, query):\n",
|
||
" \n",
|
||
" # search = SerpAPIWrapper(params=self.search_params)\n",
|
||
" # docs = search.run(query)\n",
|
||
" # search_chain = load_qa_chain(self.llm, chain_type=\"stuff\")\n",
|
||
" # answer = search_chain.run(input_documents=docs, question=query)\n",
|
||
"\n",
|
||
" # return answer\n",
|
||
" \n",
|
||
" def get_knowledge_based_answer(self, query):\n",
|
||
" \n",
|
||
" docs = self.vector_store.max_marginal_relevance_search(query)\n",
|
||
" print(f'召回的文档和相似度分数：{docs}')\n",
|
||
" # 这里 doc[1] 就是对应的score \n",
|
||
" docs = [doc[0] for doc in docs]\n",
|
||
" \n",
|
||
" document_prompt = PromptTemplate(\n",
|
||
" input_variables=[\"page_content\"], template=\"Context:\\n{page_content}\"\n",
|
||
" )\n",
|
||
" llm_chain = LLMChain(llm=self.llm, prompt=self.prompt)\n",
|
||
" combine_documents_chain = StuffDocumentsChain(\n",
|
||
" llm_chain=llm_chain,\n",
|
||
" document_variable_name=\"context\",\n",
|
||
" document_prompt=document_prompt,\n",
|
||
" )\n",
|
||
" answer = combine_documents_chain.run(\n",
|
||
" input_documents=docs, question=query\n",
|
||
" )\n",
|
||
"\n",
|
||
" self.llm.history[-1][0] = query\n",
|
||
" self.llm.history[-1][-1] = answer\n",
|
||
" return answer, docs, self.llm.history"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 2,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"application/vnd.jupyter.widget-view+json": {
|
||
"model_id": "d4342213010c4ed2ad5b04694aa436d6",
|
||
"version_major": 2,
|
||
"version_minor": 0
|
||
},
|
||
"text/plain": [
|
||
"Loading checkpoint shards: 0%| | 0/3 [00:00<?, ?it/s]"
|
||
]
|
||
},
|
||
"metadata": {},
|
||
"output_type": "display_data"
|
||
}
|
||
],
|
||
"source": [
|
||
"qa = LocalDocQA()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 3,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"name": "stdout",
|
||
"output_type": "stream",
|
||
"text": [
|
||
"召回的文档和相似度分数：[(Document(page_content='****** LOGI APT Group Intelligence Research Yearbook APT Knowledge Graph APT组织情报 研究年鉴', metadata={'source': './KnowledgeStore/APT group Intelligence Research handbook-2022.pdf', 'page': 0}), 0.45381865), (Document(page_content='9 MANDIANT APT42: Crooked Charms, Cons and Compromises FIGURE 8. APT42 impersonates University of Oxford vaccinologist. APT42 Credential harvesting page masquerading as a Yahoo login portal.', metadata={'source': './KnowledgeStore/APT42_Crooked_Charms_Cons_and_Compromises.pdf', 'page': 8}), 0.4535672), (Document(page_content='The origin story of APT32 macros T H R E A T R E S E A R C H R E P O R T R u n n i n g t h r o u g h a l l t h e S U O f i l e s t r u c t u r e s i s l a b o r i o u s a n d d i d n ’ t y i e l d m u c h m o r e t h a n a s t r i n g d u m p w o u l d h a v e d o n e a n y w a y . W e f i n d p a t h s t o s o u r c e c o d e f i l e s , p r o j e c t n a m e s , e t c . W e c a n i n f e r f r o m t h e m y r i a d o f r e f e r e n c e s i n XmlPackageOptions , O u t l i n i n g S t a t e D i r , e t c . , t h a t t h e HtaDotnet a n d ShellcodeLoader s o l u t i o n s w e r e o r i g i n a l l y u n d e r t h e f o l d e r p a t h G:\\\\WebBuilder\\\\Gift_HtaDotnet\\\\ . T h i s i s a l s o s u p p o r t e d b y t h e P D B p a t h s o f o l d e r b u i l t b i n a r i e s w i t h i n t h e b r o a d e r S t r i k e S u i t G i f t p a c k a g e . F r o m l o o k i n g a t D e b u g g e r W a t c h e s v a l u e s i n o t h e r p r o j e c t s , w e c a n s e e t h a t t h e m a l w a r e d e v e l o p e r w a s a c t i v e l y d e b u g g i n g t h e h i s t o r i c a l p r o g r a m s . S U O f i l e D e b u g g e r W a t c h e s WebBuilder/HtaDotNet/HtaDotnet.v11.suo result WebBuilder/ShellcodeLoader/.vs/L/v14/.suo (char)77 WebBuilder/ShellcodeLoader/L.suo (char)77 3 4 04/2022', metadata={'source': './KnowledgeStore/Stairwell-threat-report-The-origin-of-APT32-macros.pdf', 'page': 33}), 0.38091612), (Document(page_content='2 APTs and COVID-19: How advanced persistent threats use the coronavirus as a lureTable of contents Introduction: APT groups using COVID-19 .........................................................', metadata={'source': './KnowledgeStore/200407-MWB-COVID-White-Paper_Final.pdf', 'page': 1}), 0.44476452)]\n"
|
||
]
|
||
}
|
||
],
|
||
"source": [
|
||
"query = r\"\"\"make a brief introduction of APT?\"\"\"\n",
|
||
"ans, docs, _ = qa.get_knowledge_based_answer(query)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": 4,
|
||
"metadata": {},
|
||
"outputs": [
|
||
{
|
||
"data": {
|
||
"text/plain": [
|
||
"'\\nAnswer: APT stands for Advanced Persistent Threat, which is a type of malicious cyberattack that is carried out by a sophisticated hacker group or state-sponsored organization. APTs are designed to remain undetected for a long period of time and are often used to steal sensitive data or disrupt critical infrastructure.'"
|
||
]
|
||
},
|
||
"execution_count": 4,
|
||
"metadata": {},
|
||
"output_type": "execute_result"
|
||
}
|
||
],
|
||
"source": [
|
||
"ans"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "chatgpt",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"codemirror_mode": {
|
||
"name": "ipython",
|
||
"version": 3
|
||
},
|
||
"file_extension": ".py",
|
||
"mimetype": "text/x-python",
|
||
"name": "python",
|
||
"nbconvert_exporter": "python",
|
||
"pygments_lexer": "ipython3",
|
||
"version": "3.10.9"
|
||
},
|
||
"orig_nbformat": 4
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 2
|
||
}
|