cleanup

Author: vpj

python_autocomplete/evaluate/__init__.py

@@ -19,12 +19,9 @@ def __call__(self, text, token_str: str):
 
 
 class NextWordPredictionComplete(PredictionComplete):
-    def __init__(self, prompt: str, rest: str, min_length: int):
+    def __init__(self, rest: str, min_length: int):
         self.min_length = min_length
         self.rest = rest
-        self.is_id = False
-        if prompt and prompt[-1] in ID_CHARS:
-            self.is_id = True
 
     def __call__(self, text, token_str: str):
         if len(text) - len(self.rest) < self.min_length:
@@ -159,7 +156,7 @@ def update(self, next_token, itos: List[str], state, old_state):
                         break
                     # if not self.add_prediction(self.probs[b] * tokens[token].item(), b, token_str, state):
                     #     break
-                elif not self.add_beam(self.probs[b] * tokens[token].item(), b, token):
+                if not self.add_beam(self.probs[b] * tokens[token].item(), b, token):
                     break
 
 

python_autocomplete/evaluate/eval_sample.py

@@ -18,7 +18,7 @@ def evaluate(predictor: Predictor, text: str):
         prefix = text[:i + 1]
         stripped, prompt = predictor.rstrip(prefix)
         rest = prefix[len(stripped):]
-        prediction_complete = NextWordPredictionComplete(stripped, rest, 5)
+        prediction_complete = NextWordPredictionComplete(rest, 5)
         prompt = torch.tensor(prompt, dtype=torch.long).unsqueeze(-1)
 
         predictions = predictor.get_next_word(prompt, None, rest, [1.], prediction_complete, 5)