|
class Styleformer(): |
|
|
|
def __init__( |
|
self, |
|
style=0, |
|
ctf_model_tag="prithivida/informal_to_formal_styletransfer", |
|
ftc_model_tag="prithivida/formal_to_informal_styletransfer", |
|
atp_model_tag="prithivida/active_to_passive_styletransfer", |
|
pta_model_tag="prithivida/passive_to_active_styletransfer", |
|
adequacy_model_tag="prithivida/parrot_adequacy_model", |
|
): |
|
from transformers import AutoTokenizer |
|
from transformers import AutoModelForSeq2SeqLM |
|
from styleformer import Adequacy |
|
|
|
self.style = style |
|
self.adequacy = adequacy_model_tag and Adequacy(model_tag=adequacy_model_tag) |
|
self.model_loaded = False |
|
|
|
if self.style == 0: |
|
self.ctf_tokenizer = AutoTokenizer.from_pretrained(ctf_model_tag, use_auth_token=False) |
|
self.ctf_model = AutoModelForSeq2SeqLM.from_pretrained(ctf_model_tag, use_auth_token=False) |
|
print("Casual to Formal model loaded...") |
|
self.model_loaded = True |
|
elif self.style == 1: |
|
self.ftc_tokenizer = AutoTokenizer.from_pretrained(ftc_model_tag, use_auth_token=False) |
|
self.ftc_model = AutoModelForSeq2SeqLM.from_pretrained(ftc_model_tag, use_auth_token=False) |
|
print("Formal to Casual model loaded...") |
|
self.model_loaded = True |
|
elif self.style == 2: |
|
self.atp_tokenizer = AutoTokenizer.from_pretrained(atp_model_tag, use_auth_token=False) |
|
self.atp_model = AutoModelForSeq2SeqLM.from_pretrained(atp_model_tag, use_auth_token=False) |
|
print("Active to Passive model loaded...") |
|
self.model_loaded = True |
|
elif self.style == 3: |
|
self.pta_tokenizer = AutoTokenizer.from_pretrained(pta_model_tag, use_auth_token=False) |
|
self.pta_model = AutoModelForSeq2SeqLM.from_pretrained(pta_model_tag, use_auth_token=False) |
|
print("Passive to Active model loaded...") |
|
self.model_loaded = True |
|
else: |
|
print("Only CTF, FTC, ATP and PTA are supported in the pre-release...stay tuned") |
|
|
|
def transfer(self, input_sentence, quality_filter=0.95, max_candidates=1: |
|
if self.model_loaded: |
|
device = "cpu" |
|
|
|
if self.style == 0: |
|
output_sentence = self._casual_to_formal(input_sentence, device, quality_filter, max_candidates) |
|
return output_sentence |
|
elif self.style == 1: |
|
output_sentence = self._formal_to_casual(input_sentence, device, quality_filter, max_candidates) |
|
return output_sentence |
|
elif self.style == 2: |
|
output_sentence = self._active_to_passive(input_sentence, device) |
|
return output_sentence |
|
elif self.style == 3: |
|
output_sentence = self._passive_to_active(input_sentence, device) |
|
return output_sentence |
|
else: |
|
print("Models aren't loaded for this style, please use the right style during init") |
|
|
|
|
|
def _formal_to_casual(self, input_sentence, device, quality_filter, max_candidates): |
|
ftc_prefix = "transfer Formal to Casual: " |
|
src_sentence = input_sentence |
|
input_sentence = ftc_prefix + input_sentence |
|
input_ids = self.ftc_tokenizer.encode(input_sentence, return_tensors='pt') |
|
self.ftc_model = self.ftc_model.to(device) |
|
input_ids = input_ids.to(device) |
|
|
|
preds = self.ftc_model.generate( |
|
input_ids, |
|
do_sample=True, |
|
max_length=32, |
|
top_k=50, |
|
top_p=0.95, |
|
early_stopping=True, |
|
num_return_sequences=max_candidates) |
|
|
|
gen_sentences = set() |
|
for pred in preds: |
|
gen_sentences.add(self.ftc_tokenizer.decode(pred, skip_special_tokens=True).strip()) |
|
|
|
adequacy_scored_phrases = self.adequacy.score(src_sentence, list(gen_sentences), quality_filter, device) |
|
ranked_sentences = sorted(adequacy_scored_phrases.items(), key = lambda x:x[1], reverse=True) |
|
if len(ranked_sentences) > 0: |
|
return ranked_sentences[0][0] |
|
else: |
|
return None |
|
|
|
def _casual_to_formal(self, input_sentence, device, quality_filter, max_candidates): |
|
ctf_prefix = "transfer Casual to Formal: " |
|
src_sentence = input_sentence |
|
input_sentence = ctf_prefix + input_sentence |
|
input_ids = self.ctf_tokenizer.encode(input_sentence, return_tensors='pt') |
|
self.ctf_model = self.ctf_model.to(device) |
|
input_ids = input_ids.to(device) |
|
|
|
preds = self.ctf_model.generate( |
|
input_ids, |
|
do_sample=True, |
|
max_length=32, |
|
top_k=50, |
|
top_p=0.95, |
|
early_stopping=True, |
|
num_return_sequences=max_candidates) |
|
|
|
gen_sentences = set() |
|
for pred in preds: |
|
gen_sentences.add(self.ctf_tokenizer.decode(pred, skip_special_tokens=True).strip()) |
|
|
|
adequacy_scored_phrases = self.adequacy.score(src_sentence, list(gen_sentences), quality_filter, device) |
|
ranked_sentences = sorted(adequacy_scored_phrases.items(), key = lambda x:x[1], reverse=True) |
|
if len(ranked_sentences) > 0: |
|
return ranked_sentences[0][0] |
|
else: |
|
return None |
|
|
|
def _active_to_passive(self, input_sentence, device): |
|
atp_prefix = "transfer Active to Passive: " |
|
src_sentence = input_sentence |
|
input_sentence = atp_prefix + input_sentence |
|
input_ids = self.atp_tokenizer.encode(input_sentence, return_tensors='pt') |
|
self.atp_model = self.atp_model.to(device) |
|
input_ids = input_ids.to(device) |
|
|
|
preds = self.atp_model.generate( |
|
input_ids, |
|
do_sample=True, |
|
max_length=32, |
|
top_k=50, |
|
top_p=0.95, |
|
early_stopping=True, |
|
num_return_sequences=1) |
|
|
|
return self.atp_tokenizer.decode(preds[0], skip_special_tokens=True).strip() |
|
|
|
def _passive_to_active(self, input_sentence, device): |
|
pta_prefix = "transfer Passive to Active: " |
|
src_sentence = input_sentence |
|
input_sentence = pta_prefix + input_sentence |
|
input_ids = self.pta_tokenizer.encode(input_sentence, return_tensors='pt') |
|
self.pta_model = self.pta_model.to(device) |
|
input_ids = input_ids.to(device) |
|
|
|
preds = self.pta_model.generate( |
|
input_ids, |
|
do_sample=True, |
|
max_length=32, |
|
top_k=50, |
|
top_p=0.95, |
|
early_stopping=True, |
|
num_return_sequences=1) |
|
|
|
return self.pta_tokenizer.decode(preds[0], skip_special_tokens=True).strip() |
|
|
|
|
|
|