fix

generator.py

@@ -4,6 +4,8 @@ import torch
 import transformers
 import streamlit as st
 
+from plotly import graph_objects as go
+
 
 class Generator:
     def __init__(self, lm_model_name, device, entropy=None):
@@ -55,6 +57,8 @@ class Generator:
             num_samples,
         )
         text = st.empty()
+        plot = st.empty()
+        gen_history = []
         for i in range(max_length):
             is_caif_step = (
                 i % caif_period == 0 and self.caif_sampler is not None
@@ -70,7 +74,20 @@ class Generator:
             progress_bar.progress((i+1)/max_length)
             if ended_sequences.all():
                 break
-            text.text(self.tokenizer.decode(input_ids[0]))
+            current_decoded = self.tokenizer.decode(input_ids[0])
+            if self.caif_sampler is not None:
+                probs = torch.exp(
+                    self.caif_sampler.get_classifier_log_probs(
+                        current_decoded, target_cls_id=sampler_kwargs["target_cls_id"]
+                    )
+                ).item()
+                gen_history += [probs]
+                scatter_data = go.Scatter({
+                    "x": list(range(len(gen_history))),
+                    "y": gen_history
+                })
+                plot.plotly_chart(scatter_data, use_container_width=True)
+            text.text(current_decoded)
 
         return (
             [

sampling.py

@@ -1,3 +1,5 @@
+from typing import List
+
 import torch
 from torch.nn import functional as F