add text2text generation (#93)

fixes issue #90
This commit is contained in:
Harrison Chase
2022-11-08 18:08:46 -08:00
committed by GitHub
parent e48e562ea5
commit b9f61390e9
3 changed files with 39 additions and 45 deletions

View File

@@ -5,15 +5,22 @@ import pytest
from langchain.llms.huggingface_hub import HuggingFaceHub
def test_huggingface_call() -> None:
"""Test valid call to HuggingFace."""
llm = HuggingFaceHub(max_new_tokens=10)
def test_huggingface_text_generation() -> None:
"""Test valid call to HuggingFace text generation model."""
llm = HuggingFaceHub(repo_id="gpt2", model_kwargs={"max_new_tokens": 10})
output = llm("Say foo:")
assert isinstance(output, str)
def test_huggingface_text2text_generation() -> None:
"""Test valid call to HuggingFace text2text model."""
llm = HuggingFaceHub(repo_id="google/flan-t5-xl")
output = llm("The capital of New York is")
assert output == "Albany"
def test_huggingface_call_error() -> None:
"""Test valid call to HuggingFace that errors."""
llm = HuggingFaceHub(max_new_tokens=-1)
llm = HuggingFaceHub(model_kwargs={"max_new_tokens": -1})
with pytest.raises(ValueError):
llm("Say foo:")