fix chunking

2024-09-26 22:45:02 +08:00 · 2024-09-26 22:45:02 +08:00 · 542562dd5b
commit 542562dd5b
parent dd6fb786e2
7 changed files with 122 additions and 75 deletions
--- a/src/.env.template
+++ b/src/.env.template
@ -4,5 +4,4 @@ ZHIPUAPI=optional
 QIANFAN_ACCESS_KEY=optional
 QIANFAN_SECRET_KEY=optional
 HF_ENDPOINT='https://hf-mirror.com'
-CUDA_VISIBLE_DEVICES=0
-MODEL_ROOT_DIR=MODEL_ROOT_DIR_UNSET
+CUDA_VISIBLE_DEVICES=0
--- a/src/core/database.py
+++ b/src/core/database.py
@ -115,21 +115,16 @@ class DataBaseManager:
            db.files.append(new_file)
            new_files.append(new_file)

-        from src.core.indexing import chunk_file, chunk_text
+        from src.core.indexing import chunk
        for new_file in new_files:
            file_id = new_file["file_id"]
            idx = self.get_idx_by_fileid(db, file_id)
            db.files[idx]["status"] = "processing"

            try:
-                if new_file["type"] in ["txt", "docx", "md"]:
-                    nodes = chunk_file(new_file["path"], params=params)
-                else:
-                    text = self.read_text(new_file["path"])
-                    nodes = chunk_text(text, params)
-
+                nodes = chunk(new_file["path"], params=params)
                self.knowledge_base.add_documents(
-                    docs=[node.to_dict() for node in nodes],
+                    docs=[node.text for node in nodes],
                    collection_name=db.metaname,
                    file_id=file_id)

--- a/src/core/indexing.py
+++ b/src/core/indexing.py
@ -1,37 +1,38 @@
 import os
 from pathlib import Path
+from llama_index.core import Document
 from llama_index.core.node_parser import SimpleFileNodeParser
-from llama_index.readers.file import FlatReader
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.readers.file import FlatReader, DocxReader

 from src.utils import hashstr

-def chunk_text(text, params=None):
+def chunk(text_or_path, params=None):
    params = params or {}
-    from llama_index.core import Document
-    from llama_index.core.node_parser import SentenceSplitter
    chunk_size = int(params.get("chunk_size", 500))
    chunk_overlap = int(params.get("chunk_overlap", 20))
    splitter = SentenceSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )
-    doc = Document(id_=hashstr(text), text=text)
-    nodes = splitter.get_nodes_from_documents([doc])
-    return nodes

-def chunk_file(file, params=None):
-    parser = SimpleFileNodeParser()
-    if file.endswith(".txt"):
-        from llama_index.readers.file import FlatReader
-        docs = FlatReader().load_data(Path(file))
-    elif file.endswith(".docx"):
-        from llama_index.readers.file import DocxReader
-        docs = DocxReader().load_data(Path(file))
-    elif file.endswith(".md"):
-        from llama_index.readers.file import MarkdownReader
-        docs = MarkdownReader().load_data(Path(file))
+    if os.path.isfile(text_or_path) and "uploads" in text_or_path:
+        parser = SimpleFileNodeParser()
+        file_type = Path(text_or_path).suffix
+        if file_type in [".txt", ".json", ".md"]:
+            docs = FlatReader().load_data(Path(text_or_path))
+        elif file_type in [".docx"]:
+            docs = DocxReader().load_data(Path(text_or_path))
+        else:
+            raise ValueError(f"Unsupported file type `{file_type}`")
+
+        if params.get("use_parser"):
+            nodes = parser.get_nodes_from_documents(docs)
+        else:
+            nodes = splitter.get_nodes_from_documents(docs)
+
    else:
-        raise ValueError("Unsupported file type")
+        docs = [Document(id_=hashstr(text_or_path), text=text_or_path)]
+        nodes = splitter.get_nodes_from_documents(docs)

-    nodes = parser.get_nodes_from_documents(docs)
-    return nodes  # 返回节点
+    return nodes
--- a/src/views/tools_view.py
+++ b/src/views/tools_view.py
@ -16,21 +16,10 @@ def route_index():
    tools = [
        {
            "name": "text_chunking",
-            "title": "Text Chunking",
-            "description": "Chunking text into smaller pieces for better understanding.",
+            "title": "文本分块",
+            "description": "将文本分块以更好地理解。可以输入文本或者上传文件。",
            "url": "/tools/text_chunking",
            "method": "POST",
-            "params": [
-                {
-                    "name": "text",
-                    "type": "string",
-                    "description": "Text to be chunked."
-                },
-                {
-                    "name": "chunk_size",
-                    "type": "int",
-                }
-            ]
        }
    ]

@ -39,7 +28,7 @@ def route_index():

@tools.route("/text_chunking", methods=["POST"])
 def text_chunking():
-    from src.core.indexing import chunk_text
+    from src.core.indexing import chunk
    text = request.json.get("text")
-    nodes = chunk_text(text, params=request.json)
+    nodes = chunk(text, params=request.json)
    return jsonify({"nodes": [node.to_dict() for node in nodes]})
--- a/web/src/assets/base.css
+++ b/web/src/assets/base.css
@ -25,9 +25,9 @@
  --gray-500: #A7ACAF;
  --gray-400: #C2C7CA;
  --gray-300: #DEE3E6;
-  --gray-200: #EDF1F4;
-  --gray-100: #F5FAFD;
-  --gray-50: #F9FDFF;
+  --gray-200: #EFF1F4;
+  --gray-100: #F8FAFD;
+  --gray-50: #FBFDFF;

  --main-color: #1c6586;
  --main-color-dark: #004d5c;
--- a/web/src/components/TextChunkingComponent.vue
+++ b/web/src/components/TextChunkingComponent.vue
@ -11,10 +11,13 @@
          @finish="chunkText"
        >
          <a-form-item label="Chunk Size" name="chunkSize" >
-            <a-input v-model:value="params.chunkSize" />
+            <a-input v-model:value="params.chunkSize" :disabled="params.useParser" />
          </a-form-item>
          <a-form-item label="Chunk Overlap" name="chunkOverlap" >
-            <a-input v-model:value="params.chunkOverlap" />
+            <a-input v-model:value="params.chunkOverlap" :disabled="params.useParser" />
+          </a-form-item>
+          <a-form-item label="使用文件节点解析器" name="useParser" v-if="state.useFile">
+            <a-switch v-model:checked="params.useParser" />
          </a-form-item>

          <a-form-item>
@ -26,16 +29,38 @@
    </div>
    <div class="result-container">
      <div class="input-container">
-        <textarea v-model="text" placeholder="Enter text to chunk"></textarea>
-        <div class="infos">
+        <div class="actions">
+          <span :class="{'active': !state.useFile}" @click="state.useFile = false">上传文本</span>
+          <span :class="{'active': state.useFile}" @click="state.useFile = true">上传文件</span>
+        </div>
+        <div class="upload" v-if="state.useFile">
+          <a-upload-dragger
+            class="upload-dragger"
+            v-model:fileList="fileList"
+            name="file"
+            :max-count="1"
+            :disabled="state.uploading"
+            action="/api/database/upload"
+            @change="handleFileUpload"
+            @drop="handleDrop"
+          >
+            <p class="ant-upload-text">点击或者把文件拖拽到这里上传</p>
+            <p class="ant-upload-hint">
+              目前仅支持上传文本文件，如 .pdf, .txt, .md。且同名文件无法重复添加
+            </p>
+          </a-upload-dragger>
+        </div>
+        <textarea v-if="!state.useFile" v-model="text" placeholder="Enter text to chunk"></textarea>
+        <div class="infos" v-if="!state.useFile">
          <!-- <span>字数: {{ wordCount }}</span> -->
          <span>字符数: {{ charCount }}</span>
          <span>Token 数：{{ estimatedTokenCount }}</span>
        </div>
      </div>
+      <!-- <div>{{ chunks[0] }}</div> -->
      <div id="result-cards" class="result-cards">
        <div v-for="(chunk, index) in chunks" :key="index" class="chunk">
-          <p>{{ chunk }}</p>
+          <p><strong>#{{ index + 1 }}</strong> {{ chunk.text }}</p>
        </div>
      </div>
    </div>
@ -49,13 +74,17 @@ import { message } from 'ant-design-vue'
 const text = ref('');
 const state = reactive({
  loading: true,
+  uploading: false,
+  useFile: false,
 })

 const params = reactive({
  chunkSize: 500,
-  chunkOverlap: 20
+  chunkOverlap: 20,
+  useParser: false,
 })
 const chunks = ref([]);
+const fileList = ref([]);

 const wordCount = computed(() => text.value.split(/\s+/).filter(Boolean).length);
 const charCount = computed(() => text.value.length)
@ -78,19 +107,32 @@ const estimatedTokenCount = computed(() => {
 })

 const chunkText = async () => {
-  if (text.value.length === 0) {
-    message.error("请输入文本")
-    return
+  let text_or_file = ''
+  if (state.useFile) {
+    if (fileList.value.length === 0) {
+      message.error("请上传文件")
+      return
+    }
+    console.log(fileList.value)
+    text_or_file = fileList.value.filter(file => file.status === 'done').map(file => file.response.file_path)[0]
+  } else {
+    if (text.value.length === 0) {
+      message.error("请输入文本")
+      return
+    }
+    text_or_file = text.value
  }
+
  try {
    state.loading = true
    const response = await fetch('/api/tools/text_chunking', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
-        text: text.value,
+        text: text_or_file,
        chunk_size: params.chunkSize,
-        chunk_overlap: params.chunkOverlap
+        chunk_overlap: params.chunkOverlap,
+        use_parser: params.useParser
      })
    });

@ -99,7 +141,7 @@ const chunkText = async () => {
    }

    const data = await response.json();
-    chunks.value = data.nodes.map(node => node.text);
+    chunks.value = data.nodes;
    state.loading = false
  } catch (error) {
    console.error('Error chunking text:', error);
@ -141,6 +183,28 @@ const chunkText = async () => {
      flex-direction: column;
      margin-bottom: 15px;

+      .actions {
+        display: flex;
+        justify-content: space-between;
+        margin-bottom: 10px;
+        width: fit-content;
+        background-color: var(--gray-200);
+        padding: 8px 6px;
+        border-radius: 8px;
+
+        span {
+          color: var(--gray-900);
+          cursor: pointer;
+          padding: 4px 10px;
+          border-radius: 8px;
+          transition: background-color 0.3s;
+
+          &.active {
+            background-color: white;
+          }
+        }
+      }
+
      textarea {
        resize: vertical;
        height: 300px;
@ -149,7 +213,7 @@ const chunkText = async () => {
        border-radius: 8px;
        font-size: 1rem;
        transition: border-color 0.3s;
-        background-color: var(--gray-50);
+        background-color: var(--gray-100);

        &:focus {
          border-color: var(--main-color);
@ -174,38 +238,36 @@ const chunkText = async () => {
    }

    .chunk {
-      background-color: #f9fcff;
-      border-radius: 4px;
-      padding: 10px;
+      background-color: var(--main-5);
+      border: 1px solid var(--main-light-3);
+      border-radius: 8px;
+      padding: 16px;
      margin-bottom: 10px;
-      box-shadow: 0 1px 3px rgba(0, 0, 0, 0.1);
+      box-shadow: 0 0 10px 2px rgba(0, 0, 0, 0.01);
      break-inside: avoid; /* 避免元素被分割到不同列 */
+      // 强制换行
+      word-wrap: break-word;
+      word-break: break-all;
    }
  }
 }


-@media (max-width: 768px) {
+@media (max-width: 980px) {
  #result-cards {
    column-count: 1;
  }
 }

-@media (min-width: 769px) and (max-width: 1200px) {
+@media (min-width: 981px) and (max-width: 1500px) {
  #result-cards {
    column-count: 2;
  }
 }

-@media (min-width: 1201px) and (max-width: 1900px) {
+@media (min-width: 1501px){
  #result-cards {
    column-count: 3;
  }
 }
-
-@media (min-width: 1901px){
-  #result-cards {
-    column-count: 4;
-  }
-}
 </style>
--- a/web/src/views/ToolsView.vue
+++ b/web/src/views/ToolsView.vue
@ -107,6 +107,7 @@ onMounted(() => {

      p {
        margin: 0;
+        color: var(--gray-800);
      }
    }
  }