combine llm and support multi inputs and warp support batch_flags (#511)

Author: wzh1994

lazyllm/engine/engine.py

@@ -7,6 +7,7 @@
 from .node_meta_hook import NodeMetaHook
 import inspect
 import functools
+from itertools import repeat
 import copy
 from abc import ABC, abstractclassmethod
 from enum import Enum
@@ -346,8 +347,16 @@ def make_diverter(nodes: List[dict]):
 
 
 @NodeConstructor.register('Warp', subitems=['nodes', 'resources'])
-def make_warp(nodes: List[dict], edges: List[dict] = [], resources: List[dict] = []):
-    return lazyllm.warp(make_graph(nodes, edges, resources, enable_server=False))
+def make_warp(nodes: List[dict], edges: List[dict] = [], resources: List[dict] = [],
+              batch_flags: Optional[List[int]] = None):
+    wp = lazyllm.warp(make_graph(nodes, edges, resources, enable_server=False))
+    if batch_flags and len(batch_flags) > 1:
+        def transform(*args):
+            args = [a if b else repeat(a) for a, b in zip(args, batch_flags)]
+            args = [lazyllm.package(a) for a in zip(*args)]
+            return args
+        wp = lazyllm.pipeline(transform, wp)
+    return wp
 
 
 @NodeConstructor.register('Loop', subitems=['nodes', 'resources'])
@@ -723,6 +732,30 @@ def make_online_llm(source: str, base_model: Optional[str] = None, prompt: Optio
                                         api_key=api_key, secret_key=secret_key).prompt(prompt, history=history)
 
 
+class LLM(lazyllm.ModuleBase):
+    def __init__(self, m: lazyllm.ModuleBase, keys: Optional[List[str]] = None):
+        super().__init__()
+        self._m = m
+        self._keys = keys
+
+    def forward(self, *args, **kw):
+        if self._keys and len(self._keys) > 1:
+            assert len(args) == len(self._keys)
+            args = ({k: a for k, a in zip(self._keys, args)},)
+        else:
+            assert len(args) == 1
+        return self._m(*args, **kw)
+
+
+@NodeConstructor.register('LLM')
+def make_llm(kw: dict):
+    type: str = kw.pop('type')
+    keys: Optional[List[str]] = kw.pop('keys', None)
+    assert type in ('local', 'online'), f'Invalid type {type} given'
+    if type == 'local': return LLM(make_local_llm(**kw), keys)
+    elif type == 'online': return LLM(make_online_llm(**kw), keys)
+
+
 class STT(lazyllm.Module):
     def __init__(self, base_model: Union[str, lazyllm.TrainableModule]):
         super().__init__()

tests/basic_tests/test_engine.py

@@ -289,6 +289,32 @@ def test_engine_warp(self):
         assert engine.run(gid, 2, 3, 4, 5) == (4, 9, 16, 25)
         assert "prompt_tokens" in self.get_last_report()
 
+    def test_engine_warp_transform(self):
+        nodes = [dict(id='1', kind='Code', name='code', args=dict(
+            code='def sum(x: int, y: int, z: int): return x + y + z'))]
+        edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')]
+
+        nodes = [
+            dict(
+                id="2",
+                kind="Warp",
+                name="warp",
+                args=dict(
+                    nodes=nodes,
+                    edges=edges,
+                    batch_flags=[True, False, True],
+                    _lazyllm_enable_report=True,
+                ),
+            )
+        ]
+        edges = [dict(iid='__start__', oid='2'), dict(iid='2', oid='__end__')]
+
+        engine = LightEngine()
+        engine.set_report_url(self.report_url)
+        gid = engine.start(nodes, edges)
+        assert engine.run(gid, [2, 3, 4, 5], 1, [1, 2, 3, 1]) == (4, 6, 8, 7)
+        assert "prompt_tokens" in self.get_last_report()
+
     def test_engine_formatter(self):
         nodes = [dict(id='1', kind='Formatter', name='f1', args=dict(ftype='python', rule='[:]'))]
         edges = [dict(iid='__start__', oid='1'), dict(iid='1', oid='__end__')]

tests/charge_tests/test_engine.py

@@ -232,20 +232,21 @@ def test_stream_and_hostory(self):
         builtin_history = [['水的沸点是多少？', '您好，我的答案是：水的沸点在标准大气压下是100摄氏度。'],
                            ['世界上最大的动物是什么？', '您好，我的答案是：蓝鲸是世界上最大的动物。'],
                            ['人一天需要喝多少水？', '您好，我的答案是：一般建议每天喝8杯水，大约2升。']]
-        nodes = [dict(id='1', kind='OnlineLLM', name='m1', args=dict(source='glm', stream=True, prompt=dict(
+        nodes = [dict(id='1', kind='LLM', name='m1', args=dict(source='glm', stream=True, type='online', keys=['query'],
+                                                               prompt=dict(
                       system='请将我的问题翻译成中文。请注意，请直接输出翻译后的问题，不要反问和发挥',
                       user='问题: {query} \n, 翻译:'))),
-                 dict(id='2', kind='OnlineLLM', name='m2',
-                      args=dict(source='glm', stream=True,
+                 dict(id='2', kind='LLM', name='m2',
+                      args=dict(source='glm', stream=True, type='online',
                                 prompt=dict(system='请参考历史对话，回答问题，并保持格式不变。', user='{query}'))),
-                 dict(id='3', kind='JoinFormatter', name='join', args=dict(type='to_dict', names=['query', 'answer'])),
-                 dict(id='4', kind='OnlineLLM', stream=False, name='m3',
-                      args=dict(source='glm', history=builtin_history, prompt=dict(
+                 dict(id='3', kind='LLM', stream=False, name='m3',
+                      args=dict(source='glm', type='online', keys=['query', 'answer'], history=builtin_history,
+                                prompt=dict(
                           system='你是一个问答机器人，会根据用户的问题作出回答。',
                           user='请结合历史对话和本轮的问题，总结我们的全部对话。本轮情况如下:\n {query}, 回答: {answer}')))]
         engine = LightEngine()
         gid = engine.start(nodes, edges=[['__start__', '1'], ['1', '2'], ['1', '3'], ['2', '3'],
-                                         ['3', '4'], ['4', '__end__']], _history_ids=['2', '4'])
+                                         ['3', '__end__']], _history_ids=['2', '3'])
         history = [['雨后为什么会有彩虹？', '您好，我的答案是：雨后阳光通过水滴发生折射和反射形成了彩虹。'],
                    ['月亮会发光吗？', '您好，我的答案是：月亮本身不会发光，它反射太阳光。'],
                    ['一年有多少天', '您好，我的答案是：一年有365天，闰年有366天。']]