| from io import BytesIO |
| from typing import Dict, Any |
| from transformers import NougatProcessor, VisionEncoderDecoderModel, StoppingCriteria, StoppingCriteriaList |
| from transformers.image_utils import base64 |
| from PIL import Image |
| import torch |
|
|
|
|
| class RunningVarTorch: |
| def __init__(self, L=15, norm=False): |
| self.values = None |
| self.L = L |
| self.norm = norm |
|
|
| def push(self, x: torch.Tensor): |
| assert x.dim() == 1 |
| if self.values is None: |
| self.values = x[:, None] |
| elif self.values.shape[1] < self.L: |
| self.values = torch.cat((self.values, x[:, None]), 1) |
| else: |
| self.values = torch.cat((self.values[:, 1:], x[:, None]), 1) |
|
|
| def variance(self): |
| if self.values is None: |
| return |
| if self.norm: |
| return torch.var(self.values, 1) / self.values.shape[1] |
| else: |
| return torch.var(self.values, 1) |
|
|
|
|
| class StoppingCriteriaScores(StoppingCriteria): |
| def __init__(self, threshold: float = 0.015, window_size: int = 200): |
| super().__init__() |
| self.threshold = threshold |
| self.vars = RunningVarTorch(norm=True) |
| self.varvars = RunningVarTorch(L=window_size) |
| self.stop_inds = defaultdict(int) |
| self.stopped = defaultdict(bool) |
| self.size = 0 |
| self.window_size = window_size |
|
|
| @torch.no_grad() |
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor): |
| last_scores = scores[-1] |
| self.vars.push(last_scores.max(1)[0].float().cpu()) |
| self.varvars.push(self.vars.variance()) |
| self.size += 1 |
| if self.size < self.window_size: |
| return False |
|
|
| varvar = self.varvars.variance() |
| for b in range(len(last_scores)): |
| if varvar[b] < self.threshold: |
| if self.stop_inds[b] > 0 and not self.stopped[b]: |
| self.stopped[b] = self.stop_inds[b] >= self.size |
| else: |
| self.stop_inds[b] = int( |
| min(max(self.size, 1) * 1.15 + 150 + self.window_size, 4095) |
| ) |
| else: |
| self.stop_inds[b] = 0 |
| self.stopped[b] = False |
| return all(self.stopped.values()) and len(self.stopped) > 0 |
|
|
| class EndpointHandler(): |
| def __init__(self, path="facebook/nougat-small") -> None: |
| self.processor = NougatProcessor.from_pretrained(path) |
| self.model = VisionEncoderDecoderModel.from_pretrained(path) |
|
|
| self.device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.model.to(self.device) |
|
|
| def __call__(self, data: Dict[str, Any]) -> str: |
| image = data.pop("inputs", data) |
| image_data = Image.open(BytesIO(base64.b64decode(image))) |
| pixel_values = self.processor(image_data, return_tensors="pt").pixel_values |
|
|
| outputs = self.model.generate( |
| pixel_values.to(self.device), |
| min_length=1, |
| max_length=3584, |
| bad_words_ids=[[self.processor.tokenizer.unk_token_id]], |
| return_dict_in_generate=True, |
| output_scores=True, |
| stopping_criteria=StoppingCriteriaList([StoppingCriteriaScores()]) |
| ) |
|
|
| text = self.processor.batch_decode(outputs[0], skip_special_tokens=True)[0] |
| text = self.processor.post_process_generation(text, fix_markdown=False) |
|
|
| return outputs |
|
|
