sarang-shrivastava commited on
Commit
9747a5e
·
1 Parent(s): 8c66032

Upload handler and req.txt

Browse files
Files changed (2) hide show
  1. handler (1).py +116 -0
  2. requirements (1).txt +6 -0
handler (1).py ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ # import transformers
3
+ # from transformers import AutoTokenizer
4
+ # import torch
5
+ from datetime import datetime
6
+
7
+ import torch
8
+
9
+ import logging
10
+ logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.DEBUG)
11
+
12
+
13
+ import requests
14
+ from PIL import Image
15
+ from transformers import Blip2Processor, Blip2ForConditionalGeneration
16
+
17
+
18
+ class EndpointHandler():
19
+
20
+ def __init__(self, path=""):
21
+
22
+ self.processor = Blip2Processor.from_pretrained(path)
23
+ self.model = Blip2ForConditionalGeneration.from_pretrained(path, load_in_8bit=True, device_map="auto")
24
+
25
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
26
+
27
+ # self.model.to(self.device)
28
+
29
+ logging.info('Model moved to device-' + self.device)
30
+
31
+ # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
32
+ # self.model.eval()
33
+ # self.model.to(device=device, dtype=self.torch_dtype)
34
+
35
+ # self.generate_kwargs = {
36
+ # 'max_new_tokens': 512,
37
+ # 'temperature': 0.0001,
38
+ # 'top_p': 1.0,
39
+ # 'top_k': 0,
40
+ # 'use_cache': True,
41
+ # 'do_sample': True,
42
+ # 'eos_token_id': self.tokenizer.eos_token_id,
43
+ # 'pad_token_id': self.tokenizer.pad_token_id,
44
+ # "repetition_penalty": 1.1
45
+ # }
46
+
47
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
48
+ """
49
+ data args:
50
+ inputs (:obj: `str` | `PIL.Image` | `np.array`)
51
+ kwargs
52
+ Return:
53
+ A :obj:`list` | `dict`: will be serialized and returned
54
+ """
55
+
56
+ # streamer = TextIteratorStreamer(
57
+ # self.tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
58
+ # )
59
+
60
+ ## Model Parameters
61
+ # self.generate_kwargs['max_new_tokens'] = data['max_new_tokens'] if 'max_new_tokens' in data else self.generate_kwargs['max_new_tokens']
62
+ # self.generate_kwargs['temperature'] = data['temperature'] if 'temperature' in data else self.generate_kwargs['temperature']
63
+ # self.generate_kwargs['top_p'] = data['top_p'] if 'top_p' in data else self.generate_kwargs['top_p']
64
+ # self.generate_kwargs['top_k'] = data['top_k'] if 'top_k' in data else self.generate_kwargs['top_k']
65
+ # self.generate_kwargs['do_sample'] = data['do_sample'] if 'do_sample' in data else self.generate_kwargs['do_sample']
66
+ # self.generate_kwargs['repetition_penalty'] = data['repetition_penalty'] if 'repetition_penalty' in data else self.generate_kwargs['repetition_penalty']
67
+
68
+
69
+ ## Prepare the inputs
70
+ batch_size = data.pop("batch_size",data)
71
+ # input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids
72
+ # input_ids = input_ids.to(self.model.device)
73
+
74
+
75
+ # pip install accelerate
76
+
77
+ img_url = 'https://storage.googleapis.com/sfr-vision-language-research/BLIP/demo.jpg'
78
+
79
+ now = datetime.now()
80
+
81
+ raw_image = Image.open(requests.get(img_url, stream=True).raw).convert('RGB')
82
+
83
+ # question = "how many dogs are in the picture?"
84
+ # inputs = self.processor(raw_image, question, return_tensors="pt").to("cuda")
85
+
86
+ inputs = self.processor([raw_image]*batch_size, return_tensors="pt").to("cuda", torch.float16)
87
+
88
+ out = self.model.generate(**inputs)
89
+
90
+ # generated_text = self.processor.batch_decode(out, skip_special_tokens=True)[0].strip()
91
+ generated_text = self.processor.batch_decode(out, skip_special_tokens=True)
92
+
93
+ current = datetime.now()
94
+
95
+ # encoded_inp = self.tokenizer(inputs, return_tensors='pt', padding=True)
96
+ # for key, value in encoded_inp.items():
97
+ # encoded_inp[key] = value.to('cuda:0')
98
+
99
+ ## Invoke the model
100
+ # with torch.no_grad():
101
+ # gen_tokens = self.model.generate(
102
+ # input_ids=encoded_inp['input_ids'],
103
+ # attention_mask=encoded_inp['attention_mask'],
104
+ # **generate_kwargs,
105
+ # )
106
+
107
+ # ## Decode using tokenizer
108
+ # decoded_gen = self.tokenizer.batch_decode(gen_tokens, skip_special_tokens=True)
109
+
110
+ # with torch.no_grad():
111
+ # output_ids = self.model.generate(input_ids, **self.generate_kwargs)
112
+ # # Slice the output_ids tensor to get only new tokens
113
+ # new_tokens = output_ids[0, len(input_ids[0]) :]
114
+ # output_text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
115
+
116
+ return [{"gen_text":generated_text, "time_elapsed": str(current-now)}]
requirements (1).txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ Pillow
2
+ requests
3
+ accelerate
4
+ torch
5
+ transformers
6
+ bitsandbytes