fredaddy commited on
Commit
e56d8f8
·
verified ·
1 Parent(s): cd5795f

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +3 -3
handler.py CHANGED
@@ -6,18 +6,18 @@ class ModelHandler:
6
  def __init__(self):
7
  # Load the model and tokenizer with appropriate weights
8
  self.model = AutoModel.from_pretrained(
9
- 'openbmb/MiniCPM-V-2_6',
10
  trust_remote_code=True,
11
  attn_implementation='sdpa',
12
  torch_dtype=torch.bfloat16
13
  ).eval().cuda()
14
 
15
- self.tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True)
16
 
17
  def preprocess(self, inputs):
18
  # Preprocess image input
19
  image = Image.open(inputs['image'].file).convert('RGB')
20
- question = inputs.get("question", "What is in the image?")
21
  msgs = [{'role': 'user', 'content': [image, question]}]
22
  return msgs
23
 
 
6
  def __init__(self):
7
  # Load the model and tokenizer with appropriate weights
8
  self.model = AutoModel.from_pretrained(
9
+ 'fredaddy/MiniCPM-V-2_6',
10
  trust_remote_code=True,
11
  attn_implementation='sdpa',
12
  torch_dtype=torch.bfloat16
13
  ).eval().cuda()
14
 
15
+ self.tokenizer = AutoTokenizer.from_pretrained('fredaddy/MiniCPM-V-2_6', trust_remote_code=True)
16
 
17
  def preprocess(self, inputs):
18
  # Preprocess image input
19
  image = Image.open(inputs['image'].file).convert('RGB')
20
+ question = inputs.get("question", "Extract all data in the image. Be extremely careful to ensure that you don't miss anything. It's imperative that you extract and digitize everything on that page.")
21
  msgs = [{'role': 'user', 'content': [image, question]}]
22
  return msgs
23