add share parameter to webmodule and add callable condition for AdaptiveTransform [skip ci] (#564)

Author: wzh1994

lazyllm/tools/rag/doc_processor.py

@@ -65,7 +65,7 @@ def _create_nodes_impl(self, p_nodes, group_name):
         # NOTE transform.batch_forward will set children for p_nodes, but when calling
         # transform.batch_forward, p_nodes has been upsert in the store.
         t = self._node_groups[group_name]['transform']
-        transform = AdaptiveTransform(t) if isinstance(t, list) or t.pattern else make_transform(t)
+        transform = AdaptiveTransform(t) if isinstance(t, list) or t.pattern else make_transform(t, group_name)
         nodes = transform.batch_forward(p_nodes, group_name)
         self._store.update_nodes(nodes)
         return nodes
@@ -118,7 +118,7 @@ def _reparse_group_recursive(self, p_nodes: List[DocNode], cur_name: str, doc_id
             raise Exception(f"Failed to remove nodes for docs {doc_ids} group {cur_name} from store")
 
         t = self._node_groups[cur_name]['transform']
-        transform = AdaptiveTransform(t) if isinstance(t, list) or t.pattern else make_transform(t)
+        transform = AdaptiveTransform(t) if isinstance(t, list) or t.pattern else make_transform(t, cur_name)
         nodes = transform.batch_forward(p_nodes, cur_name)
         # reparse need set global_metadata
         self._store.update_nodes(nodes)

lazyllm/tools/rag/transform.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from copy import copy as lite_copy
 from dataclasses import dataclass, field
 import requests
 import os
@@ -22,7 +23,7 @@ class TransformArgs():
     trans_node: Optional[bool] = None
     num_workers: int = 0
     kwargs: Dict = field(default_factory=dict)
-    pattern: Optional[str] = None
+    pattern: Optional[Union[str, Callable[[str], bool]]] = None
 
     @staticmethod
     def from_dict(d):
@@ -73,6 +74,7 @@ def split_text_keep_separator(text: str, separator: str) -> List[str]:
 class NodeTransform(ABC):
     def __init__(self, num_workers: int = 0):
         self._number_workers = num_workers
+        self._name = None
 
     def batch_forward(
         self, documents: Union[DocNode, List[DocNode]], node_group: str, **kwargs
@@ -100,36 +102,43 @@ def impl(node: DocNode):
     def transform(self, document: DocNode, **kwargs) -> List[Union[str, DocNode]]:
         raise NotImplementedError('Not implemented')
 
+    def with_name(self, name: Optional[str], *, copy: bool = True) -> 'NodeTransform':
+        if name is not None:
+            if copy: return lite_copy(self).with_name(name, copy=False)
+            self._name = name
+        return self
+
     def __call__(self, node: DocNode, **kwargs: Any) -> List[DocNode]:
         # Parent and child should not be set here.
         results = self.transform(node, **kwargs)
         if isinstance(results, (DocNode, str)): results = [results]
         return [DocNode(text=chunk) if isinstance(chunk, str) else chunk for chunk in results if chunk]
 
 
-def make_transform(t):
+def make_transform(t: Union[TransformArgs, Dict[str, Any]], group_name: Optional[str] = None) -> NodeTransform:
     if isinstance(t, dict): t = TransformArgs.from_dict(t)
     transform, trans_node, num_workers = t['f'], t['trans_node'], t['num_workers']
     num_workers = dict(num_workers=num_workers) if num_workers > 0 else dict()
-    return (transform(**t['kwargs'], **num_workers)
-            if isinstance(transform, type)
-            else transform if isinstance(transform, NodeTransform)
-            else FuncNodeTransform(transform, trans_node=trans_node, **num_workers))
+    return (transform(**t['kwargs'], **num_workers).with_name(group_name, copy=False) if isinstance(transform, type)
+            else transform.with_name(group_name) if isinstance(transform, NodeTransform)
+            else FuncNodeTransform(transform, trans_node=trans_node, **num_workers).with_name(group_name, copy=False))
 
 
 class AdaptiveTransform(NodeTransform):
-    def __init__(self, transforms: Union[List[TransformArgs], TransformArgs]):
-        super().__init__(num_workers=0)
+    def __init__(self, transforms: Union[List[Union[TransformArgs, Dict]], Union[TransformArgs, Dict]],
+                 num_workers: int = 0):
+        super().__init__(num_workers=num_workers)
         if not isinstance(transforms, (tuple, list)): transforms = [transforms]
         self._transformers = [(t.get('pattern'), make_transform(t)) for t in transforms]
 
     def transform(self, document: DocNode, **kwargs) -> List[Union[str, DocNode]]:
+        if not isinstance(document, DocNode): LOG.warning(f'Invalud document type {type(document)} got')
         for pt, transform in self._transformers:
-            if pt and not pt.startswith('*'): pt = os.path.join(str(os.cwd()), pt)
-            if not isinstance(document, DocNode):
-                LOG.warning(f'Invalud document type {type(document)} got')
-            if not pt or fnmatch.fnmatch(document.docpath, pt):
+            if pt and isinstance(pt, str) and not pt.startswith('*'): pt = os.path.join(str(os.cwd()), pt)
+            if not pt or (callable(pt) and pt(document.docpath)) or (
+                    isinstance(pt, str) and fnmatch.fnmatch(document.docpath, pt)):
                 return transform(document, **kwargs)
+        LOG.warning(f'No transform found for document {document.docpath} with group name `{self._name}`')
         return []
 
 

lazyllm/tools/webpages/webmodule.py

@@ -10,7 +10,7 @@
 import time
 import re
 from pathlib import Path
-from typing import List, Union
+from typing import List, Union, Optional, Any, Dict
 
 import lazyllm
 from lazyllm import LOG, globals, FileSystemQueue, OnlineChatModule, TrainableModule
@@ -34,10 +34,12 @@ class Mode:
         Refresh = 1
         Appendix = 2
 
-    def __init__(self, m, *, components=dict(), title='对话演示终端', port=None,
-                 history=[], text_mode=None, trace_mode=None, audio=False, stream=False,
-                 files_target=None, static_paths: Union[str, Path, List[str | Path]] = None,
-                 encode_files=False) -> None:
+    def __init__(self, m: Any, *, components: Dict[Any, Any] = dict(), title: str = '对话演示终端',
+                 port: Optional[Union[int, range, tuple, list]] = None, history: List[Any] = [],
+                 text_mode: Optional[Mode] = None, trace_mode: Optional[Mode] = None, audio: bool = False,
+                 stream: bool = False, files_target: Optional[Union[Any, List[Any]]] = None,
+                 static_paths: Optional[Union[str, Path, List[Union[str, Path]]]] = None,
+                 encode_files: bool = False, share: bool = False) -> None:
         super().__init__()
         # Set the static directory of gradio so that gradio can access local resources in the directory
         if isinstance(static_paths, (str, Path)):
@@ -66,6 +68,7 @@ def __init__(self, m, *, components=dict(), title='对话演示终端', port=Non
         self.stream = stream
         self.files_target = files_target if isinstance(files_target, list) or files_target is None else [files_target]
         self.encode_files = encode_files
+        self.share = share
         self.demo = self.init_web(components)
         self.url = None
         signal.signal(signal.SIGINT, self._signal_handler)
@@ -405,7 +408,7 @@ def _work(self):
         self.url = f'http://127.0.0.1:{port}'
         self.broadcast_url = f'http://0.0.0.0:{port}'
 
-        self.demo.queue().launch(server_name="0.0.0.0", server_port=port, prevent_thread_lock=True)
+        self.demo.queue().launch(server_name="0.0.0.0", server_port=port, prevent_thread_lock=True, share=self.share)
         LOG.success('LazyLLM webmodule launched successfully: Running on: '
                     f'{self.broadcast_url}, local URL: {self.url}')
 

tests/basic_tests/test_document.py

@@ -132,7 +132,8 @@ def test_register_with_pattern(self):
             TransformArgs(f=SentenceSplitter, pattern='*.txt', kwargs=dict(chunk_size=512, chunk_overlap=50)),
             dict(f=SentenceSplitter, kwargs=dict(chunk_size=256, chunk_overlap=25))])
         Document.create_node_group('AdaptiveChunk2', transform=AdaptiveTransform([
-            dict(f=SentenceSplitter, pattern='*.txt', kwargs=dict(chunk_size=512, chunk_overlap=50)),
+            dict(f=SentenceSplitter, pattern=(lambda x: x.endswith('.txt')),
+                 kwargs=dict(chunk_size=512, chunk_overlap=50)),
             TransformArgs(f=SentenceSplitter, pattern=None, kwargs=dict(chunk_size=256, chunk_overlap=25))]))
         doc = Document('rag_master')
         doc._impl._lazy_init()