Merge branch 'main' of https://github.com/PyRepair/hyperchain

Author: lomqer

.github/workflows/test.yml

@@ -23,4 +23,30 @@ jobs:
       
       - name: Run tests
         run: tox
+  
+  publish-on-pypi:
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    needs: test
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/')
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: '3.10'
+
+      - name: Install dependencies
+        run: |
+          pip install build twine
+
+      - name: Build package
+        run: |
+          python -m build
 
+      - name: Publish on PyPI
+        uses: pypa/gh-action-pypi-publish@v1.4.2

src/hyperchain/chain/__init__.py

@@ -1,4 +1,4 @@
-from .chain_result import ChainResult
+from .chain_result import ChainResult, ChainResultList
 
 from .chain import Chain
 from .chain_sequence import ChainSequence

src/hyperchain/chain/chain.py

@@ -4,7 +4,7 @@
 
 import asyncio
 
-from .chain_result import ChainResult
+from .chain_result import ChainResult, ChainResultList
 from ..prompt_templates import Template
 
 class Chain(ABC):
@@ -27,8 +27,10 @@ def run_multiple(self, *inputs_dict: Dict[str, Any]) -> List[ChainResult]:
     async def async_run_multiple(
         self, *inputs_dict: Dict[str, Any]
     ) -> List[ChainResult]:
-        return await asyncio.gather(
-            *[self.async_run(**input_list) for input_list in inputs_dict]
+        return await ChainResultList(
+            asyncio.gather(
+                *[self.async_run(**input_list) for input_list in inputs_dict]
+            )
         )
 
     def __add__(self, other: Any) -> Chain:

src/hyperchain/chain/chain_result.py

@@ -12,6 +12,10 @@ class ChainResult:
     def __post_init__(self):
         self.output_dict = deepcopy(self.output_dict)
 
+    def __getstate__(self): return self.output_dict
+
+    def __setstate__(self, state): self.output_dict = state
+
     def __getattr__(self, name):
         if name == "output_dict":
             return self.output_dict
@@ -24,3 +28,12 @@ def __getattr__(self, name):
             return str(self.output_dict[name])
 
         return None
+
+
+class ChainResultList(list):
+    def __getstate__(self): return self.__dict__
+
+    def __setstate__(self, state): self.__dict__ = state
+
+    def __getattr__(self, name):
+        return [chain.__getattr__(name) for chain in self]

src/hyperchain/chain/llm_chain.py

@@ -1,7 +1,7 @@
-from typing import List, Any
+from typing import List, Any, Dict
 import logging
 
-from .chain_result import ChainResult
+from .chain_result import ChainResult, ChainResultList
 from .chain import Chain
 from .chain_sequence import ChainSequence
 
@@ -40,9 +40,8 @@ def __init__(
         self.output_keys = [output_name]
         self.required_keys = template.required_keys
 
-    async def async_run(self, **inputs_dict: Any) -> ChainResult:
+    async def _run_with_error_handling(self, task):
         handlers = self.llm_runner._get_error_handlers()
-        prompt = self.template.format(**inputs_dict)
         while True:
             holds_a_lock = False
             try:
@@ -65,7 +64,7 @@ async def async_run(self, **inputs_dict: Any) -> ChainResult:
                         holds_a_lock = False
                         self._error_handling_lock.release()
 
-                    result = await self.llm_runner.async_run(prompt)
+                    result = await task
 
                     if not holds_a_lock:
                         async with self._error_handling_lock:
@@ -79,9 +78,8 @@ async def async_run(self, **inputs_dict: Any) -> ChainResult:
                         holds_a_lock = False
                         self._rate_limited_state = False
                         self._error_handling_lock.release()
-                    output_dict = inputs_dict
-                    output_dict[self.output_name] = result
-                    return ChainResult(output_dict=output_dict)
+                    
+                    return result
 
                 finally:
                     if holds_a_lock:
@@ -129,6 +127,26 @@ async def async_run(self, **inputs_dict: Any) -> ChainResult:
                 if holds_a_lock:
                     holds_a_lock = False
                     self._error_handling_lock.release()
+    
+    async def async_run_multiple(
+        self, *inputs_dict: Dict[str, Any]
+    ) -> List[ChainResult]:
+        prompts = [self.template.format(**inp) for inp in inputs_dict]
+        llm_results = await self._run_with_error_handling(asyncio.create_task(self.llm_runner.run_batch(prompts=prompts)))
+        results = ChainResultList()
+        for llm_result, input_dict in zip(llm_results, inputs_dict):
+            result = ChainResult(input_dict)
+            result.output_dict[self.output_name] = llm_result
+            results.append(result)
+        return results
 
+    async def async_run(self, **inputs_dict: Any) -> ChainResult:
+        handlers = self.llm_runner._get_error_handlers()
+        prompt = self.template.format(**inputs_dict)
+
+        result = ChainResult(output_dict=inputs_dict)
+        result.output_dict[self.output_name] = await self._run_with_error_handling(asyncio.create_task(self.llm_runner.async_run(prompt)))
+        return result
+        
     def __add__(self, other) -> Chain:
         return ChainSequence([self]) + other

src/hyperchain/llm_runners/llm_runner.py

@@ -2,12 +2,15 @@
 from .error_handler import BaseErrorHandler
 from .llm_result import LLMResult
 from typing import List, Any
-
+from asyncio import create_task, gather
 
 class LLMRunner(ABC):
     @abstractmethod
     async def async_run(self, prompt: Any) -> LLMResult:
         pass
 
+    async def run_batch(self, prompts: List[Any]) -> List[LLMResult]:
+        return await gather(*[self.async_run(prompt=prompt) for prompt in prompts])
+
     def _get_error_handlers(self) -> List[BaseErrorHandler]:
         return []

src/hyperchain/llm_runners/masked_model_runner.py

@@ -39,6 +39,20 @@ async def async_run(self, prompt: str):
             response = response.replace('<mask>', prediction[0]['token_str'], 1)
 
         return LLMResult(response, extra_llm_outputs=predictions)
+    
+    async def run_batch(self, prompts):
+        from torch import inference_mode
+        with inference_mode():
+            predictions_list = self.fill_mask(prompts, **self.pipeline_parameters)
+        
+        results = []
+        for prompt, predictions in zip(prompts, predictions_list):
+            response = prompt
+            for prediction in predictions:
+                response = response.replace('<mask>', prediction[0]['token_str'], 1)
+            results.append(LLMResult(response, extra_llm_outputs=predictions))
+
+        return results
 
     def _get_error_handlers(self):
         return []

src/hyperchain/llm_runners/t5_model_runner.py

@@ -9,7 +9,7 @@ def __init__(
         model,
         tokenizer = None,
         model_kwargs = {},
-    ):
+    ):  
         if isinstance(model, str):
             self.model = T5ForConditionalGeneration.from_pretrained(model)
         else:
@@ -60,10 +60,24 @@ def _apply_response(self, prompt, response):
         return result
 
     async def async_run(self, prompt: str):
+        from torch import inference_mode
         input_ids  = self.tokenizer(prompt, return_tensors="pt").input_ids
-        response = self.model.generate(input_ids, **self.model_kwargs)
+        with inference_mode():
+            response = self.model.generate(input_ids, **self.model_kwargs)
         decoded_response = self.tokenizer.decode(self._apply_response(input_ids, response)[0], skip_special_tokens=True)
         return LLMResult(decoded_response, extra_llm_outputs={"input_ids": input_ids, "response": response})
+    
+    async def run_batch(self, prompts):
+        from torch import inference_mode
+        input_ids_batch  = self.tokenizer(prompts, padding=True, return_tensors="pt").input_ids
+        with inference_mode():
+            responses = self.model.generate(input_ids_batch, **self.model_kwargs)
+        decoded_responses = self.tokenizer.batch_decode(self._apply_response(input_ids_batch, responses), skip_special_tokens=True)
+        return [
+            LLMResult(decoded_response, extra_llm_outputs={"input_ids": input_ids, "response": response})
+            for decoded_response, input_ids, response in zip(decoded_responses, input_ids_batch, responses)
+        ]
+
 
     def _get_error_handlers(self):
         return []

tests/llm_runners/test_t5_model_runner.py

@@ -1,5 +1,5 @@
 import pytest
-from unittest.mock import Mock
+from unittest.mock import Mock, MagicMock, patch
 from hyperchain.llm_runners.t5_model_runner import T5ConditionalModelRunner
 
 def test_apply_response():
@@ -18,6 +18,7 @@ def test_apply_response():
     assert result == [[4, 5, 6, 7, 8, 9, 10]]
 
 @pytest.mark.asyncio
+@patch.dict('sys.modules', torch=MagicMock())
 async def test_async_run():
     mocked_model = Mock()
     mocked_tokenizer = Mock()