updated app.py

app.py

@@ -13,7 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+# print
 """
 Example command with bag of words:
 python examples/run_pplm.py -B space --cond_text "The president" --length 100 --gamma 1.5 --num_iterations 3 --num_samples 10 --stepsize 0.01 --window_length 5 --kl_scale 0.01 --gm_scale 0.95
@@ -608,10 +608,9 @@ def generate_text_pplm(
     last_reps = torch.ones(50257)
     last_reps = last_reps.to(device)
     for i in range_func:
-
         # Get past/probs for current output, except for last word
         # Note that GPT takes 2 inputs: past + current_token
-
+        
         # run model forward to obtain unperturbed
         if past is None and output_so_far is not None:
             last = output_so_far[:, -1:]
@@ -739,7 +738,7 @@ def set_generic_model_params(discrim_weights, discrim_meta):
     DISCRIMINATOR_MODELS_PARAMS['generic'] = meta
 
 
-pretrained_model="gpt2-medium"
+pretrained_model="microsoft/DialoGPT-large"
 cond_text=""
 uncond=False
 num_samples=1
@@ -758,15 +757,15 @@ grad_length=10000
 horizon_length=5
 window_length=0
 decay=False
-gamma=1.5
+gamma=1.0
 gm_scale=0.95
 kl_scale=0.01
 seed=0
 no_cuda=False
 colorama=False
 verbosity="quiet"
-fp="./paper_code/discrim_models/persoothe_classifier.pt"
-model_fp="./paper_code/discrim_models/persoothe_encoder.pt"
+fp="./paper_code/discrim_models/persoothe_classifier.pt" #"/content/drive/Shareddrives/COS_IW04_ZL/COSIW04/Discriminators/3_class_lrggpt_fit_deeper_2/3_PerSoothe_classifier_head_epoch_8.pt"
+model_fp="./paper_code/discrim_models/persoothe_encoder.pt" #None 
 calc_perplexity=False
 is_deep=False
 is_deeper=True
@@ -801,10 +800,7 @@ model = GPT2LMHeadModel.from_pretrained(
     output_hidden_states=True
 )
 if model_fp != None and model_fp != "": 
-    try: 
-        model.load_state_dict(torch.load(model_fp))
-    except:
-        print("Can't load local model")
+    model.load_state_dict(torch.load(model_fp))
 model.to(device)
 model.eval()
 
@@ -817,16 +813,19 @@ for param in model.parameters():
 
 eot_token = "<|endoftext|>"
 
-def get_reply(response, history = "How are you?<|endoftext|>"):
-    if response.endswith(("bye", "Bye", "bye.", "Bye.")):
-        return
-    convo_hist = history + response + eot_token
+def get_reply(response, history = None, in_stepsize = 2.56, in_horizon_length = 5, in_num_iterations = 10, in_top_k = 2):
+    stepsize = in_stepsize
+    horizon_length = int(in_horizon_length)
+    num_iterations = int(in_num_iterations)
+    top_k = int(in_top_k)
+    if response.endswith(("bye", "Bye", "bye.", "Bye.", "bye!", "Bye!")):
+        return "<div class='chatbot'>Chatbot restarted</div>", None
+    convo_hist = (history if history != None else "How are you?<|endoftext|>") + response + eot_token
     # figure out conditioning text
     tokenized_cond_text = tokenizer.encode(
-            convo_hist,
+            eot_token + convo_hist,
             add_special_tokens=False
     )
-
     # generate perturbed texts
 
     # full_text_generation returns:
@@ -861,30 +860,21 @@ def get_reply(response, history = "How are you?<|endoftext|>"):
     )
 
     # iterate through the perturbed texts
-    try:
-        pert_gen_text = tokenizer.decode(pert_gen_tok_texts[0].tolist()[0])
-        convo_hist_split = pert_gen_text.split(eot_token)
-
-        # write some HTML
-        html = "<div class='chatbot'>"
-        for m, msg in enumerate(convo_hist_split):
-            cls = "user" if m%2 == 0 else "bot"
-            html += "<div class='msg {}'> {}</div>".format(cls, msg)
-        html += "</div>"
-
-        if len(convo_hist_split) > 4: convo_hist_split = convo_hist_split[-4:]
-        convo_hist = eot_token.join(convo_hist_split)
-    except:
-        convo_hist_split = history.split(eot_token) + ["*ai has no response*"]
-
-        # write some HTML
-        html = "<div class='chatbot'>"
-        for m, msg in enumerate(convo_hist_split):
-            cls = "user" if m%2 == 0 else "bot"
-            html += "<div class='msg {}'> {}</div>".format(cls, msg)
-        html += "</div>"
-
-        convo_hist = history + "*ai has no response*" + eot_token
+    for i, pert_gen_tok_text in enumerate(pert_gen_tok_texts):
+        try:
+            pert_gen_text = tokenizer.decode(pert_gen_tok_text.tolist()[0])
+            convo_hist_split = pert_gen_text.split(eot_token)
+            html = "<div class='chatbot'>"
+            for m, msg in enumerate(convo_hist_split[1:-1]):
+                cls = "user" if m%2 == 0 else "bot"
+                html += "<div class='msg {}'> {}</div>".format(cls, msg)
+            html += "</div>"
+
+            if len(convo_hist_split) > 4: convo_hist_split = convo_hist_split[-4:]
+            convo_hist = eot_token.join(convo_hist_split)
+
+        except:
+            return "<div class='chatbot'>Error occured, chatbot restarted</div>", None
     
     return html, convo_hist
 
@@ -898,6 +888,11 @@ css = """
 
 gr.Interface(fn=get_reply,
              theme="default",
-             inputs=[gr.inputs.Textbox(placeholder="How are you?"), "state"],
+             inputs=[gr.inputs.Textbox(placeholder="How are you?"), 
+                    "state", 
+                    gr.inputs.Number(default=2.56, label="Step"), 
+                    gr.inputs.Number(default=5, label="Horizon"), 
+                    gr.inputs.Number(default=10, label="Iterations"),
+                    gr.inputs.Number(default=2, label="Top_k")],
              outputs=["html", "state"],
-             css=css).launch() 
+             css=css).launch(share=True)