Update llama_func.py

llama_func.py

@@ -1,6 +1,7 @@
 import os
 import logging
 
+from llama_index import GPTSimpleVectorIndex
 from llama_index import download_loader
 from llama_index import (
  Document,
@@ -9,146 +10,181 @@ from llama_index import (
  QuestionAnswerPrompt,
  RefinePrompt,
 )
+from langchain.llms import OpenAI
 import colorama
-# import PyPDF2
-from tqdm import tqdm
+
 
 from presets import *
 from utils import *
-from config import local_embedding
-
-
-def get_index_name(file_src):
- file_paths = [x.name for x in file_src]
- file_paths.sort(key=lambda x: os.path.basename(x))
-
- md5_hash = hashlib.md5()
- for file_path in file_paths:
- with open(file_path, "rb") as f:
- while chunk := f.read(8192):
- md5_hash.update(chunk)
-
- return md5_hash.hexdigest()
-
-
-def block_split(text):
- blocks = []
- while len(text) > 0:
- blocks.append(Document(text[:1000]))
- text = text[1000:]
- return blocks
 
 
 def get_documents(file_src):
  documents = []
+ index_name = ""
  logging.debug("Loading documents...")
  logging.debug(f"file_src: {file_src}")
  for file in file_src:
- filepath = file.name
- filename = os.path.basename(filepath)
- file_type = os.path.splitext(filepath)[1]
- logging.info(f"loading file: {filename}")
- try:
- if file_type == ".pdf":
- logging.debug("Loading PDF...")
- CJKPDFReader = download_loader("CJKPDFReader")
- loader = CJKPDFReader()
- text_raw = loader.load_data(file=filepath)[0].text
- elif file_type == ".docx":
- logging.debug("Loading Word...")
- DocxReader = download_loader("DocxReader")
- loader = DocxReader()
- text_raw = loader.load_data(file=filepath)[0].text
- elif file_type == ".epub":
- logging.debug("Loading EPUB...")
- EpubReader = download_loader("EpubReader")
- loader = EpubReader()
- text_raw = loader.load_data(file=filepath)[0].text
- elif file_type == ".xlsx":
- logging.debug("Loading Excel...")
- text_list = excel_to_string(filepath)
- for elem in text_list:
- documents.append(Document(elem))
- continue
- else:
- logging.debug("Loading text file...")
- with open(filepath, "r", encoding="utf-8") as f:
- text_raw = f.read()
- except Exception as e:
- logging.error(f"Error loading file: {filename}")
- pass
- text = add_space(text_raw)
- # text = block_split(text)
- # documents += text
- documents += [Document(text)]
- logging.debug("Documents loaded.")
- return documents
+ logging.debug(f"file: {file.name}")
+ index_name += file.name
+ if os.path.splitext(file.name)[1] == ".pdf":
+ logging.debug("Loading PDF...")
+ CJKPDFReader = download_loader("CJKPDFReader")
+ loader = CJKPDFReader()
+ documents += loader.load_data(file=file.name)
+ elif os.path.splitext(file.name)[1] == ".docx":
+ logging.debug("Loading DOCX...")
+ DocxReader = download_loader("DocxReader")
+ loader = DocxReader()
+ documents += loader.load_data(file=file.name)
+ elif os.path.splitext(file.name)[1] == ".epub":
+ logging.debug("Loading EPUB...")
+ EpubReader = download_loader("EpubReader")
+ loader = EpubReader()
+ documents += loader.load_data(file=file.name)
+ else:
+ logging.debug("Loading text file...")
+ with open(file.name, "r", encoding="utf-8") as f:
+ text = add_space(f.read())
+ documents += [Document(text)]
+ index_name = sha1sum(index_name)
+ return documents, index_name
 
 
 def construct_index(
  api_key,
  file_src,
  max_input_size=4096,
- num_outputs=5,
+ num_outputs=1,
  max_chunk_overlap=20,
  chunk_size_limit=600,
  embedding_limit=None,
  separator=" ",
+ num_children=10,
+ max_keywords_per_chunk=10,
 ):
- from langchain.chat_models import ChatOpenAI
- from langchain.embeddings.huggingface import HuggingFaceEmbeddings
- from llama_index import GPTSimpleVectorIndex, ServiceContext, LangchainEmbedding, OpenAIEmbedding
-
- if api_key:
- os.environ["OPENAI_API_KEY"] = api_key
- else:
- # 由于一个依赖的愚蠢的设计，这里必须要有一个API KEY
- os.environ["OPENAI_API_KEY"] = "sk-xxxxxxx"
+ os.environ["OPENAI_API_KEY"] = api_key
  chunk_size_limit = None if chunk_size_limit == 0 else chunk_size_limit
  embedding_limit = None if embedding_limit == 0 else embedding_limit
  separator = " " if separator == "" else separator
 
+ llm_predictor = LLMPredictor(
+ llm=OpenAI(model_name="gpt-3.5-turbo-0301", openai_api_key=api_key)
+ )
  prompt_helper = PromptHelper(
- max_input_size=max_input_size,
- num_output=num_outputs,
- max_chunk_overlap=max_chunk_overlap,
- embedding_limit=embedding_limit,
- chunk_size_limit=600,
+ max_input_size,
+ num_outputs,
+ max_chunk_overlap,
+ embedding_limit,
+ chunk_size_limit,
  separator=separator,
  )
- index_name = get_index_name(file_src)
+ documents, index_name = get_documents(file_src)
  if os.path.exists(f"./index/{index_name}.json"):
  logging.info("找到了缓存的索引文件，加载中……")
  return GPTSimpleVectorIndex.load_from_disk(f"./index/{index_name}.json")
  else:
  try:
- documents = get_documents(file_src)
- if local_embedding:
- embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name = "sentence-transformers/distiluse-base-multilingual-cased-v2"))
- else:
- embed_model = OpenAIEmbedding()
- logging.info("构建索引中……")
- with retrieve_proxy():
- service_context = ServiceContext.from_defaults(
- prompt_helper=prompt_helper,
- chunk_size_limit=chunk_size_limit,
- embed_model=embed_model,
- )
- index = GPTSimpleVectorIndex.from_documents(
- documents, service_context=service_context
- )
- logging.debug("索引构建完成！")
+ logging.debug("构建索引中……")
+ index = GPTSimpleVectorIndex(
+ documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper
+ )
  os.makedirs("./index", exist_ok=True)
  index.save_to_disk(f"./index/{index_name}.json")
- logging.debug("索引已保存至本地!")
  return index
-
  except Exception as e:
- logging.error("索引构建失败！", e)
  print(e)
  return None
 
 
+def chat_ai(
+ api_key,
+ index,
+ question,
+ context,
+ chatbot,
+):
+ os.environ["OPENAI_API_KEY"] = api_key
+
+ logging.info(f"Question: {question}")
+
+ response, chatbot_display, status_text = ask_ai(
+ api_key,
+ index,
+ question,
+ replace_today(PROMPT_TEMPLATE),
+ REFINE_TEMPLATE,
+ SIM_K,
+ INDEX_QUERY_TEMPRATURE,
+ context,
+ )
+ if response is None:
+ status_text = "查询失败，请换个问法试试"
+ return context, chatbot
+ response = response
+
+ context.append({"role": "user", "content": question})
+ context.append({"role": "assistant", "content": response})
+ chatbot.append((question, chatbot_display))
+
+ os.environ["OPENAI_API_KEY"] = ""
+ return context, chatbot, status_text
+
+
+def ask_ai(
+ api_key,
+ index,
+ question,
+ prompt_tmpl,
+ refine_tmpl,
+ sim_k=1,
+ temprature=0,
+ prefix_messages=[],
+):
+ os.environ["OPENAI_API_KEY"] = api_key
+
+ logging.debug("Index file found")
+ logging.debug("Querying index...")
+ llm_predictor = LLMPredictor(
+ llm=OpenAI(
+ temperature=temprature,
+ model_name="gpt-3.5-turbo-0301",
+ prefix_messages=prefix_messages,
+ )
+ )
+
+ response = None # Initialize response variable to avoid UnboundLocalError
+ qa_prompt = QuestionAnswerPrompt(prompt_tmpl)
+ rf_prompt = RefinePrompt(refine_tmpl)
+ response = index.query(
+ question,
+ llm_predictor=llm_predictor,
+ similarity_top_k=sim_k,
+ text_qa_template=qa_prompt,
+ refine_template=rf_prompt,
+ response_mode="compact",
+ )
+
+ if response is not None:
+ logging.info(f"Response: {response}")
+ ret_text = response.response
+ nodes = []
+ for index, node in enumerate(response.source_nodes):
+ brief = node.source_text[:25].replace("\n", "")
+ nodes.append(
+ f"<details><summary>[{index+1}]\t{brief}...</summary><p>{node.source_text}</p></details>"
+ )
+ new_response = ret_text + "\n----------\n" + "\n\n".join(nodes)
+ logging.info(
+ f"Response: {colorama.Fore.BLUE}{ret_text}{colorama.Style.RESET_ALL}"
+ )
+ os.environ["OPENAI_API_KEY"] = ""
+ return ret_text, new_response, f"查询消耗了{llm_predictor.last_token_usage} tokens"
+ else:
+ logging.warning("No response found, returning None")
+ os.environ["OPENAI_API_KEY"] = ""
+ return None
+
+
 def add_space(text):
  punctuations = {"，": "， ", "。": "。 ", "？": "？ ", "！": "！ ", "：": "： ", "；": "； "}
  for cn_punc, en_punc in punctuations.items():