加入索引模式的实时回答功能；适配llama_index 0.5.0；加入繁体中文支持

modules/chat_func.py

@@ -13,6 +13,9 @@ import colorama
 from duckduckgo_search import ddg
 import asyncio
 import aiohttp
+from llama_index.indices.query.vector_store import GPTVectorStoreIndexQuery
+from llama_index.indices.query.schema import QueryBundle
+from langchain.llms import OpenAIChat
 
 from modules.presets import *
 from modules.llama_func import *
@@ -63,7 +66,7 @@ def get_response(
     # 如果有自定义的api-url，使用自定义url发送请求，否则使用默认设置发送请求
     if shared.state.api_url != API_URL:
         logging.info(f"使用自定义API URL: {shared.state.api_url}")
-        
+
     response = requests.post(
         shared.state.api_url,
         headers=headers,
@@ -72,7 +75,7 @@ def get_response(
         timeout=timeout,
         proxies=proxies,
     )
-    
+
     return response
 
 
@@ -103,13 +106,17 @@ def stream_predict(
     else:
         chatbot.append((inputs, ""))
     user_token_count = 0
+    if fake_input is not None:
+        input_token_count = count_token(construct_user(fake_input))
+    else:
+        input_token_count = count_token(construct_user(inputs))
     if len(all_token_counts) == 0:
         system_prompt_token_count = count_token(construct_system(system_prompt))
         user_token_count = (
-            count_token(construct_user(inputs)) + system_prompt_token_count
+            input_token_count + system_prompt_token_count
         )
     else:
-        user_token_count = count_token(construct_user(inputs))
+        user_token_count = input_token_count
     all_token_counts.append(user_token_count)
     logging.info(f"输入token计数: {user_token_count}")
     yield get_return_value()
@@ -137,6 +144,8 @@ def stream_predict(
     yield get_return_value()
     error_json_str = ""
 
+    if fake_input is not None:
+        history[-2] = construct_user(fake_input)
     for chunk in tqdm(response.iter_lines()):
         if counter == 0:
             counter += 1
@@ -201,7 +210,10 @@ def predict_all(
         chatbot.append((fake_input, ""))
     else:
         chatbot.append((inputs, ""))
-    all_token_counts.append(count_token(construct_user(inputs)))
+    if fake_input is not None:
+        all_token_counts.append(count_token(construct_user(fake_input)))
+    else:
+        all_token_counts.append(count_token(construct_user(inputs)))
     try:
         response = get_response(
             openai_api_key,
@@ -224,13 +236,22 @@ def predict_all(
         status_text = standard_error_msg + ssl_error_prompt + error_retrieve_prompt
         return chatbot, history, status_text, all_token_counts
     response = json.loads(response.text)
-    content = response["choices"][0]["message"]["content"]
-    history[-1] = construct_assistant(content)
-    chatbot[-1] = (chatbot[-1][0], content+display_append)
-    total_token_count = response["usage"]["total_tokens"]
-    all_token_counts[-1] = total_token_count - sum(all_token_counts)
-    status_text = construct_token_message(total_token_count)
-    return chatbot, history, status_text, all_token_counts
+    if fake_input is not None:
+        history[-2] = construct_user(fake_input)
+    try:
+        content = response["choices"][0]["message"]["content"]
+        history[-1] = construct_assistant(content)
+        chatbot[-1] = (chatbot[-1][0], content+display_append)
+        total_token_count = response["usage"]["total_tokens"]
+        if fake_input is not None:
+            all_token_counts[-1] += count_token(construct_assistant(content))
+        else:
+            all_token_counts[-1] = total_token_count - sum(all_token_counts)
+        status_text = construct_token_message(total_token_count)
+        return chatbot, history, status_text, all_token_counts
+    except KeyError:
+        status_text = standard_error_msg + str(response)
+        return chatbot, history, status_text, all_token_counts
 
 
 def predict(
@@ -254,37 +275,55 @@ def predict(
         yield chatbot+[(inputs, "")], history, "开始生成回答……", all_token_counts
     if reply_language == "跟随问题语言（不稳定）":
         reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."
+    old_inputs = None
+    display_reference = []
+    limited_context = False
     if files:
+        limited_context = True
+        old_inputs = inputs
         msg = "加载索引中……（这可能需要几分钟）"
         logging.info(msg)
         yield chatbot+[(inputs, "")], history, msg, all_token_counts
         index = construct_index(openai_api_key, file_src=files)
         msg = "索引构建完成，获取回答中……"
+        logging.info(msg)
         yield chatbot+[(inputs, "")], history, msg, all_token_counts
-        history, chatbot, status_text = chat_ai(openai_api_key, index, inputs, history, chatbot, reply_language)
-        yield chatbot, history, status_text, all_token_counts
-        return
-
-    old_inputs = ""
-    link_references = []
-    if use_websearch:
+        llm_predictor = LLMPredictor(llm=OpenAIChat(temperature=0, model_name=selected_model))
+        prompt_helper = PromptHelper(max_input_size = 4096, num_output = 5, max_chunk_overlap = 20, chunk_size_limit=600)
+        service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
+        query_object = GPTVectorStoreIndexQuery(index.index_struct, service_context=service_context, similarity_top_k=5, vector_store=index._vector_store, docstore=index._docstore)
+        query_bundle = QueryBundle(inputs)
+        nodes = query_object.retrieve(query_bundle)
+        reference_results = [n.node.text for n in nodes]
+        reference_results = add_source_numbers(reference_results, use_source=False)
+        display_reference = add_details(reference_results)
+        display_reference = "\n\n" + "".join(display_reference)
+        inputs = (
+            replace_today(PROMPT_TEMPLATE)
+            .replace("{query_str}", inputs)
+            .replace("{context_str}", "\n\n".join(reference_results))
+            .replace("{reply_language}", reply_language )
+        )
+    elif use_websearch:
+        limited_context = True
         search_results = ddg(inputs, max_results=5)
         old_inputs = inputs
-        web_results = []
+        reference_results = []
         for idx, result in enumerate(search_results):
             logging.info(f"搜索结果{idx + 1}：{result}")
             domain_name = urllib3.util.parse_url(result["href"]).host
-            web_results.append(f'[{idx+1}]"{result["body"]}"\nURL: {result["href"]}')
-            link_references.append(f"{idx+1}. [{domain_name}]({result['href']})\n")
-        link_references = "\n\n" + "".join(link_references)
+            reference_results.append([result["body"], result["href"]])
+            display_reference.append(f"{idx+1}. [{domain_name}]({result['href']})\n")
+        reference_results = add_source_numbers(reference_results)
+        display_reference = "\n\n" + "".join(display_reference)
         inputs = (
             replace_today(WEBSEARCH_PTOMPT_TEMPLATE)
             .replace("{query}", inputs)
-            .replace("{web_results}", "\n\n".join(web_results))
+            .replace("{web_results}", "\n\n".join(reference_results))
             .replace("{reply_language}", reply_language )
         )
     else:
-        link_references = ""
+        display_reference = ""
 
     if len(openai_api_key) != 51:
         status_text = standard_error_msg + no_apikey_msg
@@ -317,7 +356,7 @@ def predict(
             temperature,
             selected_model,
             fake_input=old_inputs,
-            display_append=link_references
+            display_append=display_reference
         )
         for chatbot, history, status_text, all_token_counts in iter:
             if shared.state.interrupted:
@@ -337,7 +376,7 @@ def predict(
             temperature,
             selected_model,
             fake_input=old_inputs,
-            display_append=link_references
+            display_append=display_reference
         )
         yield chatbot, history, status_text, all_token_counts
 
@@ -350,6 +389,11 @@ def predict(
             + colorama.Style.RESET_ALL
         )
 
+    if limited_context:
+        history = history[-4:]
+        all_token_counts = all_token_counts[-2:]
+        yield chatbot, history, status_text, all_token_counts
+
     if stream:
         max_token = max_token_streaming
     else:

modules/llama_func.py

@@ -13,6 +13,8 @@ from llama_index import (
 from langchain.llms import OpenAI
 from langchain.chat_models import ChatOpenAI
 import colorama
+import PyPDF2
+from tqdm import tqdm
 
 from modules.presets import *
 from modules.utils import *
@@ -29,6 +31,12 @@ def get_index_name(file_src):
 
     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 = []
@@ -38,9 +46,12 @@ def get_documents(file_src):
         logging.info(f"loading file: {file.name}")
         if os.path.splitext(file.name)[1] == ".pdf":
             logging.debug("Loading PDF...")
-            CJKPDFReader = download_loader("CJKPDFReader")
-            loader = CJKPDFReader()
-            text_raw = loader.load_data(file=file.name)[0].text
+            pdftext = ""
+            with open(file.name, 'rb') as pdfFileObj:
+                pdfReader = PyPDF2.PdfReader(pdfFileObj)
+                for page in tqdm(pdfReader.pages):
+                    pdftext += page.extract_text()
+            text_raw = pdftext
         elif os.path.splitext(file.name)[1] == ".docx":
             logging.debug("Loading DOCX...")
             DocxReader = download_loader("DocxReader")
@@ -56,6 +67,8 @@ def get_documents(file_src):
             with open(file.name, "r", encoding="utf-8") as f:
                 text_raw = f.read()
         text = add_space(text_raw)
+        # text = block_split(text)
+        # documents += text
         documents += [Document(text)]
     logging.debug("Documents loaded.")
     return documents
@@ -65,13 +78,11 @@ 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,
+        separator=" "
 ):
     os.environ["OPENAI_API_KEY"] = api_key
     chunk_size_limit = None if chunk_size_limit == 0 else chunk_size_limit
@@ -81,14 +92,7 @@ def construct_index(
     llm_predictor = LLMPredictor(
         llm=ChatOpenAI(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,
-        separator=separator,
-    )
+    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, separator=separator)
     index_name = get_index_name(file_src)
     if os.path.exists(f"./index/{index_name}.json"):
         logging.info("找到了缓存的索引文件，加载中……")
@@ -97,7 +101,7 @@ def construct_index(
         try:
             documents = get_documents(file_src)
             logging.info("构建索引中……")
-            service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
+            service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper, chunk_size_limit=chunk_size_limit)
             index = GPTSimpleVectorIndex.from_documents(
                 documents,  service_context=service_context
             )

modules/presets.py

@@ -57,7 +57,8 @@ MODELS = [
 ]  # 可选的模型
 
 REPLY_LANGUAGES = [
-    "中文",
+    "简体中文",
+    "繁體中文",
     "English",
     "日本語",
     "Español",

modules/utils.py

@@ -375,8 +375,8 @@ def replace_today(prompt):
 
 
 def get_geoip():
-    response = requests.get("https://ipapi.co/json/", timeout=5)
     try:
+        response = requests.get("https://ipapi.co/json/", timeout=5)
         data = response.json()
     except:
         data = {"error": True, "reason": "连接ipapi失败"}
@@ -384,7 +384,7 @@ def get_geoip():
         logging.warning(f"无法获取IP地址信息。\n{data}")
         if data["reason"] == "RateLimited":
             return (
-                f"获取IP地理位置失败，因为达到了检测IP的速率限制。聊天功能可能仍然可用，但请注意，如果您的IP地址在不受支持的地区，您可能会遇到问题。"
+                f"获取IP地理位置失败，因为达到了检测IP的速率限制。聊天功能可能仍然可用。"
             )
         else:
             return f"获取IP地理位置失败。原因：{data['reason']}。你仍然可以使用聊天功能。"
@@ -457,7 +457,7 @@ def get_proxies():
 
     if proxies == {}:
         proxies = None
-        
+
     return proxies
 
 def run(command, desc=None, errdesc=None, custom_env=None, live=False):
@@ -500,4 +500,19 @@ Python: <span title="{sys.version}">{python_version}</span>
 Gradio: {gr.__version__}
  • 
 Commit: {commit_info}
-"""
+"""
+
+def add_source_numbers(lst, source_name = "Source", use_source = True):
+    if use_source:
+        return [f'[{idx+1}]\t "{item[0]}"\n{source_name}: {item[1]}' for idx, item in enumerate(lst)]
+    else:
+        return [f'[{idx+1}]\t "{item}"' for idx, item in enumerate(lst)]
+
+def add_details(lst):
+    nodes = []
+    for index, txt in enumerate(lst):
+        brief = txt[:25].replace("\n", "")
+        nodes.append(
+            f"<details><summary>{brief}...</summary><p>{txt}</p></details>"
+        )
+    return nodes

requirements.txt

@@ -10,3 +10,4 @@ Pygments
 llama_index
 langchain
 markdown
+PyPDF2