higher timeout

app.py

@@ -150,7 +150,7 @@ def generate(
     input_ids = input_ids.to(model.device)
 
 
-    streamer = CustomStreamer(skip_prompt=True, timeout=10.0)
+    streamer = CustomStreamer(skip_prompt=True, timeout=1000.0)
     
     def generate_with_states():
         with torch.no_grad():
@@ -173,32 +173,11 @@ def generate(
     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
 
-    # 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)
@@ -219,9 +198,19 @@ def generate(
         
         yield se_highlighted_text, acc_highlighted_text
 
-    thread.join()
-
     #### 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]:]