mirror of
https://github.com/hwchase17/langchain.git
synced 2025-04-28 11:55:21 +00:00
8c6734325b
Using `pyupgrade` to get all `partners` code up to 3.9 standards (mostly, fixing old `typing` imports).
19 lines
662 B
Python
19 lines
662 B
Python
from collections.abc import Generator
|
|
|
|
from langchain_huggingface.llms import HuggingFacePipeline
|
|
|
|
|
|
def test_huggingface_pipeline_streaming() -> None:
|
|
"""Test streaming tokens from huggingface_pipeline."""
|
|
llm = HuggingFacePipeline.from_model_id(
|
|
model_id="gpt2", task="text-generation", pipeline_kwargs={"max_new_tokens": 10}
|
|
)
|
|
generator = llm.stream("Q: How do you say 'hello' in German? A:'", stop=["."])
|
|
stream_results_string = ""
|
|
assert isinstance(generator, Generator)
|
|
|
|
for chunk in generator:
|
|
assert isinstance(chunk, str)
|
|
stream_results_string = chunk
|
|
assert len(stream_results_string.strip()) > 1
|