| from SCRL_new.scrl.model import load_model |
| from transformers import AutoTokenizer |
| import re |
| from abs_compressor import AbstractCompressor |
|
|
|
|
| class SCRLCompressor(AbstractCompressor): |
|
|
| def __init__(self, model_dir: str, device: str = "cpu", tokenizer_dir: str = "sentence-transformers/paraphrase-distilroberta-base-v2"): |
| self.model_dir = model_dir |
| self.device = device |
| self.model = load_model(self.model_dir, self.device) |
| self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir) |
|
|
| def compress(self, original_prompt: str, ratio: float = 0.5, max_length: int = 256) -> dict: |
| original_tokens = len(self.gpt_tokenizer.encode(original_prompt)) |
|
|
| |
| sources = re.findall(r'.{%d}' % max_length, original_prompt.strip()) |
| |
| if sources: |
| summaries = self.model.predict(sources, self.tokenizer, self.device) |
| |
| |
|
|
| compressed_prompt = "" |
| for s in summaries: |
| compressed_prompt += s |
|
|
| 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 |
| else: |
| result = { |
| 'compressed_prompt': "", |
| 'ratio': 0, |
| 'original_tokens': "", |
| 'compressed_tokens': "", |
| } |
| return result |
|
|
|
|
