add prompt

generator.py

@@ -72,6 +72,14 @@ class Generator:
 
         })
         inp_len = len(input_ids[0])
+        if self.caif_sampler is not None:
+            current_decoded = self.tokenizer.decode(input_ids[0])
+            probs = torch.exp(
+                self.caif_sampler.get_classifier_log_probs(
+                    current_decoded, target_cls_id=sampler_kwargs["target_cls_id"]
+                )
+            ).item()
+            gen_history += [probs]
         for i in range(max_length):
             is_caif_step = (
                 i % caif_period == 0 and self.caif_sampler is not None
@@ -98,7 +106,7 @@ class Generator:
                 scatter_data = go.Scatter({
                     "x": list(range(len(gen_history))),
                     "y": gen_history,
-                    "hovertext": [self.tokenizer.decode(t) for t in input_ids[0][inp_len:]]
+                    "hovertext": ["[PROMPT]"] + [self.tokenizer.decode(t) for t in input_ids[0][inp_len:]]
                 })
                 fig = go.Figure([scatter_data], layout=layout)
                 plot.plotly_chart(fig, use_container_width=True)