|
@@ -1,128 +0,0 @@
|
|
|
-import torch
|
|
|
-from langchain.agents import ZeroShotAgent, Tool, AgentExecutor
|
|
|
-from langchain.llms import OpenAI
|
|
|
-from langchain.memory import ConversationBufferMemory, ReadOnlySharedMemory
|
|
|
-from langchain.chains import LLMChain, RetrievalQA
|
|
|
-from langchain.embeddings.huggingface import HuggingFaceEmbeddings
|
|
|
-from langchain.prompts import PromptTemplate
|
|
|
-from langchain.text_splitter import CharacterTextSplitter
|
|
|
-from langchain.vectorstores import Chroma
|
|
|
-from langchain.document_loaders import TextLoader
|
|
|
-from models import ChatGLM
|
|
|
-import sentence_transformers
|
|
|
-import os
|
|
|
-import readline
|
|
|
-from pathlib import Path
|
|
|
-
|
|
|
-
|
|
|
-class ChatglmWithSharedMemoryOpenaiLLM:
|
|
|
-
|
|
|
- def __init__(self, params: dict = None):
|
|
|
- params = params or {}
|
|
|
- self.embedding_model = params.get('embedding_model', 'text2vec')
|
|
|
- self.vector_search_top_k = params.get('vector_search_top_k', 6)
|
|
|
- self.llm_model = params.get('llm_model', 'chatglm-6b')
|
|
|
- self.llm_history_len = params.get('llm_history_len', 10)
|
|
|
- self.device = 'cuda' if params.get('use_cuda', False) else 'cpu'
|
|
|
- self._embedding_model_dict = {
|
|
|
- "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
|
|
|
- "ernie-base": "nghuyong/ernie-3.0-base-zh",
|
|
|
- "text2vec": "GanymedeNil/text2vec-large-chinese",
|
|
|
- }
|
|
|
- self._llm_model_dict = {
|
|
|
- "chatglm-6b-int4-qe": "THUDM/chatglm-6b-int4-qe",
|
|
|
- "chatglm-6b-int4": "THUDM/chatglm-6b-int4",
|
|
|
- "chatglm-6b": "THUDM/chatglm-6b",
|
|
|
- }
|
|
|
- self.init_cfg()
|
|
|
- self.init_docsearch()
|
|
|
- self.init_state_of_history()
|
|
|
- self.summry_chain, self.memory = self.agents_answer()
|
|
|
- self.agent_chain = self.create_agent_chain()
|
|
|
-
|
|
|
- def init_cfg(self):
|
|
|
- self.chatglm = ChatGLM()
|
|
|
- self.chatglm.load_model(model_name_or_path=self._llm_model_dict[self.llm_model])
|
|
|
- self.chatglm.history_len = self.llm_history_len
|
|
|
- self.embeddings = HuggingFaceEmbeddings(model_name=self._embedding_model_dict[self.embedding_model],)
|
|
|
- self.embeddings.client = sentence_transformers.SentenceTransformer(self.embeddings.model_name,
|
|
|
- device=self.device)
|
|
|
-
|
|
|
- def init_docsearch(self):
|
|
|
- doc_path = str(Path.cwd() / "content/state_of_the_search.txt")
|
|
|
- loader = TextLoader(doc_path)
|
|
|
- documents = loader.load()
|
|
|
- text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
|
|
|
- texts = text_splitter.split_documents(documents)
|
|
|
- docsearch = Chroma.from_documents(texts, self.embeddings, collection_name="state-of-search")
|
|
|
- self.state_of_search = RetrievalQA.from_chain_type(llm=self.chatglm, chain_type="stuff", retriever=docsearch.as_retriever())
|
|
|
-
|
|
|
- def init_state_of_history(self):
|
|
|
- doc_path = str(Path.cwd() / "content/state_of_the_history.txt")
|
|
|
- loader = TextLoader(doc_path)
|
|
|
- documents = loader.load()
|
|
|
- text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
|
|
|
- texts = text_splitter.split_documents(documents)
|
|
|
- docsearch = Chroma.from_documents(texts, self.embeddings, collection_name="state-of-history")
|
|
|
- self.state_of_history = RetrievalQA.from_chain_type(llm=self.chatglm, chain_type="stuff", retriever=docsearch.as_retriever())
|
|
|
-
|
|
|
- def agents_answer(self):
|
|
|
- template = """This is a conversation between a human and a bot:
|
|
|
-
|
|
|
- {chat_history}
|
|
|
-
|
|
|
- Write a summary of the conversation for {input}:
|
|
|
- """
|
|
|
-
|
|
|
- prompt = PromptTemplate(
|
|
|
- input_variables=["input", "chat_history"],
|
|
|
- template=template
|
|
|
- )
|
|
|
- memory = ConversationBufferMemory(memory_key="chat_history")
|
|
|
- readonlymemory = ReadOnlySharedMemory(memory=memory)
|
|
|
- summry_chain = LLMChain(
|
|
|
- llm=self.chatglm,
|
|
|
- prompt=prompt,
|
|
|
- verbose=True,
|
|
|
- memory=readonlymemory, # use the read-only memory to prevent the tool from modifying the memory
|
|
|
- )
|
|
|
- return summry_chain, memory
|
|
|
-
|
|
|
- def create_agent_chain(self):
|
|
|
- tools = [
|
|
|
- Tool(
|
|
|
- name="State of Search QA System",
|
|
|
- func=self.state_of_search.run,
|
|
|
- description="当您需要搜索有关问题时非常有用。输入应该是一个完整的问题。"
|
|
|
- ),
|
|
|
- Tool(
|
|
|
- name="state-of-history-qa",
|
|
|
- func=self.state_of_history.run,
|
|
|
- description="跟露露的历史对话 - 当提出我们之间发生了什么事请时,这里面的回答是很有用的"
|
|
|
- ),
|
|
|
- Tool(
|
|
|
- name="Summary",
|
|
|
- func=self.summry_chain.run,
|
|
|
- description="useful for when you summarize a conversation. The input to this tool should be a string, representing who will read this summary."
|
|
|
- )
|
|
|
- ]
|
|
|
-
|
|
|
- prefix = """你需要充当一个倾听者,尽量回答人类的问题,你可以使用这里工具,它们非常有用:"""
|
|
|
- suffix = """Begin!
|
|
|
-
|
|
|
- {chat_history}
|
|
|
- Question: {input}
|
|
|
- {agent_scratchpad}"""
|
|
|
-
|
|
|
- prompt = ZeroShotAgent.create_prompt(
|
|
|
- tools,
|
|
|
- prefix=prefix,
|
|
|
- suffix=suffix,
|
|
|
- input_variables=["input", "chat_history", "agent_scratchpad"]
|
|
|
- )
|
|
|
-
|
|
|
- llm_chain = LLMChain(llm=OpenAI(temperature=0), prompt=prompt)
|
|
|
- agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
|
|
|
- agent_chain = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True, memory=self.memory)
|
|
|
-
|
|
|
- return agent_chain
|