Can you provide a sample code snippet on how to use?
#2
by
segmond
- opened
Can you provide a sample code snippet on how to use?
I've uploaded some code to run the model over some text.
Use fix_pronouns_in_text from pronoun_fixer.py
from transformers import AutoModelForMaskedLM, AutoTokenizer, FillMaskPipeline
import pronoun_fixer
# text produced by sentence level machine translation where the pronoun was ambiguous in the source language
# and is wrong in the target language
MTL_TEXT = """
Cadence Lee thought he was a normal girl, perhaps a little well to do, but not exceptionally so.
"""
device = 'cuda'
pronoun_checkpoint = "thefrigidliquidation/roberta-base-pronouns"
pronoun_model = AutoModelForMaskedLM.from_pretrained(pronoun_checkpoint).to(device)
pronoun_tokenizer = AutoTokenizer.from_pretrained(pronoun_checkpoint)
unmasker = FillMaskPipeline(model=pronoun_model, tokenizer=pronoun_tokenizer, device=device, top_k=10)
fixed_text = pronoun_fixer.fix_pronouns_in_text(unmasker, pronoun_tokenizer, MTL_TEXT)
print(fixed_text)
# Cadence Lee thought she was a normal girl, perhaps a little well to do, but not exceptionally so.
# now the pronoun is fixed