Add code to use the model

pronoun_fixer.py

@@ -0,0 +1,219 @@
+from tqdm import tqdm
+from transformers import FillMaskPipeline, RobertaTokenizerFast
+
+
+MAX_CTX_LEN = 512
+
+SPACE_PREFIX = 'Ġ'
+
+PRONOUN_TOKENS = {
+    'I', 'ĠI',
+    'you', 'You', 'Ġyou', 'ĠYou',
+    'he', 'He', 'Ġhe', 'ĠHe',
+    'she', 'She', 'Ġshe', 'ĠShe',
+    'it', 'It', 'Ġit', 'ĠIt',
+    'we', 'We', 'Ġwe', 'ĠWe',
+    'they', 'They', 'Ġthey', 'ĠThey',
+    'my', 'My', 'Ġmy', 'ĠMy',
+    'your', 'Your', 'Ġyour', 'ĠYour',
+    'his', 'His', 'Ġhis', 'ĠHis',
+    'her', 'Her', 'Ġher', 'ĠHer',
+    'its', 'Its', 'Ġits', 'ĠIts',
+    'our', 'Our', 'Ġour', 'ĠOur',
+    'their', 'Their', 'Ġtheir', 'ĠTheir',
+    'mine', 'Mine', 'Ġmine', 'ĠMine',
+    'yours', 'Yours', 'Ġyours', 'ĠYours',
+    'hers', 'Hers', 'Ġhers', 'ĠHers',
+    'ours', 'Ours', 'Ġours', 'ĠOurs',
+    'theirs', 'Theirs', 'Ġtheirs', 'ĠTheirs',
+}
+
+
+def count_tokens(tokenizer, text: str) -> int:
+    """ return number of tokens in string """
+    return len(tokenizer(text)['input_ids'])
+
+
+def text_to_token_names(tokenizer, text: str) -> list[str]:
+    inputs = tokenizer(text)
+    ref_tokens = []
+    for id in inputs["input_ids"]:
+        token = tokenizer._convert_id_to_token(id)
+        ref_tokens.append(token)
+    return ref_tokens
+
+
+def text_to_token_ids(tokenizer, text: str) -> list[str]:
+    return tokenizer(text)["input_ids"]
+
+
+def has_at_least_one_pronoun(tokenizer, text: str) -> bool:
+    token_names = text_to_token_names(tokenizer, text)
+    for pronoun_token in PRONOUN_TOKENS:
+        if pronoun_token in token_names:
+            return True
+    return False
+
+
+def chunk_to_contexts(tokenizer, text: str):
+    lines = text.splitlines()
+    
+    for i in range(len(lines)):
+        
+        # add lines before and after for context
+        ctx = [lines[i]]
+        focus_line_idx = 0
+        before_line_idx = i
+        after_line_idx = i
+        
+        # try adding lines as context until we reach MAX_CTX_LEN
+        while True:
+            
+            something_done = False
+            
+            # try adding a line before
+            if before_line_idx - 1 >= 0:
+                before_candidate = [lines[before_line_idx - 1]] + ctx
+                assert len(before_candidate) == len(ctx) + 1
+                if count_tokens(tokenizer, "\n".join(before_candidate)) < MAX_CTX_LEN:
+                    ctx = before_candidate
+                    focus_line_idx += 1
+                    before_line_idx -= 1
+                    something_done = True
+            
+            # try adding a line after
+            if after_line_idx + 1 < len(lines):
+                # after_candidate = ctx + "\n" + lines[after_line_idx + 1]
+                after_candidate = ctx + [lines[after_line_idx + 1]]
+                if count_tokens(tokenizer, "\n".join(after_candidate)) < MAX_CTX_LEN:
+                    ctx = after_candidate
+                    after_line_idx += 1
+                    something_done = True
+            
+            # if we can't add any line, we're done
+            if not something_done:
+                break
+        
+        assert len("".join(f"{x}\n" for x in ctx).splitlines()) == len(ctx)
+        
+        yield "".join(f"{x}\n" for x in ctx), focus_line_idx
+
+
+def mask_pronouns(tokenizer: RobertaTokenizerFast, text: str) -> tuple[str, list[str]]:
+    """ replaces all pronouns in text with <mask> """
+    token_names = text_to_token_names(tokenizer, text)
+    
+    masked_token_names = []
+    original_pronouns = []
+    
+    for token_name in token_names:
+        if token_name in PRONOUN_TOKENS:
+            masked_token_names.append(tokenizer.mask_token)
+            original_pronouns.append(token_name)
+        else:
+            masked_token_names.append(token_name)
+    
+    masked_text = tokenizer.decode(tokenizer.convert_tokens_to_ids(masked_token_names), skip_special_tokens=False)
+    # remove start and end tokens
+    return masked_text[len(tokenizer.bos_token):-len(tokenizer.eos_token)], original_pronouns
+
+
+def uncase_token(token_name: str) -> str:
+    token_name = token_name.replace(' ', '')
+    token_name = token_name.replace(SPACE_PREFIX, '')
+    return token_name.lower()
+
+
+def uncase_mask_result(mask_result):
+    uncased_token_probs = {uncase_token(k): 0 for k in PRONOUN_TOKENS}
+    for guess in mask_result:
+        uncased_token_str = uncase_token(guess['token_str'])
+        if uncased_token_str not in uncased_token_probs:
+            continue
+        uncased_token_probs[uncased_token_str] += guess['score']
+    return uncased_token_probs
+
+
+def case_token_like(best_token_uncased: str, original_token: str, best_token_cased_str: str) -> str:
+    """
+    :param best_token_uncased: the uncased, unspaced token that's the best match
+    :param original_token: the original token we are replacing
+    :param best_token_cased_str: the token str that's the best match. used for some cap
+    :return: 
+    """
+    space = (SPACE_PREFIX == original_token[0]) or (' ' == original_token[0])
+    cap = original_token[1 if space else 0].isupper()
+    if best_token_uncased == 'i':
+        cap = True
+    # if the original token was 'I', we can't use it for cap info
+    if original_token in ['I', 'ĠI']:
+        cap = best_token_cased_str.strip()[0].isupper()
+    if cap:
+        best_token_uncased = best_token_uncased[0].upper() + best_token_uncased[1:]
+    if space:
+        best_token_uncased = ' ' + best_token_uncased
+    return best_token_uncased
+
+
+def fix_pronouns_in_text(
+        unmasker: FillMaskPipeline,
+        tokenizer,
+        text: str,
+        alpha: float = 0.05,
+        use_tqdm: bool = False,
+        tqdm_kwargs=None
+) -> str:
+    """
+    Fixes pronouns in MTL text
+    :param unmasker: unmasker pipeline
+    :param tokenizer: model tokenizer
+    :param text: text to fix
+    :param alpha: only replace the existing pronouns with probability less than alpha
+    :param use_tqdm: show tqdm progress bar
+    :param tqdm_kwargs: any tqdm args
+    :return: the fixed text
+    """
+    if tqdm_kwargs is None:
+        tqdm_kwargs = {}
+    
+    fixed_lines = []
+    ctxs = list(chunk_to_contexts(tokenizer, text))
+    ctxs_iter = tqdm(ctxs, smoothing=0.0, desc="Fixing pronouns", **tqdm_kwargs) if use_tqdm else ctxs
+    
+    for ctx, focus_line_idx in ctxs_iter:
+        
+        ctx_lines = ctx.splitlines()
+        focus_line = ctx_lines[focus_line_idx]
+        
+        # we can skip focusing on lines without a pronoun
+        if not has_at_least_one_pronoun(tokenizer, focus_line):
+            fixed_lines.append(focus_line)
+            continue
+        
+        # mask all pronouns
+        masked_ctx, original_pronouns = mask_pronouns(tokenizer, ctx)
+        
+        # unmask pronouns
+        mask_results = unmasker(masked_ctx)
+        if isinstance(mask_results[0], dict):
+            mask_results = [mask_results]
+        
+        unmasked_ctx = masked_ctx
+        
+        for i, mask_result in enumerate(mask_results):
+            original_pronoun = original_pronouns[i]
+            uncased_original = uncase_token(original_pronoun)
+            uncased_result = uncase_mask_result(mask_result)
+            
+            # if what was there doesn't make any sense, replace it
+            if uncased_result[uncased_original] < alpha:
+                best_uncased_pronoun = max(uncased_result, key=uncased_result.get)
+                # TODO: ensure correct type, possessive, adj, subject
+                best_cased_pronoun = case_token_like(best_uncased_pronoun, original_pronoun, mask_result[0]['token_str'])
+                unmasked_ctx = unmasked_ctx.replace(tokenizer.mask_token, best_cased_pronoun, 1)
+            else:
+                unmasked_ctx = unmasked_ctx.replace(tokenizer.mask_token, original_pronoun.replace(SPACE_PREFIX, ' '), 1)
+        
+        fixed_lines.append(unmasked_ctx.splitlines()[focus_line_idx])
+    
+    return "\n".join(fixed_lines)