Update handler.py

handler.py

@@ -101,21 +101,25 @@ class EndpointHandler():
         if not inputs.endswith("<cite>"):
             inputs += "<cite>"
         logger.info(inputs)
-        
         inputs = self.tokenizer(inputs, return_tensors="pt").to("cuda")
         with torch.no_grad():
             outputs = self.model(**inputs)
         outputs_logits = outputs.logits[0, -1, self.law_token_ids]
-        base_logits = outputs.logits[0, -1, self.law_token_ids]
 
+        base_input = tokenizer("<cite>", return_tensors="pt").to("cuda")
+        with torch.no_grad():
+            base_output = self.model(**base_input)
+            
+        base_logits = base_output.logits[0, -1, self.law_token_ids]
         raw_mean = outputs_logits.mean()
         outputs_logits = outputs_logits - base_lambda * base_logits
         outputs_logits = outputs_logits + (raw_mean - outputs_logits.mean())
 
         law_token_probs = outputs_logits.softmax(dim=0)
         sorted_ids = torch.argsort(law_token_probs, descending=True)[:topk]
-        print([self.law_token_names[x] for x in sorted_ids])
+        logger.info([self.law_token_names[x] for x in sorted_ids])
         token_objects = [
             self.law_lookup.get_law_from_token(self.law_token_names[x])
             for x in sorted_ids.tolist()]
+        
         return {"tokens": token_objects}