Update llama_func.py

llama_func.py

@@ -1,7 +1,6 @@
 import os
 import logging
 
-from llama_index import GPTSimpleVectorIndex
 from llama_index import download_loader
 from llama_index import (
  Document,
@@ -10,181 +9,156 @@ from llama_index import (
  QuestionAnswerPrompt,
  RefinePrompt,
 )
-from langchain.llms import OpenAI
 import colorama
+import PyPDF2
+from tqdm import tqdm
 
+from modules.presets import *
+from modules.utils import *
+from modules.config import local_embedding
 
-from presets import *
-from utils import *
+
+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:
- 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
+ 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...")
+ try:
+ from modules.pdf_func import parse_pdf
+ from modules.config import advance_docs
+
+ two_column = advance_docs["pdf"].get("two_column", False)
+ pdftext = parse_pdf(filepath, two_column).text
+ except:
+ pdftext = ""
+ with open(filepath, "rb") as pdfFileObj:
+ pdfReader = PyPDF2.PdfReader(pdfFileObj)
+ for page in tqdm(pdfReader.pages):
+ pdftext += page.extract_text()
+ text_raw = pdftext
+ 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
 
 
 def construct_index(
  api_key,
  file_src,
  max_input_size=4096,
- num_outputs=1,
+ num_outputs=5,
  max_chunk_overlap=20,
  chunk_size_limit=600,
  embedding_limit=None,
  separator=" ",
- num_children=10,
- max_keywords_per_chunk=10,
 ):
- os.environ["OPENAI_API_KEY"] = api_key
+ 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"
  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,
- num_outputs,
- max_chunk_overlap,
- embedding_limit,
- chunk_size_limit,
+ max_input_size=max_input_size,
+ num_output=num_outputs,
+ max_chunk_overlap=max_chunk_overlap,
+ embedding_limit=embedding_limit,
+ chunk_size_limit=600,
  separator=separator,
  )
- documents, index_name = get_documents(file_src)
+ index_name = get_index_name(file_src)
  if os.path.exists(f"./index/{index_name}.json"):
- logging.info("找到了缓存的索引文件，加载中…")
+ logging.info("找到了缓存的索引文件，加载中……")
  return GPTSimpleVectorIndex.load_from_disk(f"./index/{index_name}.json")
  else:
  try:
- logging.debug("构建索引中……")
- index = GPTSimpleVectorIndex(
- documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper
- )
+ 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("索引构建完成！")
  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():