add plotly chart

app.py

@@ -13,6 +13,10 @@ import tokenizers
 from sampling import CAIFSampler, TopKWithTemperatureSampler
 from generator import Generator
 
+import pickle
+
+from plotly import graph_objects as go
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 ATTRIBUTE_MODELS = {
@@ -34,7 +38,7 @@ LANGUAGE_MODELS = {
         'sberbank-ai/rugpt3small_based_on_gpt2',
         "sberbank-ai/rugpt3large_based_on_gpt2"
     ),
-    "English": ("distilgpt2", "gpt2", "EleutherAI/gpt-neo-1.3B")
+    "English": ("gpt2", "distilgpt2", "EleutherAI/gpt-neo-1.3B")
 }
 
 ATTRIBUTE_MODEL_LABEL = {
@@ -65,6 +69,21 @@ PROMPT_EXAMPLE = {
 
 def main():
     st.header("CAIF")
+    with open("entropy_cdf.pkl", "rb") as inp:
+        x_s, y_s = pickle.load(inp)
+    scatter = go.Scatter({
+        "x": x_s,
+        "y": y_s,
+        "name": "GPT2",
+        "mode": "lines",
+    }
+    )
+    layout = go.Layout({
+        "yaxis": {"title": "CAIF step probability"},
+        "xaxis": {"title": "Entropy threshold"},
+        "template": "plotly_white",
+    })
+    figure = go.Figure(data=[scatter], layout=layout)
     language = st.selectbox("Language", ("English", "Russian"))
     cls_model_name = st.selectbox(
         ATTRIBUTE_MODEL_LABEL[language],
@@ -87,6 +106,7 @@ def main():
         target_label_id = 1
     prompt = st.text_input(TEXT_PROMPT_LABEL[language], PROMPT_EXAMPLE[language])
     alpha = st.slider("Alpha", min_value=-10, max_value=10, step=1, value=0)
+    st.plotly_chart(figure)
     entropy_threshold = st.slider("Entropy threshold", min_value=0., max_value=5., step=.1, value=2.)
     auth_token = os.environ.get('TOKEN') or True
     fp16 = st.checkbox("FP16", value=True)
@@ -103,9 +123,6 @@ def main():
     st.subheader("Generated text:")
     st.markdown(text)
 
-
-
-
 @st.cache(hash_funcs={tokenizers.Tokenizer: lambda lm_tokenizer: hash(lm_tokenizer.to_str)}, allow_output_mutation=True)
 def load_generator(lm_model_name: str) -> Generator:
     with st.spinner('Loading language model...'):

requirements.txt

@@ -1,3 +1,4 @@
 streamlit
 transformers
 torch
+plotly