| import torch |
| from transformers import AutoTokenizer, PreTrainedModel |
| from transformers.models.t5.configuration_t5 import T5Config |
|
|
| from .configuration_hyperstyler import HyperStylerConfig |
| from .hyperstyler import HyperStyler |
| |
| from .modeling_t5_hyperstyler import T5ForConditionalGeneration_HyperStyler |
| from .config_t5 import T5Config as _T5Config |
| from .style_utils import load_style_model, text_to_style |
|
|
|
|
| class HyperStylerModel(PreTrainedModel): |
| config_class = HyperStylerConfig |
| base_model_prefix = "hyperstyler" |
| |
| _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) |
| style_embs = style.unsqueeze(0).expand(len(sources), -1, -1) |
| 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 |
|
|
