hermanhelf commited on
Commit
6cb3558
1 Parent(s): 29f27f5
Files changed (1) hide show
  1. handler.py +13 -0
handler.py CHANGED
@@ -1,11 +1,14 @@
1
  #!/usr/bin/env python3
2
  from typing import Dict, List, Any
 
3
 
4
  class EndpointHandler():
5
  def __init__(self, path=""):
6
  # Preload all the elements you are going to need at inference.
7
  # pseudo:
8
  # self.model= load_model(path)
 
 
9
  pass
10
 
11
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
@@ -19,3 +22,13 @@ class EndpointHandler():
19
 
20
  # pseudo
21
  # self.model(input)
 
 
 
 
 
 
 
 
 
 
 
1
  #!/usr/bin/env python3
2
  from typing import Dict, List, Any
3
+ from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
4
 
5
  class EndpointHandler():
6
  def __init__(self, path=""):
7
  # Preload all the elements you are going to need at inference.
8
  # pseudo:
9
  # self.model= load_model(path)
10
+ self.model = PaliGemmaForConditionalGeneration.from_pretrained(path)
11
+ self.processor = AutoProcessor.from_pretrained(path)
12
  pass
13
 
14
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
 
22
 
23
  # pseudo
24
  # self.model(input)
25
+
26
+ raw_inputs = data.pop("inputs", data)
27
+
28
+ image_file = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg?download=true"
29
+ raw_image = Image.open(requests.get(image_file, stream=True).raw)
30
+ inputs = self.processor(raw_inputs["prompt"], raw_image, return_tensors="pt")
31
+ output = self.model.generate(**inputs, max_new_tokens=20)
32
+ response = processor.decode(output[0], skip_special_tokens=True)
33
+ return response
34
+ # print(processor.decode(output[0], skip_special_tokens=True)[len(prompt):])