gabrielearmento commited on
Commit
655ba26
1 Parent(s): 562c8e2

Create handler.py

Browse files
Files changed (1) hide show
  1. handler.py +23 -0
handler.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import AutoModel, AutoTokenizer
3
+
4
+ class EndpointHandler():
5
+ def __init__(self, path=""):
6
+ # Preload all the elements you are going to need at inference.
7
+ self.model = AutoModel.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5-int4', trust_remote_code=True)
8
+ self.tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-Llama3-V-2_5-int4', trust_remote_code=True)
9
+
10
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
11
+ image_url = data.pop("image_url")
12
+ image = Image.open(image_url).convert("RGB")
13
+ message = data.pop("message")
14
+ messages = [{'role': 'user', 'content': message}]
15
+ return model.chat(
16
+ image=image,
17
+ msgs=msgs,
18
+ tokenizer=self.tokenizer,
19
+ sampling=True, # if sampling=False, beam_search will be used by default
20
+ temperature=0.7,
21
+ # system_prompt='' # pass system_prompt if needed
22
+ )
23
+