update

app.py

@@ -9,12 +9,13 @@ model = GPT2LMHeadModel.from_pretrained(model_name)
 
 
 class DialogGPT:
-    def __init__(self, tokenizer, model, n_candidate=4, param_lambda=0.1):
+    def __init__(self, tokenizer, model, n_candidate=4, param_lambda=0.10):
         self.tokenizer = tokenizer
         self.model = model
         self.model.eval()
         self.n_candidate = n_candidate
         self.param_lambda = param_lambda
+        self.param_gamma: int = 2
 
     def _calc_single_scores(self, token_ids):
         with torch.inference_mode():
@@ -33,7 +34,7 @@ class DialogGPT:
         # log_likelihood (b, l)
         log_likelihood = logit_at_target
         log_likelihood.masked_fill_(mask_at_pad, 0.0)
-        log_likelihood_per_candidate = log_likelihood.sum(dim=1)
+        log_likelihood_per_candidate = log_likelihood[:, self.param_gamma:].sum(dim=1)
         # normalize by length
         # log_likelihood_per_candidate = log_likelihood_per_candidate / (candidate_token_ids.shape[1] - mask_at_pad.sum(dim=1))
         return log_likelihood_per_candidate
@@ -85,7 +86,7 @@ class DialogGPT:
                 max_time=10,
                 num_return_sequences=self.n_candidate,
                 max_length=512,
-                min_length=2,
+                min_length=4,
                 forced_eos_token_id=self.tokenizer.pad_token_id,
                 return_dict_in_generate=True,
                 output_scores=True,