# # Copyright 2025 The InfiniFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import logging import random import re import numpy as np from common.constants import LLMType from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_by_id, get_model_config_by_type_and_name from common.connection_utils import timeout from rag.flow.base import ProcessBase, ProcessParamBase from rag.flow.parser.pdf_chunk_metadata import finalize_pdf_chunk from rag.flow.tokenizer.schema import TokenizerFromUpstream from rag.nlp import rag_tokenizer from common import settings from rag.svr.task_executor import embed_limiter from common.token_utils import truncate from common.misc_utils import thread_pool_exec class TokenizerParam(ProcessParamBase): def __init__(self): super().__init__() self.search_method = ["full_text", "embedding"] self.filename_embd_weight = 0.1 self.fields = ["text"] def check(self): for v in self.search_method: self.check_valid_value(v.lower(), "Chunk method abnormal.", ["full_text", "embedding"]) def get_input_form(self) -> dict[str, dict]: return {} class Tokenizer(ProcessBase): component_name = "Tokenizer" async def _embedding(self, name, chunks): # Tokenization may legitimately produce zero chunks; embedding should be a no-op. if not chunks: return [], 0 parts = sum(["full_text" in self._param.search_method, "embedding" in self._param.search_method]) token_count = 0 if self._canvas._kb_id: e, kb = KnowledgebaseService.get_by_id(self._canvas._kb_id) if kb.tenant_embd_id: embd_model_config = get_model_config_by_id(kb.tenant_embd_id) else: embd_model_config = get_model_config_by_type_and_name(self._canvas._tenant_id, LLMType.EMBEDDING, kb.embd_id) else: embd_model_config = get_tenant_default_model_by_type(self._canvas._tenant_id, LLMType.EMBEDDING) embedding_model = LLMBundle(self._canvas._tenant_id, embd_model_config) texts = [] valid_pairs = [] for i, c in enumerate(chunks): txt = "" if isinstance(self._param.fields, str): self._param.fields=[self._param.fields] for f in self._param.fields: f = c.get(f) if isinstance(f, str): txt += f elif isinstance(f, list): txt += "\n".join(f) cleaned_txt = re.sub(r"]{0,12})?>", " ", txt).strip() if not cleaned_txt: continue texts.append(cleaned_txt) valid_pairs.append((i, c)) if not texts: return chunks, token_count vts, c = embedding_model.encode([name]) token_count += c tts = np.tile(vts[0], (len(texts), 1)) @timeout(60) def batch_encode(txts): nonlocal embedding_model return embedding_model.encode([truncate(c, embedding_model.max_length - 10) for c in txts]) cnts_batches = [] for i in range(0, len(texts), settings.EMBEDDING_BATCH_SIZE): async with embed_limiter: vts, c = await thread_pool_exec(batch_encode,texts[i : i + settings.EMBEDDING_BATCH_SIZE],) cnts_batches.append(vts) token_count += c if i % 33 == 32: self.callback(i * 1.0 / len(texts) / parts / settings.EMBEDDING_BATCH_SIZE + 0.5 * (parts - 1)) cnts_ = np.vstack(cnts_batches) if cnts_batches else np.array([]) cnts = cnts_ title_w = float(self._param.filename_embd_weight) vects = (title_w * tts + (1 - title_w) * cnts) if len(tts) == len(cnts) else cnts assert len(vects) == len(valid_pairs) for i, (_, ck) in enumerate(valid_pairs): v = vects[i].tolist() ck["q_%d_vec" % len(v)] = v return chunks, token_count async def _invoke(self, **kwargs): try: chunks = kwargs.get("chunks") if chunks is not None: kwargs["chunks"] = [c for c in chunks if c is not None] from_upstream = TokenizerFromUpstream.model_validate(kwargs) except Exception as e: self.set_output("_ERROR", f"Input error: {str(e)}") return self.set_output("output_format", "chunks") parts = sum(["full_text" in self._param.search_method, "embedding" in self._param.search_method]) if "full_text" in self._param.search_method: self.callback(random.randint(1, 5) / 100.0, "Start to tokenize.") # Branch on the declared upstream format so an empty chunk list stays on the chunk path. if from_upstream.output_format == "chunks": chunks = from_upstream.chunks or [] for i, ck in enumerate(chunks): ck["chunk_order_int"] = i ck["title_tks"] = rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", from_upstream.name)) ck["title_sm_tks"] = rag_tokenizer.fine_grained_tokenize(ck["title_tks"]) if ck.get("questions"): ck["question_kwd"] = ck["questions"].split("\n") ck["question_tks"] = rag_tokenizer.tokenize(str(ck["questions"])) if ck.get("keywords"): ck["important_kwd"] = ck["keywords"].split(",") ck["important_tks"] = rag_tokenizer.tokenize(str(ck["keywords"])) if ck.get("summary"): ck["content_ltks"] = rag_tokenizer.tokenize(str(ck["summary"])) ck["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(ck["content_ltks"]) elif ck.get("text"): ck["content_ltks"] = rag_tokenizer.tokenize(ck["text"]) ck["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(ck["content_ltks"]) if i % 100 == 99: self.callback(i * 1.0 / len(chunks) / parts) elif from_upstream.output_format in ["markdown", "text", "html"]: if from_upstream.output_format == "markdown": payload = from_upstream.markdown_result elif from_upstream.output_format == "text": payload = from_upstream.text_result else: payload = from_upstream.html_result if not payload: return "" ck = {"text": payload} if "full_text" in self._param.search_method: ck["title_tks"] = rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", from_upstream.name)) ck["title_sm_tks"] = rag_tokenizer.fine_grained_tokenize(ck["title_tks"]) ck["content_ltks"] = rag_tokenizer.tokenize(payload) ck["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(ck["content_ltks"]) chunks = [ck] else: # Empty JSON payloads are valid and should remain empty downstream. chunks = from_upstream.json_result or [] for i, ck in enumerate(chunks): ck["title_tks"] = rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", from_upstream.name)) ck["title_sm_tks"] = rag_tokenizer.fine_grained_tokenize(ck["title_tks"]) if not ck.get("text"): continue ck["content_ltks"] = rag_tokenizer.tokenize(ck["text"]) ck["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(ck["content_ltks"]) if i % 100 == 99: self.callback(i * 1.0 / len(chunks) / parts) self.callback(1.0 / parts, "Finish tokenizing.") if "embedding" in self._param.search_method: self.callback(random.randint(1, 5) / 100.0 + 0.5 * (parts - 1), "Start embedding inference.") if from_upstream.name.strip() == "": logging.warning("Tokenizer: empty name provided from upstream, embedding may be not accurate.") chunks, token_count = await self._embedding(from_upstream.name, chunks) self.set_output("embedding_token_consumption", token_count) self.callback(1.0, "Finish embedding.") chunks = [finalize_pdf_chunk(ck) for ck in chunks] self.set_output("chunks", chunks)