Update async aconvbot

README.md

@@ -24,6 +24,15 @@ prin(convertbot("How are you?"))
 # I am good  # or along that line
 ```
 
+The async version `aconvbot`, potentialy for `fastapi` or `Nonebot` plugins and such,  is rather artificial since it's based on `ThreadPoolExecutor`. Hence it's not intended for production. You probably should not spawn too many instances.
+```python
+from convbot import aconvbot
+
+async def afunc(text):
+    resp = await aconvbot(text)
+    ...
+```
+
 Interactive
 
 ```bash
@@ -31,4 +40,4 @@ python -m convbot
 ```
 ## Not tested in Windows 10 and Mac
 
-The module uses pytorch that is installed differently in Windows than in Linux. To run in Windows or Mac, you can probably just try to install pytorch manually.
+The module uses pytorch that is installed differently in Windows than in Linux. To run `convbot` in Windows or Mac, you may give it a spin  by cloning the repo (git clone [https://github.com/ffreemt/convbot](https://github.com/ffreemt/convbot)) and installing pytorch manually.

convbot/__init__.py

@@ -1,5 +1,8 @@
 """Init."""
 __version__ = "0.1.0"
-from .convbot import convbot
+from .convbot import convbot, aconvbot
 
-__all__ = ("convbot",)
+__all__ = (
+    "convbot",
+    "aconvbot",
+)

convbot/convbot.py

@@ -4,6 +4,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 from logzero import logger
 
+from .force_async import force_async
+
 # model_name = "microsoft/DialoGPT-large"
 # model_name = "microsoft/DialoGPT-small"
 # pylint: disable=invalid-name
@@ -39,7 +41,7 @@ def _convbot(
         chat_history_ids = ""
 
     input_ids = tokenizer.encode(text + tokenizer.eos_token, return_tensors="pt")
-    if chat_history_ids:
+    if isinstance(chat_history_ids, torch.Tensor):
         bot_input_ids = torch.cat([chat_history_ids, input_ids], dim=-1)
     else:
         bot_input_ids = input_ids
@@ -113,6 +115,24 @@ def convbot(
     return resp
 
 
+@force_async
+def aconvbot(
+    text: str,
+    n_retries: int = 3,
+    max_length: int = 1000,
+    do_sample: bool = True,
+    top_p: float = 0.95,
+    top_k: int = 0,
+    temperature: float = 0.75,
+) -> str:
+    try:
+        _ = convbot(text,n_retries, max_length, do_sample, top_p, top_k, temperature)
+    except Exception as e:
+        logger.error(e)
+        raise
+    return _
+
+
 def main():
     print("Bot: Talk to me")
     while 1:

convbot/force_async.py

@@ -0,0 +1,43 @@
+"""Turn a sync func to async."""
+from concurrent.futures import ThreadPoolExecutor
+import asyncio
+import functools
+
+def force_async(func):
+    """Turn a sync func to async.
+
+    Args:
+        func: a sync func
+
+    Return:
+        async func
+
+    Usage:
+        @force_async
+        def normal_func():
+            ...
+        loop = asyncio.get_event_loop()
+        #~ tasks = [sync_loop1(1, 5), sync_loop1(2, 10)]
+        #~ res = loop.run_until_complete(asyncio.gather(*tasks))  # OK
+        res = loop.run_until_complete(
+            asyncio.gather(
+                *[
+                    sync_loop1(1, 7),
+                    sync_loop1(2, 6),
+                    sync_loop1(2, 6),
+                    async_func,
+                ]
+            )
+        )
+    """
+    # executor = ThreadPoolExecutor()
+    # from concurrent.futures import ThreadPoolExecutor
+    executor = ThreadPoolExecutor(max_workers=10)
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        """Preserve func info."""
+        future = executor.submit(func, *args, **kwargs)
+        return asyncio.wrap_future(future)  # make it awaitable
+
+    return wrapper

tests/test_convbot.py

@@ -1,6 +1,11 @@
 """Test convbot."""
+# import asyncio
+import pytest
+
 from convbot import __version__
-from convbot import convbot
+from convbot import convbot, aconvbot
+
+pytestmark = pytest.mark.asyncio
 
 
 def test_version():
@@ -11,12 +16,26 @@ def test_version():
 def test_sanity():
     """Sanity check."""
     try:
-        assert not convbot()
+        assert not convbot("")
     except Exception:
         assert True
 
 
 def test_convbot():
+    """Test convbot."""
     resp = convbot("How are you?")
     assert len(resp) > 3
 
+    # 2nd call uses chat_history_ids
+    resp = convbot("How old are you?")
+    assert len(resp) > 3
+
+
+async def tests_aconvbot():
+    """Test aconvbot."""
+    resp = await aconvbot("How are you?")
+    assert len(resp) > 3
+
+    # 2nd call uses chat_history_ids
+    resp = await aconvbot("How old are you?")
+    assert len(resp) > 3