LanguageMachines/blip2-opt-2.7b

1from typing import Dict, List, Any2# import transformers3# from transformers import AutoTokenizer4# import torch5from datetime import datetime6import torch7 8import logging9logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)10 11import requests12from PIL import Image13from transformers import Blip2Processor, Blip2ForConditionalGeneration14 15 16class EndpointHandler():17    18    def __init__(self, path=""):19        20        self.processor = Blip2Processor.from_pretrained(path)21        self.model = Blip2ForConditionalGeneration.from_pretrained(path, device_map="auto")22 23        self.device = "cuda" if torch.cuda.is_available() else "cpu"24 25        self.model.to(self.device)26 27        logging.info('Model moved to device-' + self.device)28        29        # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")30        # self.model.eval()31        # self.model.to(device=device, dtype=self.torch_dtype)32        33        # self.generate_kwargs = {34        #     'max_new_tokens': 512,35        #     'temperature': 0.0001,36        #     'top_p': 1.0,37        #     'top_k': 0,38        #     'use_cache': True,39        #     'do_sample': True,40        #     'eos_token_id': self.tokenizer.eos_token_id,41        #     'pad_token_id': self.tokenizer.pad_token_id,42        #     "repetition_penalty": 1.143        # }44    45    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:46        """47       data args:48            inputs (:obj: `str` | `PIL.Image` | `np.array`)49            kwargs50      Return:51            A :obj:`list` | `dict`: will be serialized and returned52        """53 54        # streamer = TextIteratorStreamer(55        # self.tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True56        # )57        58        ## Model Parameters59        # self.generate_kwargs['max_new_tokens'] = data['max_new_tokens'] if 'max_new_tokens' in data else self.generate_kwargs['max_new_tokens']60        # self.generate_kwargs['temperature'] = data['temperature'] if 'temperature' in data else self.generate_kwargs['temperature']61        # self.generate_kwargs['top_p'] = data['top_p'] if 'top_p' in data else self.generate_kwargs['top_p']62        # self.generate_kwargs['top_k'] = data['top_k'] if 'top_k' in data else self.generate_kwargs['top_k']63        # self.generate_kwargs['do_sample'] = data['do_sample'] if 'do_sample' in data else self.generate_kwargs['do_sample']64        # self.generate_kwargs['repetition_penalty'] = data['repetition_penalty'] if 'repetition_penalty' in data else self.generate_kwargs['repetition_penalty']65        66        67        ## Prepare the inputs68        # inputs = data.pop("inputs",data)69        # input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids70        # input_ids = input_ids.to(self.model.device)71 72 73        # pip install accelerate74        75        img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg' 76 77        now = datetime.now()78 79        raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')80        81        question = "how many dogs are in the picture?"82        inputs = self.processor(raw_image, question, return_tensors="pt").to(self.device)83        84        out = self.model.generate(**inputs)85        output_text = self.processor.decode(out[0], skip_special_tokens=True)86 87        current = datetime.now()88 89        # encoded_inp = self.tokenizer(inputs, return_tensors='pt', padding=True)90        # for key, value in encoded_inp.items():91        #     encoded_inp[key] = value.to('cuda:0')92 93        ## Invoke the model     94        # with torch.no_grad():95        #     gen_tokens =  self.model.generate(96        #         input_ids=encoded_inp['input_ids'],97        #         attention_mask=encoded_inp['attention_mask'],98        #         **generate_kwargs,99        #     )100 101        # ## Decode using tokenizer102        # decoded_gen = self.tokenizer.batch_decode(gen_tokens, skip_special_tokens=True)        103 104        # with torch.no_grad():105        #     output_ids = self.model.generate(input_ids, **self.generate_kwargs)106        # # Slice the output_ids tensor to get only new tokens107        # new_tokens = output_ids[0, len(input_ids[0]) :]108        # output_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)109        110        return [{"gen_text":output_text, "time_elapsed": str(current-now)}]111