feat: ctransformers support async chain (#6859)

- Description: Adding async method for CTransformers 
- Issue: I've found impossible without this code to run Websockets
inside a FastAPI micro service and a CTransformers model.
  - Tag maintainer: Not necessary yet, I don't like to mention directly 
  - Twitter handle: @_semoal
This commit is contained in:
Sergio Moreno
2023-07-10 10:23:41 +02:00
committed by GitHub
parent d2cf0d16b3
commit 21a353e9c2
2 changed files with 57 additions and 2 deletions

View File

@@ -1,4 +1,5 @@
"""Test C Transformers wrapper."""
import pytest
from langchain.llms import CTransformers
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
@@ -19,3 +20,20 @@ def test_ctransformers_call() -> None:
assert isinstance(output, str)
assert len(output) > 1
assert 0 < callback_handler.llm_streams <= config["max_new_tokens"]
@pytest.mark.asyncio
async def test_ctransformers_async_inference() -> None:
config = {"max_new_tokens": 5}
callback_handler = FakeCallbackHandler()
llm = CTransformers(
model="marella/gpt-2-ggml",
config=config,
callbacks=[callback_handler],
)
output = await llm._acall(prompt="Say foo:")
assert isinstance(output, str)
assert len(output) > 1
assert 0 < callback_handler.llm_streams <= config["max_new_tokens"]