| | from transformers import AutoTokenizer, AutoModelForCausalLM |
| | import torch |
| | from abs_compressor import AbstractCompressor |
| |
|
| |
|
| | class KiSCompressor(AbstractCompressor): |
| | def __init__(self, DEVICE: str = 'cpu', model_dir: str = 'philippelaban/keep_it_simple'): |
| | self.DEVICE = DEVICE |
| | self.tokenizer = AutoTokenizer.from_pretrained(model_dir, padding_side='right', pad_token='<|endoftext|') |
| | self.tokenizer.pad_token = self.tokenizer.eos_token |
| | self.tokenizer.padding_side = 'right' |
| | self.kis_model = AutoModelForCausalLM.from_pretrained(model_dir) |
| | self.kis_model.to(self.DEVICE) |
| | |
| | |
| | |
| |
|
| | def compress(self, original_prompt: str, ratio: float = 0.5, max_length: int = 150, num_beams: int = 4, do_sample: bool = True, num_return_sequences: int = 1, target_index: int = 0) -> dict: |
| |
|
| | original_tokens = len(self.gpt_tokenizer.encode(original_prompt)) |
| |
|
| | start_id = self.tokenizer.bos_token_id |
| | print(self.tokenizer.padding_side) |
| | tokenized_paragraph = [(self.tokenizer.encode(text=original_prompt) + [start_id])] |
| | input_ids = torch.LongTensor(tokenized_paragraph) |
| | if self.DEVICE == 'cuda': |
| | input_ids = input_ids.type(torch.cuda.LongTensor) |
| | output_ids = self.kis_model.generate(input_ids, max_length=max_length, num_beams=num_beams, do_sample=do_sample, |
| | num_return_sequences=num_return_sequences, |
| | pad_token_id=self.tokenizer.eos_token_id) |
| | output_ids = output_ids[:, input_ids.shape[1]:] |
| | output = self.tokenizer.batch_decode(output_ids) |
| | output = [o.replace(self.tokenizer.eos_token, "") for o in output] |
| | compressed_prompt = output[target_index] |
| |
|
| | compressed_tokens = len(self.gpt_tokenizer.encode(compressed_prompt)) |
| |
|
| | result = { |
| | 'compressed_prompt': compressed_prompt, |
| | 'ratio': compressed_tokens / original_tokens, |
| | 'original_tokens': original_tokens, |
| | 'compressed_tokens': compressed_tokens, |
| | } |
| |
|
| | return result |
| |
|
| |
|
| |
|
