"""

Helpers to support streaming generate output.

Borrowed from https://github.com/oobabooga/text-generation-webui/blob/ad37f396fc8bcbab90e11ecf17c56c97bfbd4a9c/modules/callbacks.py

"""

import gc

import traceback

from queue import Queue

from threading import Thread

import torch

import transformers

class Stream(transformers.StoppingCriteria):

    def __init__(self, callback_func=None):

        self.callback_func = callback_func

    def __call__(self, input_ids, scores) -> bool:

        if self.callback_func is not None:

            self.callback_func(input_ids[0])

        return False

class Iteratorize:

    """

    Transforms a function that takes a callback

    into a lazy iterator (generator).

    """

    def __init__(self, func, kwargs={}, callback=None):

        self.mfunc = func

        self.c_callback = callback

        self.q = Queue()

        self.sentinel = object()

        self.kwargs = kwargs

        self.stop_now = False

        def _callback(val):

            if self.stop_now:

                raise ValueError

            self.q.put(val)

        def gentask():

            try:

                ret = self.mfunc(callback=_callback, **self.kwargs)

            except ValueError:

                pass

            except:

                traceback.print_exc()

                pass

            self.q.put(self.sentinel)

            if self.c_callback:

                self.c_callback(ret)

        self.thread = Thread(target=gentask)

        self.thread.start()

    def __iter__(self):

        return self

    def __next__(self):

        obj = self.q.get(True, None)

        if obj is self.sentinel:

            raise StopIteration

        else:

            return obj

    def __enter__(self):

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):

        self.stop_now = True