Skip to content

Add streaming functionality ! #6

@Sylvester-Anthony

Description

@Sylvester-Anthony

Hello Team ,
I used the package for one of our use cases while working with a RAG application. I kinda figured out that streaming was not really possible with the current code, so I implemented an additional function on my own. I just wanted ask your feedback on this and would you add this as a feature.

The code:

class StreamableContextCiter(ContextCiter):
   @property
    def _output(self):
        if self._cache.get("output") is None:
            streamer = transformers.TextIteratorStreamer(
                self.tokenizer, skip_special_tokens=True, skip_prompt=True
            )
            prompt_ids, prompt = self._get_prompt_ids(return_prompt=True)
            input_ids = torch.tensor([prompt_ids], device=self.model.device)

            def generate_tokens():
                self.model.generate(
                    input_ids=input_ids,
                    streamer=streamer,
                    **self.generate_kwargs,
                )

            thread = threading.Thread(target=generate_tokens)
            thread.start()
            generated_output = prompt
            for token in streamer:
                generated_output += token
                self.streamer_queue.put(token)
            thread.join()
            self._cache["output"] = generated_output
        return self._cache["output"]

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions