| from parlai.core.opt import Opt |
| from parlai.utils.typing import TShared |
| from parlai.agents.transformer.transformer import TransformerGeneratorAgent |
|
|
| from .generation_methods import VocabTopKSampling, RerankedTopKSampling |
| from .generation_utils import Wordlist, Reranker, load_wordlist, cefr_to_int |
|
|
| class ControllableBlender(TransformerGeneratorAgent): |
| def __init__(self, opt: Opt, shared: TShared = None): |
| super().__init__(opt, shared) |
|
|
| if opt.get("inference", None) == "vocab": |
| wordlist_path = opt.get("wordlist_path", None) |
| assert wordlist_path, "Please provide path to vocab list, in order to use inference method 'vocab'" |
|
|
| allowed_words = load_wordlist(wordlist_path) |
| self.wordlist = Wordlist(allowed_words, self.dict) |
|
|
| elif opt.get("inference", None) == "rerank": |
| cefr = opt.get("rerank_cefr", None) |
| assert cefr, "Please provide CEFR level, in order to use inference method 'rerank'" |
|
|
| rerank_tokenizer = opt.get("rerank_tokenizer", None) |
| rerank_model = opt.get("rerank_model", None) |
| assert rerank_model, "Please provide path to directory containing model weights, in order to use inference method 'rerank'" |
|
|
| device = opt.get("complexity_model_device", None) |
| penalty_stddev = opt.get("penalty_stddev", None) |
| text_truncate = opt.get("text_truncate", None) |
|
|
| word_filter = None |
| filter_path = opt.get("filter_path", "") |
| if filter_path: |
| word_filter = load_wordlist(filter_path) |
|
|
| exempt_tokens = [self.dict.tok2ind.get(self.dict.null_token), |
| self.dict.tok2ind.get(self.dict.start_token), |
| self.dict.tok2ind.get(self.dict.end_token), |
| self.dict.tok2ind.get(self.dict.unk_token)] |
|
|
| if penalty_stddev < 0: |
| exempt_tokens = "all" |
|
|
| self.reranker = Reranker(cefr=cefr_to_int(cefr), |
| model=rerank_model, |
| tokenizer=rerank_tokenizer, |
| device=device, |
| text_truncate=text_truncate, |
| exempt_tokens=exempt_tokens, |
| penalty_stddev=penalty_stddev, |
| vocab_size=len(self.dict), |
| word_filter=word_filter) |
|
|
| else: |
| raise ValueError(f"Inference method {opt.get('inference', None)} does not exist. " |
| f"Please use 'vocab' or 'rerank'.") |
|
|
|
|
| def _treesearch_factory(self, device, verbose=False): |
| method = self.opt.get('inference', 'greedy') |
| beam_size = self.opt.get('beam_size', 1) |
| if method == 'vocab': |
| return VocabTopKSampling( |
| k=self.opt.get('topk', 40), |
| wordlist=self.wordlist, |
| beam_size=beam_size, |
| min_length=self.beam_min_length, |
| block_ngram=self.beam_block_ngram, |
| context_block_ngram=self.beam_context_block_ngram, |
| length_penalty=self.opt.get('beam_length_penalty', 0.65), |
| padding_token=self.NULL_IDX, |
| bos_token=self.START_IDX, |
| eos_token=self.END_IDX, |
| device=device, |
| verbose=verbose, |
| ) |
| elif method == "rerank": |
| return RerankedTopKSampling( |
| k=self.opt.get('topk', 40), |
| reranker=self.reranker, |
| tokenids_to_text=self._v2t, |
| beam_size=beam_size, |
| min_length=self.beam_min_length, |
| block_ngram=self.beam_block_ngram, |
| context_block_ngram=self.beam_context_block_ngram, |
| length_penalty=self.opt.get('beam_length_penalty', 0.65), |
| padding_token=self.NULL_IDX, |
| bos_token=self.START_IDX, |
| eos_token=self.END_IDX, |
| device=device, |
| verbose=verbose, |
| ) |
| else: |
| return super()._treesearch_factory(device, verbose=verbose) |
|
|
| def share(self): |
| """ |
| Share internal states between parent and child instances. |
| """ |
| shared = super().share() |
| if hasattr(self, 'wordlist'): |
| shared['wordlist'] = self.wordlist |
| if hasattr(self, 'reranker'): |
| shared['reranker'] = self.reranker |
| return shared |
|
|
|
|
|
|
