cli_demo.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637
  1. from configs.model_config import *
  2. from chains.local_doc_qa import LocalDocQA
  3. import os
  4. import nltk
  5. nltk.data.path = [os.path.join(os.path.dirname(__file__), "nltk_data")] + nltk.data.path
  6. # return top-k text chunk from vector store
  7. VECTOR_SEARCH_TOP_K = 10
  8. # LLM input history length
  9. LLM_HISTORY_LEN = 3
  10. # Show reply with source text from input document
  11. REPLY_WITH_SOURCE = True
  12. if __name__ == "__main__":
  13. local_doc_qa = LocalDocQA()
  14. local_doc_qa.init_cfg(llm_model=LLM_MODEL,
  15. embedding_model=EMBEDDING_MODEL,
  16. embedding_device=EMBEDDING_DEVICE,
  17. llm_history_len=LLM_HISTORY_LEN,
  18. top_k=VECTOR_SEARCH_TOP_K)
  19. vs_path = None
  20. while not vs_path:
  21. filepath = input("Input your local knowledge file path 请输入本地知识文件路径:")
  22. vs_path = local_doc_qa.init_knowledge_vector_store(filepath)
  23. history = []
  24. while True:
  25. query = input("Input your question 请输入问题:")
  26. resp, history = local_doc_qa.get_knowledge_based_answer(query=query,
  27. vs_path=vs_path,
  28. chat_history=history)
  29. if REPLY_WITH_SOURCE:
  30. print(resp)
  31. else:
  32. print(resp["result"])