fix bug of HFTransformer.generate (#118)

2024-01-30 22:52:43 +08:00
parent 94ba3a1a75
commit c9973b9558
1 changed files with 2 additions and 1 deletions
--- a/lagent/llms/huggingface.py
+++ b/lagent/llms/huggingface.py
@@ -105,7 +105,8 @@ class HFTransformer(BaseModel):
        Returns:
            (a list of/batched) text/chat completion
        """
-        for chunk in self.stream_generate(inputs, do_sample, **kwargs):
+        for status, chunk, _ in self.stream_generate(inputs, do_sample,
+                                                     **kwargs):
            response = chunk
        return response