Merge pull request #331 from NVIDIA/feature/huggingface-pipeline-stre…

…aming

Feature/huggingface pipeline streaming

docs/user_guides/advanced/streaming.md

@@ -86,3 +86,32 @@ POST /v1/chat/completions
     "stream": true
 }
 ```
+
+### Streaming for LLMs deployed using HuggingFacePipeline
+
+We also support streaming for LLMs deployed using `HuggingFacePipeline`.
+One example is provided in the [HF Pipeline Dolly](./../../../examples/configs/llm/hf_pipeline_dolly/README.md) configuration.
+
+To use streaming for HF Pipeline LLMs, you first need to set the streaming flag in your `config.yml`.
+
+```yaml
+streaming: True
+```
+
+Then you need to create an `nemoguardrails.llm.providers.huggingface.AsyncTextIteratorStreamer` streamer object,
+add it to the `kwargs` of the pipeline and to the `model_kwargs` of the `HuggingFacePipelineCompatible` object.
+
+```python
+from nemoguardrails.llm.providers.huggingface import AsyncTextIteratorStreamer
+
+# instantiate tokenizer object required by LLM
+streamer = AsyncTextIteratorStreamer(tokenizer, skip_prompt=True)
+params = {"temperature": 0.01, "max_new_tokens": 100, "streamer": streamer}
+
+pipe = pipeline(
+    # all other parameters
+    **params,
+)
+
+llm = HuggingFacePipelineCompatible(pipeline=pipe, model_kwargs=params)
+```

docs/user_guides/llm-support.md

@@ -25,19 +25,21 @@ If you want to use an LLM and you cannot see a prompt in the [prompts folder](..
 | Dialog Rails                                       | :heavy_check_mark: (0.74) | :heavy_check_mark: (0.83) | :heavy_check_mark: (0.82) | :heavy_check_mark: (0.77) | :heavy_check_mark: (0.76) | :exclamation: (0.45)      | :exclamation:      | :exclamation: (0.54) | :exclamation: (0.54) | :exclamation: (0.50) | :exclamation: (0.40) | :exclamation: _(DEPENDS ON MODEL)_ |
 | • Single LLM call                             | :heavy_check_mark: (0.83) | :heavy_check_mark: (0.81) | :heavy_check_mark:        | :x:                       | :x:                       | :x:                       | :x:                | :x:                  | :x:                  | :x:                  | :x:                  | :x:                                |
 | • Multi-step flow generation                  | _EXPERIMENTAL_            | _EXPERIMENTAL_            | :x:                       | :x:                       | :x:                       | :x:                       | :x:                | :x:                  | :x:                  | :x:                  | :x:                  | :x:                                |
-| Streaming  	                                       | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :x:                       | :x:                       | :heavy_check_mark:        | :heavy_check_mark: | :x:                  | :x:                  | :x:                  | :x:                  | :x:                                |
+| Streaming  	                                       | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | -                         | -                         | :heavy_check_mark:        | :heavy_check_mark: | -                    | -                    | -                    | -                    | :heavy_check_mark:                 |
 | Hallucination detection (SelfCheckGPT with AskLLM) | :heavy_check_mark:        | :heavy_check_mark:        | :x:                       | :x:                       | :x:                       | :x:                       | :x:                | :x:                  | :x:                  | :x:                  | :x:                  | :x:                                |
 | AskLLM rails                                       |                           |                           |                           |                           |                           |                           |                    |                      |                      |                      |                      |                                    |
 | • Jailbreak detection                         | :heavy_check_mark: (0.88) | :heavy_check_mark: (0.88) | :heavy_check_mark: (0.86) | :x:                       | :x:                       | :heavy_check_mark: (0.85) | :x:                | :x:                  | :x:                  | :x:                  | :x:                  | :x:                                |
 | • Output moderation                           | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :x:                       | :x:                       | :heavy_check_mark: (0.85) | :x:                | :x:                  | :x:                  | :x:                  | :x:                  | :x:                                |
 | • Fact-checking                               | :heavy_check_mark: (0.81) | :heavy_check_mark: (0.82) | :heavy_check_mark: (0.81) | :heavy_check_mark: (0.80) | :x:                       | :heavy_check_mark: (0.83) | :x:                | :x:                  | :x:                  | :x:                  | :x:                  | :exclamation: _(DEPENDS ON MODEL)_ |
  | AlignScore fact-checking _(LLM independent)_       | :heavy_check_mark: (0.89) | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark: | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:                 |
 | ActiveFence moderation _(LLM independent)_         | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark: | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:                 |
+| Llama Guard moderation _(LLM independent)_         | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark:        | :heavy_check_mark: | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:   | :heavy_check_mark:                 |
 
 Table legend:
 - :heavy_check_mark: - Supported (_The feature is fully supported by the LLM based on our experiments and tests_)
 - :exclamation: - Limited Support (_Experiments and tests show that the LLM is under-performing for that feature_)
 - :x: - Not Supported (_Experiments show very poor performance or no experiments have been done for the LLM-feature pair_)
+- \- - Not Applicable (_e.g. models support streaming, it depends how they are deployed_)
 
 The performance numbers reported in the table above for each LLM-feature pair are as follows:
 - the banking dataset evaluation for dialog (topical) rails

examples/configs/llm/hf_pipeline_dolly/README.md

@@ -1,6 +1,7 @@
 # HuggingFace Pipeline with Dolly models
 
 This configuration uses the HuggingFace Pipeline LLM with the [dolly-v2-3b](https://huggingface.co/databricks/dolly-v2-3b) model.
+It also shows how to support streaming in NeMo Guardrails for LLMs deployed using HuggingFacePipeline.
 
 The `dolly-v2-3b` LLM model has been tested on the topical rails evaluation sets, results are available [here](../../../../nemoguardrails/eval/README.md).
 

examples/configs/llm/hf_pipeline_dolly/config.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 from functools import lru_cache
 
-from torch.cuda import device_count
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, pipeline
 
 from nemoguardrails.llm.helpers import get_llm_instance_wrapper
 from nemoguardrails.llm.providers import (
@@ -24,16 +24,40 @@
 
 
 @lru_cache
-def get_dolly_v2_3b_llm():
-    repo_id = "databricks/dolly-v2-3b"
-    params = {"temperature": 0, "max_length": 1024}
+def get_dolly_v2_3b_llm(streaming: bool = True):
+    name = "databricks/dolly-v2-3b"
 
-    # Use the first CUDA-enabled GPU, if any
-    device = 0 if device_count() else -1
+    config = AutoConfig.from_pretrained(name, trust_remote_code=True)
+    device = "cpu"
+    config.init_device = device
+    config.max_seq_len = 45
 
-    llm = HuggingFacePipelineCompatible.from_model_id(
-        model_id=repo_id, device=device, task="text-generation", model_kwargs=params
+    model = AutoModelForCausalLM.from_pretrained(
+        name,
+        config=config,
+        trust_remote_code=True,
     )
+    tokenizer = AutoTokenizer.from_pretrained(name)
+    params = {"temperature": 0.01, "max_new_tokens": 100}
+
+    # If we want streaming, we create a streamer.
+    if streaming:
+        from nemoguardrails.llm.providers.huggingface import AsyncTextIteratorStreamer
+
+        streamer = AsyncTextIteratorStreamer(tokenizer, skip_prompt=True)
+        params["streamer"] = streamer
+
+    pipe = pipeline(
+        model=model,
+        task="text-generation",
+        tokenizer=tokenizer,
+        device=device,
+        do_sample=True,
+        use_cache=True,
+        **params,
+    )
+
+    llm = HuggingFacePipelineCompatible(pipeline=pipe, model_kwargs=params)
 
     return llm
 

examples/configs/llm/hf_pipeline_dolly/config.yml

@@ -2,6 +2,9 @@ models:
   - type: main
     engine: hf_pipeline_dolly
 
+# Remove attribute / set to False if streaming is not required
+streaming: True
+
 instructions:
   - type: general
     content: |

examples/configs/llm/hf_pipeline_dolly/rails.co

@@ -1,20 +1,27 @@
-define user express greeting
-  "Hello"
-  "Hi"
+# Uncomment these to test dialog rails
 
-define user ask capabilities
-  "What can you do?"
-  "What can you help me with?"
-  "tell me what you can do"
-  "tell me about you"
-
-define flow
-  user express greeting
-  bot express greeting
-
-define flow
-  user ask capabilities
-  bot inform capabilities
-
-define bot inform capabilities
-  "I am an AI assistant and I'm here to help."
+#define user express greeting
+#  "Hello"
+#  "Hi"
+#
+#define user ask capabilities
+#  "What can you do?"
+#  "What can you help me with?"
+#  "tell me what you can do"
+#  "tell me about you"
+#
+#define user ask general question
+#  "How is the weather tomorrow?"
+#  "Can you tell me which is the best movie this week"
+#  "i would like to know the best scifi books of all time."
+#
+#define flow
+#  user express greeting
+#  bot express greeting
+#
+#define flow
+#  user ask capabilities
+#  bot inform capabilities
+#
+#define bot inform capabilities
+#  "I am an AI assistant and I'm here to help."

examples/scripts/demo_streaming.py

@@ -66,6 +66,19 @@ async def process_tokens():
     print(result)
 
 
+async def demo_hf_pipeline():
+    """Demo for streaming of response chunks directly with HuggingFacePipline deployed LLMs."""
+    config = RailsConfig.from_path(config_path="./../configs/llm/hf_pipeline_dolly")
+    app = LLMRails(config)
+
+    history = [{"role": "user", "content": "What is the capital of France?"}]
+
+    async for chunk in app.stream_async(messages=history):
+        print(f"CHUNK: {chunk}")
+        # Or do something else with the token
+
+
 if __name__ == "__main__":
     asyncio.run(demo_1())
     asyncio.run(demo_2())
+    # asyncio.run(demo_hf_pipeline())

nemoguardrails/llm/helpers.py

@@ -32,6 +32,14 @@ def get_llm_instance_wrapper(
     """
 
     class WrapperLLM(LLM):
+        """The wrapper class needs to have defined any parameters we need to be set by NeMo Guardrails.
+
+        Currently added only streaming and temperature.
+        """
+
+        streaming: Optional[bool] = False
+        temperature: Optional[float] = 1.0
+
         @property
         def model_kwargs(self):
             """Return the model's kwargs.
@@ -50,12 +58,24 @@ def _llm_type(self) -> str:
             """
             return llm_type
 
+        def _modify_instance_kwargs(self):
+            """Modify the parameters of the llm_instance with the attributes set for the wrapper.
+
+            This will allow the actual LLM instance to use the parameters at generation.
+            TODO: Make this function more generic if needed.
+            """
+
+            if hasattr(llm_instance, "model_kwargs"):
+                llm_instance.model_kwargs["temperature"] = self.temperature
+                llm_instance.model_kwargs["streaming"] = self.streaming
+
         def _call(
             self,
             prompt: str,
             stop: Optional[List[str]] = None,
             run_manager: Optional[CallbackManagerForLLMRun] = None,
         ) -> str:
+            self._modify_instance_kwargs()
             return llm_instance._call(prompt, stop, run_manager)
 
         async def _acall(
@@ -64,6 +84,7 @@ async def _acall(
             stop: Optional[List[str]] = None,
             run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         ) -> str:
+            self._modify_instance_kwargs()
             return await llm_instance._acall(prompt, stop, run_manager)
 
     return WrapperLLM

nemoguardrails/llm/providers/huggingface/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+from .streamers import AsyncTextIteratorStreamer

nemoguardrails/llm/providers/huggingface/streamers.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 asyncio
+
+from transformers.generation.streamers import TextStreamer
+
+
+class AsyncTextIteratorStreamer(TextStreamer):
+    """
+    Simple async implementation for HuggingFace Transformers streamers.
+
+    This follows closely how transformers.generation.streamers.TextIteratorStreamer works,
+    with minor modifications to make it async.
+    """
+
+    def __init__(
+        self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, **decode_kwargs
+    ):
+        super().__init__(tokenizer, skip_prompt, **decode_kwargs)
+        self.text_queue = asyncio.Queue()
+        self.stop_signal = None
+        self.loop = None
+
+    def on_finalized_text(self, text: str, stream_end: bool = False):
+        """Put the new text in the queue. If the stream is ending, also put a stop signal in the queue."""
+        if len(text) > 0:
+            asyncio.run_coroutine_threadsafe(self.text_queue.put(text), self.loop)
+
+        if stream_end:
+            asyncio.run_coroutine_threadsafe(self.text_queue.put(text), self.loop)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        value = await self.text_queue.get()
+        if value == self.stop_signal:
+            raise StopAsyncIteration()
+        else:
+            return value