| 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 |
|
|
|
|
|
|
