Harrison/add huggingface hub (#23)

Add support for huggingface hub

I could not find a good way to enforce stop tokens over the huggingface
hub api - that needs to hopefully be cleaned up in the future
This commit is contained in:
Harrison Chase
2022-10-25 22:00:33 -07:00
committed by GitHub
parent 316aae8223
commit 020c42dcae
11 changed files with 228 additions and 28 deletions

View File

@@ -0,0 +1,19 @@
"""Test HuggingFace API wrapper."""
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)
output = llm("Say foo:")
assert isinstance(output, str)
def test_huggingface_call_error() -> None:
"""Test valid call to HuggingFace that errors."""
llm = HuggingFaceHub(max_new_tokens=-1)
with pytest.raises(ValueError):
llm("Say foo:")