todos

app.py

@@ -10,6 +10,12 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+# TODO Sentence level highlighting instead (prediction after every word is not what it was trained on). Also solves token-level highlighting issues.
+# TODO log prob output scaling highlighting instead?
+# TODO make it look nicer
+# TODO streaming output (need custom generation function because of probes)
+# TODO add options to switch between models, SLT/TBG, layers?
+# TODO full semantic entropy calculation
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -40,7 +46,6 @@ if torch.cuda.is_available():
     tokenizer.use_default_system_prompt = False
 
     # load the probe data
-    # TODO compare accuracy and SE probe in different tabs/sections
     with open("./model/20240625-131035_demo.pkl", "rb") as f:
         probe_data = pkl.load(f)
     # take the NQ open one
@@ -52,7 +57,6 @@ if torch.cuda.is_available():
 else:
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
 
-
 @spaces.GPU
 def generate(
     message: str,
@@ -62,7 +66,7 @@ def generate(
     top_p: float = 0.9,
     top_k: int = 50,
     repetition_penalty: float = 1.2,
-) -> Iterator[str]:
+) -> Tuple[str, str]:
     conversation = []
     if system_prompt:
         conversation.append({"role": "system", "content": system_prompt})
@@ -74,55 +78,6 @@ 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 = 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 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,