mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
118 lines
4.6 KiB
118 lines
4.6 KiB
import os
|
|
|
|
from colossalqa.data_loader.document_loader import DocumentLoader
|
|
from colossalqa.local.llm import ColossalAPI, ColossalLLM
|
|
from colossalqa.memory import ConversationBufferWithSummary
|
|
from colossalqa.prompt.prompt import PROMPT_RETRIEVAL_QA_ZH
|
|
from colossalqa.retriever import CustomRetriever
|
|
from langchain.embeddings import HuggingFaceEmbeddings
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
|
|
|
|
def test_memory_long():
|
|
model_path = os.environ.get("EN_MODEL_PATH")
|
|
data_path = os.environ.get("TEST_DATA_PATH_EN")
|
|
model_name = os.environ.get("EN_MODEL_NAME")
|
|
sql_file_path = os.environ.get("SQL_FILE_PATH")
|
|
|
|
if not os.path.exists(sql_file_path):
|
|
os.makedirs(sql_file_path)
|
|
|
|
colossal_api = ColossalAPI.get_api(model_name, model_path)
|
|
llm = ColossalLLM(n=4, api=colossal_api)
|
|
memory = ConversationBufferWithSummary(
|
|
llm=llm, max_tokens=600, llm_kwargs={"max_new_tokens": 50, "temperature": 0.6, "do_sample": True}
|
|
)
|
|
retriever_data = DocumentLoader([[data_path, "company information"]]).all_data
|
|
|
|
# Split
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
|
|
splits = text_splitter.split_documents(retriever_data)
|
|
|
|
embedding = HuggingFaceEmbeddings(
|
|
model_name="moka-ai/m3e-base", model_kwargs={"device": "cpu"}, encode_kwargs={"normalize_embeddings": False}
|
|
)
|
|
|
|
# Create retriever
|
|
information_retriever = CustomRetriever(k=3, sql_file_path=sql_file_path)
|
|
information_retriever.add_documents(docs=splits, cleanup="incremental", mode="by_source", embedding=embedding)
|
|
|
|
memory.initiate_document_retrieval_chain(
|
|
llm,
|
|
PROMPT_RETRIEVAL_QA_ZH,
|
|
information_retriever,
|
|
chain_type_kwargs={
|
|
"chat_history": "",
|
|
},
|
|
)
|
|
|
|
# This keep the prompt length excluding dialogues the same
|
|
docs = information_retriever.get_relevant_documents("this is a test input.")
|
|
prompt_length = memory.chain.prompt_length(docs, **{"question": "this is a test input.", "chat_history": ""})
|
|
remain = 600 - prompt_length
|
|
have_summarization_flag = False
|
|
for i in range(40):
|
|
chat_history = memory.load_memory_variables({"question": "this is a test input.", "input_documents": docs})[
|
|
"chat_history"
|
|
]
|
|
|
|
assert memory.get_conversation_length() <= remain
|
|
memory.save_context({"question": "this is a test input."}, {"output": "this is a test output."})
|
|
if "A summarization of historical conversation:" in chat_history:
|
|
have_summarization_flag = True
|
|
assert have_summarization_flag == True
|
|
|
|
|
|
def test_memory_short():
|
|
model_path = os.environ.get("EN_MODEL_PATH")
|
|
data_path = os.environ.get("TEST_DATA_PATH_EN")
|
|
model_name = os.environ.get("EN_MODEL_NAME")
|
|
sql_file_path = os.environ.get("SQL_FILE_PATH")
|
|
|
|
if not os.path.exists(sql_file_path):
|
|
os.makedirs(sql_file_path)
|
|
|
|
colossal_api = ColossalAPI.get_api(model_name, model_path)
|
|
llm = ColossalLLM(n=4, api=colossal_api)
|
|
memory = ConversationBufferWithSummary(
|
|
llm=llm, llm_kwargs={"max_new_tokens": 50, "temperature": 0.6, "do_sample": True}
|
|
)
|
|
retriever_data = DocumentLoader([[data_path, "company information"]]).all_data
|
|
|
|
# Split
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
|
|
splits = text_splitter.split_documents(retriever_data)
|
|
|
|
embedding = HuggingFaceEmbeddings(
|
|
model_name="moka-ai/m3e-base", model_kwargs={"device": "cpu"}, encode_kwargs={"normalize_embeddings": False}
|
|
)
|
|
|
|
# create retriever
|
|
information_retriever = CustomRetriever(k=3, sql_file_path=sql_file_path)
|
|
information_retriever.add_documents(docs=splits, cleanup="incremental", mode="by_source", embedding=embedding)
|
|
|
|
memory.initiate_document_retrieval_chain(
|
|
llm,
|
|
PROMPT_RETRIEVAL_QA_ZH,
|
|
information_retriever,
|
|
chain_type_kwargs={
|
|
"chat_history": "",
|
|
},
|
|
)
|
|
|
|
# This keep the prompt length excluding dialogues the same
|
|
docs = information_retriever.get_relevant_documents("this is a test input.", return_scores=True)
|
|
|
|
for i in range(4):
|
|
chat_history = memory.load_memory_variables({"question": "this is a test input.", "input_documents": docs})[
|
|
"chat_history"
|
|
]
|
|
assert chat_history.count("Assistant: this is a test output.") == i
|
|
assert chat_history.count("Human: this is a test input.") == i
|
|
memory.save_context({"question": "this is a test input."}, {"output": "this is a test output."})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_memory_short()
|
|
test_memory_long()
|