basestreamer

app.py

@@ -8,7 +8,7 @@ from queue import Queue
 import spaces
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BaseStreamer
 
 
 MAX_MAX_NEW_TOKENS = 2048
@@ -53,19 +53,51 @@ else:
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
 
-class CustomStreamer(TextIteratorStreamer):
-    """
-    Streamer to also store hidden states in a queue.
-    TODO check this works
-    """
-    def __init__(self, tokenizer, skip_prompt: bool = False, skip_special_tokens: bool = False, **decode_kwargs):
-        super().__init__(tokenizer, skip_prompt, skip_special_tokens, **decode_kwargs)
+
+class CustomStreamer(BaseStreamer):
+    def __init__(self, skip_prompt: bool = False, timeout: Optional[float] = None):
+        self.skip_prompt = skip_prompt
+        self.timeout = timeout
+
+        self.token_queue = Queue()
         self.hidden_states_queue = Queue()
+        self.stop_signal = None
+
+        self.next_tokens_are_prompt = True
 
     def put(self, value):
-        if isinstance(value, dict) and 'hidden_states' in value:
-            self.hidden_states_queue.put(value['hidden_states'])
-        super().put(value)
+        """Receives tokens and adds them to the token queue."""
+        if len(value.shape) > 1 and value.shape[0] > 1:
+            raise ValueError("CustomStreamer 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
+
+        for token in value.tolist():
+            self.token_queue.put(token, timeout=self.timeout)
+
+    def put_hidden_states(self, hidden_states):
+        """Receives hidden states and adds them to the hidden states queue."""
+        self.hidden_states_queue.put(hidden_states, timeout=self.timeout)
+
+    def end(self):
+        """Signals the end of the stream."""
+        self.next_tokens_are_prompt = True
+        self.token_queue.put(self.stop_signal, timeout=self.timeout)
+        self.hidden_states_queue.put(self.stop_signal, timeout=self.timeout)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        token = self.token_queue.get(timeout=self.timeout)
+        if token == self.stop_signal:
+            raise StopIteration()
+        else:
+            return token
 
 # Streamer claude
 # def generate(
@@ -116,27 +148,56 @@ def generate(
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     input_ids = input_ids.to(model.device)
 
-    # streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    streamer = CustomStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = dict(
-        input_ids=input_ids,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        repetition_penalty=repetition_penalty,
-        streamer=streamer,
-        output_hidden_states=True,
-        return_dict_in_generate=True,
-    )
-    # with threading
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+
+    streamer = CustomStreamer(skip_prompt=True, timeout=10.0)
+    
+    def generate_with_states():
+        with torch.no_grad():
+            model.generate(
+                input_ids=input_ids,
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                top_p=top_p,
+                top_k=top_k,
+                temperature=temperature,
+                repetition_penalty=repetition_penalty,
+                output_hidden_states=True,
+                return_dict_in_generate=True,
+                streamer=streamer
+            )
+
+    thread = Thread(target=generate_with_states)
     thread.start()
+
     se_highlighted_text = ""
     acc_highlighted_text = ""
-    for new_text in streamer:
+    for token_id in streamer:
         hidden_states = streamer.hidden_states_queue.get()
+        if hidden_states is streamer.stop_signal:
+            break
+
+    # streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+    # streamer = CustomStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    # generation_kwargs = dict(
+    #     input_ids=input_ids,
+    #     max_new_tokens=max_new_tokens,
+    #     do_sample=True,
+    #     top_p=top_p,
+    #     top_k=top_k,
+    #     temperature=temperature,
+    #     repetition_penalty=repetition_penalty,
+    #     streamer=streamer,
+    #     output_hidden_states=True,
+    #     return_dict_in_generate=True,
+    # )
+    # #### with threading
+    # thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    # thread.start()
+    # se_highlighted_text = ""
+    # acc_highlighted_text = ""
+
+    # for new_text in streamer:
+        # hidden_states = streamer.hidden_states_queue.get()
         # Semantic Uncertainty Probe
         token_embeddings = torch.stack([generated_token[0, 0, :].cpu() for generated_token in hidden_states]).numpy()   # (num_layers, hidden_size)
         se_concat_layers = token_embeddings[se_layer_range[0]:se_layer_range[1]].reshape(-1)
@@ -146,6 +207,8 @@ def generate(
         acc_concat_layers = token_embeddings[acc_layer_range[0]:acc_layer_range[1]].reshape(-1)
         acc_probe_pred = (1 - acc_probe.predict_proba(acc_concat_layers.reshape(1, -1))[0][1]) * 2 - 1
         
+        # decode latest token
+        new_test = tokenizer.decode(token_id)
         print(new_text, se_probe_pred, acc_probe_pred)  
 
         se_new_highlighted_text = highlight_text(new_text, se_probe_pred)
@@ -155,25 +218,9 @@ def generate(
         
         yield se_highlighted_text, acc_highlighted_text
 
-        # Semantic Uncertainty Probe
-        # se_token_embeddings = torch.stack([layer[0, -1, :].cpu() for layer in hidden_states])
-        # se_concat_layers = se_token_embeddings.numpy()[se_layer_range[0]:se_layer_range[1]].reshape(-1)
-        # se_probe_pred = se_probe.predict_proba(se_concat_layers.reshape(1, -1))[0][1] * 2 - 1
-        
-        # # Accuracy Probe
-        # acc_token_embeddings = torch.stack([layer[0, -1, :].cpu() for layer in hidden_states])
-        # acc_concat_layers = acc_token_embeddings.numpy()[acc_layer_range[0]:acc_layer_range[1]].reshape(-1)
-        # acc_probe_pred = acc_probe.predict_proba(acc_concat_layers.reshape(1, -1))[0][1] * 2 - 1
-        
-        # se_new_highlighted_text = highlight_text(new_text, se_probe_pred)
-        # acc_new_highlighted_text = highlight_text(new_text, acc_probe_pred)
-        
-        # se_highlighted_text += se_new_highlighted_text
-        # acc_highlighted_text += acc_new_highlighted_text
-        
-        # yield se_highlighted_text, acc_highlighted_text
+    thread.join()
 
-    # Generate without threading
+    #### Generate without threading
     # with torch.no_grad():
     #     outputs = model.generate(**generation_kwargs)
     # generated_tokens = outputs.sequences[0, input_ids.shape[1]:]
@@ -206,7 +253,6 @@ def generate(
     #     se_highlighted_text += f" {se_new_highlighted_text}"
     #     acc_highlighted_text += f" {acc_new_highlighted_text}"
         
-    #     # yield se_highlighted_text, acc_highlighted_text
     # return se_highlighted_text, acc_highlighted_text