File size: 3,966 Bytes
834e995 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | import torch
from transformers import AutoTokenizer, PreTrainedModel
from transformers.models.t5.configuration_t5 import T5Config
from .configuration_hyperstyler import HyperStylerConfig
from .hyperstyler import HyperStyler
# Imported here so that Transformers copies these files with the remote code.
from .modeling_t5_hyperstyler import T5ForConditionalGeneration_HyperStyler # noqa: F401
from .config_t5 import T5Config as _T5Config # noqa: F401
from .style_utils import load_style_model, text_to_style
class HyperStylerModel(PreTrainedModel):
config_class = HyperStylerConfig
base_model_prefix = "hyperstyler"
# T5 shares these with hyperstyler.model.shared.weight.
_tied_weights_keys = [
"hyperstyler.model.encoder.embed_tokens.weight",
"hyperstyler.model.decoder.embed_tokens.weight",
]
def __init__(self, config):
super().__init__(config)
t5_config = T5Config.from_pretrained(config.base_model_name)
self.hyperstyler = HyperStyler(
config.base_model_name, config.styler_config, use_style=True, t5_config=t5_config
)
self._tokenizer = None
self._style_model = None
def _init_weights(self, module):
pass
@property
def tokenizer(self):
if self._tokenizer is None:
self._tokenizer = AutoTokenizer.from_pretrained(self.config.base_model_name, legacy=True)
return self._tokenizer
def style_embeddings(self, texts):
"""AnnaWegmann/Style-Embedding vectors for a list of texts: [len(texts), 768]."""
if self._style_model is None:
model, tokenizer, _ = load_style_model(self.config.style_model_name)
self._style_model = (model.to(self.device).eval(), tokenizer)
model, tokenizer = self._style_model
embeds = text_to_style(model=model, tokenizer=tokenizer, texts=texts, device=self.device, model_type="style")
return torch.stack(embeds)
def forward(self, input_ids, attention_mask, labels=None, style_embs=None, style_masks=None, **kwargs):
return self.hyperstyler(input_ids=input_ids, attention_mask=attention_mask, labels=labels,
style_embs=style_embs, style_masks=style_masks, **kwargs)
def generate(self, input_ids, attention_mask, style_embs, style_masks, **kwargs):
return self.hyperstyler.generate(input_ids=input_ids, attention_mask=attention_mask,
style_embs=style_embs, style_masks=style_masks, **kwargs)
@torch.no_grad()
def transfer(self, source_texts, target_texts, temperature=0.8, top_p=1.0, do_sample=True, **generate_kwargs):
"""Rewrite source text(s) in the style of the author of ``target_texts``.
Args:
source_texts: a string or a list of strings.
target_texts: example texts by the target author.
Returns:
a string, or a list of strings if ``source_texts`` is a list.
"""
single = isinstance(source_texts, str)
sources = [source_texts] if single else list(source_texts)
style = self.style_embeddings(list(target_texts)).to(self.device) # [n, 768]
style_embs = style.unsqueeze(0).expand(len(sources), -1, -1) # [B, n, 768]
style_masks = torch.ones(style_embs.shape[:2], device=self.device)
inputs = self.tokenizer(sources, return_tensors="pt", padding="max_length",
max_length=self.config.max_encoder_len, truncation=True).to(self.device)
generate_kwargs.setdefault("max_new_tokens", self.config.max_decoder_len)
out = self.generate(inputs["input_ids"], inputs["attention_mask"], style_embs, style_masks,
do_sample=do_sample, temperature=temperature, top_p=top_p, **generate_kwargs)
texts = self.tokenizer.batch_decode(out, skip_special_tokens=True)
return texts[0] if single else texts
|