mirror of https://github.com/hpcaitech/ColossalAI
63 lines
2.9 KiB
Python
63 lines
2.9 KiB
Python
import os
|
|
|
|
from colossalqa.retrieval_conversation_universal import UniversalRetrievalConversation
|
|
|
|
|
|
def test_en_retrievalQA():
|
|
data_path_en = os.environ.get('TEST_DATA_PATH_EN')
|
|
data_path_zh = os.environ.get('TEST_DATA_PATH_ZH')
|
|
en_model_path = os.environ.get('EN_MODEL_PATH')
|
|
zh_model_path = os.environ.get('ZH_MODEL_PATH')
|
|
zh_model_name = os.environ.get('ZH_MODEL_NAME')
|
|
en_model_name = os.environ.get('EN_MODEL_NAME')
|
|
sql_file_path = os.environ.get('SQL_FILE_PATH')
|
|
qa_session = UniversalRetrievalConversation(files_en=[{
|
|
'data_path': data_path_en,
|
|
'name': 'company information',
|
|
'separator': '\n'
|
|
}],
|
|
files_zh=[{
|
|
'data_path': data_path_zh,
|
|
'name': 'company information',
|
|
'separator': '\n'
|
|
}],
|
|
zh_model_path=zh_model_path,
|
|
en_model_path=en_model_path,
|
|
zh_model_name=zh_model_name,
|
|
en_model_name=en_model_name,
|
|
sql_file_path=sql_file_path)
|
|
ans = qa_session.run("which company runs business in hotel industry?", which_language='en')
|
|
print(ans)
|
|
|
|
|
|
def test_zh_retrievalQA():
|
|
data_path_en = os.environ.get('TEST_DATA_PATH_EN')
|
|
data_path_zh = os.environ.get('TEST_DATA_PATH_ZH')
|
|
en_model_path = os.environ.get('EN_MODEL_PATH')
|
|
zh_model_path = os.environ.get('ZH_MODEL_PATH')
|
|
zh_model_name = os.environ.get('ZH_MODEL_NAME')
|
|
en_model_name = os.environ.get('EN_MODEL_NAME')
|
|
sql_file_path = os.environ.get('SQL_FILE_PATH')
|
|
qa_session = UniversalRetrievalConversation(files_en=[{
|
|
'data_path': data_path_en,
|
|
'name': 'company information',
|
|
'separator': '\n'
|
|
}],
|
|
files_zh=[{
|
|
'data_path': data_path_zh,
|
|
'name': 'company information',
|
|
'separator': '\n'
|
|
}],
|
|
zh_model_path=zh_model_path,
|
|
en_model_path=en_model_path,
|
|
zh_model_name=zh_model_name,
|
|
en_model_name=en_model_name,
|
|
sql_file_path=sql_file_path)
|
|
ans = qa_session.run("哪家公司在经营酒店业务？", which_language='zh')
|
|
print(ans)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_en_retrievalQA()
|
|
test_zh_retrievalQA()
|