remove streaming

app.py

@@ -54,80 +54,6 @@ else:
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
 
-
-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):
-        """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(
-#     message: str,
-#     system_prompt: str,
-#     chat_history: List[Tuple[str, str]],
-#     max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
-#     temperature: float = 0.6,
-#     top_p: float = 0.9,
-#     top_k: int = 50,
-#     repetition_penalty: float = 1.2,
-# ) -> Iterator[Tuple[str, str]]:
-#     conversation = []
-#     if system_prompt:
-#         conversation.append({"role": "system", "content": system_prompt})
-#     for user, assistant in chat_history:
-#         conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
-#     conversation.append({"role": "user", "content": message})
-
-#     input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
-#     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-#         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-#         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-#     input_ids = input_ids.to(model.device)
-
-
-
-
-
 @spaces.GPU
 def generate(
     message: str,
@@ -150,36 +76,82 @@ def generate(
     input_ids = input_ids.to(model.device)
 
 
-    streamer = CustomStreamer(skip_prompt=True, timeout=1000.0)
+    # 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
-            )
+    # 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()
 
-    thread = Thread(target=generate_with_states)
-    thread.start()
+    # se_highlighted_text = ""
+    # acc_highlighted_text = ""
+    # for token_id in streamer:
+    #     print
+    #     hidden_states = streamer.hidden_states_queue.get()
+    #     if hidden_states is streamer.stop_signal:
+    #         break
+
+    #     # 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)
+    #     se_probe_pred = se_probe.predict_proba(se_concat_layers.reshape(1, -1))[0][1] * 2 - 1
+        
+    #     # Accuracy Probe
+    #     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_text = tokenizer.decode(token_id)
+    #     print(new_text, se_probe_pred, acc_probe_pred)  
+
+    #     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 += f" {se_new_highlighted_text}"
+    #     acc_highlighted_text += f" {acc_new_highlighted_text}"
+        
+    #     yield se_highlighted_text, acc_highlighted_text
+
+    #### Generate without threading
+    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 torch.no_grad():
+        outputs = model.generate(**generation_kwargs)
+    generated_tokens = outputs.sequences[0, input_ids.shape[1]:]
+    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+    print(generated_text)
+    # hidden states
+    hidden = outputs.hidden_states  # list of tensors, one for each token, then (batch size, sequence length, hidden size)
 
     se_highlighted_text = ""
     acc_highlighted_text = ""
-    for token_id in streamer:
-        print(token_id)
-        hidden_states = streamer.hidden_states_queue.get()
-        if hidden_states is streamer.stop_signal:
-            break
+
+    for i in range(1, len(hidden)):
 
         # Semantic Uncertainty Probe
-        token_embeddings = torch.stack([generated_token[0, 0, :].cpu() for generated_token in hidden_states]).numpy()   # (num_layers, hidden_size)
+        token_embeddings = torch.stack([generated_token[0, 0, :].cpu() for generated_token in hidden[i]]).numpy()   # (num_layers, hidden_size)
         se_concat_layers = token_embeddings[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
         
@@ -187,63 +159,16 @@ 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_text = tokenizer.decode(token_id)
-        print(new_text, se_probe_pred, acc_probe_pred)  
+        output_id = outputs.sequences[0, input_ids.shape[1]+i]
+        output_word = tokenizer.decode(output_id)
+        print(output_id, output_word, se_probe_pred, acc_probe_pred)  
 
-        se_new_highlighted_text = highlight_text(new_text, se_probe_pred)
-        acc_new_highlighted_text = highlight_text(new_text, acc_probe_pred)
+        se_new_highlighted_text = highlight_text(output_word, se_probe_pred)
+        acc_new_highlighted_text = highlight_text(output_word, acc_probe_pred)
         se_highlighted_text += f" {se_new_highlighted_text}"
         acc_highlighted_text += f" {acc_new_highlighted_text}"
         
-        yield se_highlighted_text, acc_highlighted_text
-
-    #### Generate without threading
-    # 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 torch.no_grad():
-    #     outputs = model.generate(**generation_kwargs)
-    # generated_tokens = outputs.sequences[0, input_ids.shape[1]:]
-    # generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
-    # # hidden states
-    # hidden = outputs.hidden_states  # list of tensors, one for each token, then (batch size, sequence length, hidden size)
-
-    # # TODO do this loop on the fly instead of waiting for the whole generation
-    # se_highlighted_text = ""
-    # acc_highlighted_text = ""
-
-    # for i in range(1, len(hidden)):
-
-    #     # Semantic Uncertainty Probe
-    #     token_embeddings = torch.stack([generated_token[0, 0, :].cpu() for generated_token in hidden[i]]).numpy()   # (num_layers, hidden_size)
-    #     se_concat_layers = token_embeddings[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 = 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
-        
-    #     output_id = outputs.sequences[0, input_ids.shape[1]+i]
-    #     output_word = tokenizer.decode(output_id)
-    #     print(output_id, output_word, se_probe_pred, acc_probe_pred)  
-
-    #     se_new_highlighted_text = highlight_text(output_word, se_probe_pred)
-    #     acc_new_highlighted_text = highlight_text(output_word, acc_probe_pred)
-    #     se_highlighted_text += f" {se_new_highlighted_text}"
-    #     acc_highlighted_text += f" {acc_new_highlighted_text}"
-        
-    # return se_highlighted_text, acc_highlighted_text
+    return se_highlighted_text, acc_highlighted_text