Update app.py

app.py

@@ -6,15 +6,13 @@ import os
 import random
 import time
 from functools import partial
-from threading import Thread
 
 import gradio as gr
 import nncore
-import spaces
 import torch
 from huggingface_hub import snapshot_download
-from transformers import TextIteratorStreamer
 
+import spaces
 from videomind.constants import GROUNDER_PROMPT, PLANNER_PROMPT, VERIFIER_PROMPT
 from videomind.dataset.utils import process_vision_info
 from videomind.model.builder import build_model
@@ -63,43 +61,6 @@ function init() {
 """
 
 
-class CustomStreamer(TextIteratorStreamer):
-
-    def put(self, value):
-        if len(value.shape) > 1 and value.shape[0] > 1:
-            raise ValueError('TextStreamer only supports batch size 1')
-        elif len(value.shape) > 1:
-            value = value[0]
-
-        if self.skip_prompt and self.next_tokens_are_prompt:
-            self.next_tokens_are_prompt = False
-            return
-
-        self.token_cache.extend(value.tolist())
-
-        # force skipping eos token
-        if self.token_cache[-1] == self.tokenizer.eos_token_id:
-            self.token_cache = self.token_cache[:-1]
-
-        text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs)
-
-        # cache decoded text for future use
-        self.text_cache = text
-
-        if text.endswith('\n'):
-            printable_text = text[self.print_len:]
-            self.token_cache = []
-            self.print_len = 0
-        elif len(text) > 0 and self._is_chinese_char(ord(text[-1])):
-            printable_text = text[self.print_len:]
-            self.print_len += len(printable_text)
-        else:
-            printable_text = text[self.print_len:text.rfind(' ') + 1]
-            self.print_len += len(printable_text)
-
-        self.on_finalized_text(printable_text)
-
-
 def seconds_to_hms(seconds):
     hours, remainder = divmod(round(seconds), 3600)
     minutes, seconds = divmod(remainder, 60)
@@ -128,7 +89,7 @@ def reset_components():
 
 
 @spaces.GPU
-def main(video, prompt, role, temperature, max_new_tokens, model, processor, streamer, device):
+def main(video, prompt, role, temperature, max_new_tokens, model, processor, device):
     history = []
 
     if not video:
@@ -204,9 +165,8 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, str
         model.base_model.enable_adapter_layers()
         model.set_adapter('planner')
 
-        generation_kwargs = dict(
+        output_ids = model.generate(
             **data,
-            streamer=streamer,
             do_sample=temperature > 0,
             temperature=temperature if temperature > 0 else None,
             top_p=None,
@@ -214,15 +174,18 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, str
             repetition_penalty=None,
             max_new_tokens=max_new_tokens)
 
-        t = Thread(target=model.generate, kwargs=generation_kwargs)
-        t.start()
+        assert data.input_ids.size(0) == output_ids.size(0) == 1
+        output_ids = output_ids[0, data.input_ids.size(1):]
+        if output_ids[-1] == processor.tokenizer.eos_token_id:
+            output_ids = output_ids[:-1]
+        response = processor.decode(output_ids, clean_up_tokenization_spaces=False)
 
-        skipped = False
-        for i, text in enumerate(streamer):
-            if text and not skipped:
+        for i, text in enumerate(response.split(' ')):
+            if i == 0:
                 history[-1]['content'] = history[-1]['content'].rstrip('.')
-                skipped = True
-            history[-1]['content'] += text
+                history[-1]['content'] += text
+            else:
+                history[-1]['content'] += ' ' + text
             yield history
 
         elapsed_time = round(time.perf_counter() - start_time, 1)
@@ -230,7 +193,7 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, str
         yield history
 
         try:
-            parsed = json.loads(streamer.text_cache)
+            parsed = json.loads(response)
             action = parsed[0] if isinstance(parsed, list) else parsed
             if action['type'].lower() == 'grounder' and action['value']:
                 query = action['value']
@@ -301,9 +264,8 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, str
         model.base_model.enable_adapter_layers()
         model.set_adapter('grounder')
 
-        generation_kwargs = dict(
+        output_ids = model.generate(
             **data,
-            streamer=streamer,
             do_sample=temperature > 0,
             temperature=temperature if temperature > 0 else None,
             top_p=None,
@@ -311,15 +273,18 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, str
             repetition_penalty=None,
             max_new_tokens=max_new_tokens)
 
-        t = Thread(target=model.generate, kwargs=generation_kwargs)
-        t.start()
+        assert data.input_ids.size(0) == output_ids.size(0) == 1
+        output_ids = output_ids[0, data.input_ids.size(1):]
+        if output_ids[-1] == processor.tokenizer.eos_token_id:
+            output_ids = output_ids[:-1]
+        response = processor.decode(output_ids, clean_up_tokenization_spaces=False)
 
-        skipped = False
-        for i, text in enumerate(streamer):
-            if text and not skipped:
+        for i, text in enumerate(response.split(' ')):
+            if i == 0:
                 history[-1]['content'] = history[-1]['content'].rstrip('.')
-                skipped = True
-            history[-1]['content'] += text
+                history[-1]['content'] += text
+            else:
+                history[-1]['content'] += ' ' + text
             yield history
 
         elapsed_time = round(time.perf_counter() - start_time, 1)
@@ -520,9 +485,8 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, str
         data = data.to(device)
 
         with model.disable_adapter():
-            generation_kwargs = dict(
+            output_ids = model.generate(
                 **data,
-                streamer=streamer,
                 do_sample=temperature > 0,
                 temperature=temperature if temperature > 0 else None,
                 top_p=None,
@@ -530,25 +494,28 @@ def main(video, prompt, role, temperature, max_new_tokens, model, processor, str
                 repetition_penalty=None,
                 max_new_tokens=max_new_tokens)
 
-            t = Thread(target=model.generate, kwargs=generation_kwargs)
-            t.start()
+        assert data.input_ids.size(0) == output_ids.size(0) == 1
+        output_ids = output_ids[0, data.input_ids.size(1):]
+        if output_ids[-1] == processor.tokenizer.eos_token_id:
+            output_ids = output_ids[:-1]
+        response = processor.decode(output_ids, clean_up_tokenization_spaces=False)
 
-            skipped = False
-            for i, text in enumerate(streamer):
-                if text and not skipped:
-                    history[-1]['content'] = history[-1]['content'].rstrip('.')
-                    skipped = True
+        for i, text in enumerate(response.split(' ')):
+            if i == 0:
+                history[-1]['content'] = history[-1]['content'].rstrip('.')
                 history[-1]['content'] += text
-                yield history
+            else:
+                history[-1]['content'] += ' ' + text
+            yield history
 
         elapsed_time = round(time.perf_counter() - start_time, 1)
         history[-1]['metadata']['title'] += f' ({elapsed_time} seconds)'
         yield history
 
         if 'gnd' in role and do_grounding:
-            response = f'After zooming in and analyzing the target moment, I finalize my answer: <span style="color:green">{streamer.text_cache}</span>'
+            response = f'After zooming in and analyzing the target moment, I finalize my answer: <span style="color:green">{response}</span>'
         else:
-            response = f'After watching the whole video, my answer is: <span style="color:green">{streamer.text_cache}</span>'
+            response = f'After watching the whole video, my answer is: <span style="color:green">{response}</span>'
 
         history.append({'role': 'assistant', 'content': ''})
         for i, text in enumerate(response.split(' ')):
@@ -572,11 +539,9 @@ if __name__ == '__main__':
     print('Initializing role *verifier*')
     model.load_adapter(nncore.join(MODEL, 'verifier'), adapter_name='verifier')
 
-    streamer = CustomStreamer(processor.tokenizer, skip_prompt=True)
-
     device = next(model.parameters()).device
 
-    main = partial(main, model=model, processor=processor, streamer=streamer, device=device)
+    main = partial(main, model=model, processor=processor, device=device)
 
     path = os.path.dirname(os.path.realpath(__file__))