Skip to content

Commit 074b768

Browse files
author
Rossdan Craig rossdan@lastmileai.dev
committed
[HF][streaming][1/n] Text Summarization
TSIA Adding streaming functionality to text summarization model parser ## Test Plan Rebase onto and test it with 11ace0a. Follow the README from AIConfig Editor https://github.com/lastmile-ai/aiconfig/tree/main/python/src/aiconfig/editor#dev, then run these command ```bash aiconfig_path=/Users/rossdancraig/Projects/aiconfig/cookbooks/Gradio/huggingface.aiconfig.json parsers_path=/Users/rossdancraig/Projects/aiconfig/cookbooks/Gradio/hf_model_parsers.py alias aiconfig="python3 -m 'aiconfig.scripts.aiconfig_cli'" aiconfig edit --aiconfig-path=$aiconfig_path --server-port=8080 --server-mode=debug_servers --parsers-module-path=$parsers_path ``` Then in AIConfig Editor run the prompt (it will be streaming format by default) https://github.com/lastmile-ai/aiconfig/assets/151060367/e91a1d8b-a3e9-459c-9eb1-2d8e5ec58e73
1 parent 0c9e704 commit 074b768

File tree

2 files changed

+10
-3
lines changed

2 files changed

+10
-3
lines changed

extensions/HuggingFace/python/src/aiconfig_extension_hugging_face/local_inference/text_generation.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,7 @@ async def run_inference(
251251
not "stream" in completion_data or completion_data.get("stream") != False
252252
)
253253
if should_stream:
254-
tokenizer : AutoTokenizer = AutoTokenizer.from_pretrained(model_name)
254+
tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(model_name)
255255
streamer = TextIteratorStreamer(tokenizer)
256256
completion_data["streamer"] = streamer
257257

extensions/HuggingFace/python/src/aiconfig_extension_hugging_face/local_inference/text_summarization.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,13 +128,18 @@ def construct_stream_output(
128128
"metadata": {},
129129
}
130130
)
131+
131132
accumulated_message = ""
132133
for new_text in streamer:
133134
if isinstance(new_text, str):
135+
# For some reason these symbols aren't filtered out by the streamer
136+
new_text = new_text.replace("</s>", "")
137+
new_text = new_text.replace("<s>", "")
138+
134139
accumulated_message += new_text
135140
options.stream_callback(new_text, accumulated_message, 0)
136-
137141
output.data = accumulated_message
142+
138143
return output
139144

140145

@@ -245,7 +250,9 @@ async def run_inference(self, prompt: Prompt, aiconfig: "AIConfigRuntime", optio
245250

246251
# if stream enabled in runtime options and config, then stream. Otherwise don't stream.
247252
streamer = None
248-
should_stream = (options.stream if options else False) and (not "stream" in completion_data or completion_data.get("stream") != False)
253+
should_stream = (options.stream if options else False) and (
254+
not "stream" in completion_data or completion_data.get("stream") != False
255+
)
249256
if should_stream:
250257
tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(model_name)
251258
streamer = TextIteratorStreamer(tokenizer)

0 commit comments

Comments
 (0)