Use 'with' statement to better handle streaming memory

2023-03-12 02:04:28 -03:00 · 2023-03-12 02:04:28 -03:00 · 0bd5430988
commit 0bd5430988
parent 37f0166b2d
3 changed files with 38 additions and 18 deletions
--- a/modules/RWKV.py
+++ b/modules/RWKV.py
@ -50,11 +50,11 @@ class RWKVModel:
        return context+self.pipeline.generate(context, token_count=token_count, args=args, callback=callback)

    def generate_with_streaming(self, **kwargs):
-        iterable = Iteratorize(self.generate, kwargs, callback=None)
-        reply = kwargs['context']
-        for token in iterable:
-            reply += token
-            yield reply
+        with Iteratorize(self.generate, kwargs, callback=None) as generator:
+            reply = kwargs['context']
+            for token in generator:
+                reply += token
+                yield reply

 class RWKVTokenizer:
    def __init__(self):