cli_demo.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  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 = 6
  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. last_print_len = 0
  27. for resp, history in local_doc_qa.get_knowledge_based_answer(query=query,
  28. vs_path=vs_path,
  29. chat_history=history,
  30. streaming=STREAMING):
  31. if STREAMING:
  32. print(resp["result"][last_print_len:], end="", flush=True)
  33. last_print_len = len(resp["result"])
  34. else:
  35. print(resp["result"])
  36. if REPLY_WITH_SOURCE:
  37. source_text = [f"""出处 [{inum + 1}] {os.path.split(doc.metadata['source'])[-1]}:\n\n{doc.page_content}\n\n"""
  38. # f"""相关度:{doc.metadata['score']}\n\n"""
  39. for inum, doc in
  40. enumerate(resp["source_documents"])]
  41. print("\n\n" + "\n\n".join(source_text))