Upload modeling_OneChart.py

modeling_OneChart.py

@@ -393,8 +393,10 @@ class OneChartOPTForCausalLM(OPTForCausalLM):
         setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
 
     def chat(self, tokenizer, image_file, reliable_check=True, print_prompt=False):
-        dtype=torch.bfloat16
-        device="cuda"
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        # dtype = torch.bfloat16 if device=="cuda" else next(self.get_model().parameters()).dtype
+        dtype=torch.float16 if device=="cuda" else torch.float32
+        # print(device, dtype)
         def list_json_value(json_dict):
             rst_str = []
             sort_flag = True
@@ -456,17 +458,29 @@ class OneChartOPTForCausalLM(OPTForCausalLM):
         stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
         streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-        with torch.autocast(device, dtype=dtype):
+        if device=='cuda':
+            with torch.autocast(device, dtype=dtype):
+                output_ids = self.generate(
+                    input_ids,
+                    images=[image_tensor_1.unsqueeze(0)],
+                    do_sample=False,
+                    num_beams = 1,
+                    # no_repeat_ngram_size = 20,
+                    # streamer=streamer,
+                    max_new_tokens=4096,
+                    stopping_criteria=[stopping_criteria]
+                    )
+        else:
             output_ids = self.generate(
-                input_ids,
-                images=[image_tensor_1.unsqueeze(0).half()],
-                do_sample=False,
-                num_beams = 1,
-                # no_repeat_ngram_size = 20,
-                # streamer=streamer,
-                max_new_tokens=4096,
-                stopping_criteria=[stopping_criteria]
-                )
+                    input_ids,
+                    images=[image_tensor_1.unsqueeze(0)],
+                    do_sample=False,
+                    num_beams = 1,
+                    # no_repeat_ngram_size = 20,
+                    # streamer=streamer,
+                    max_new_tokens=4096,
+                    stopping_criteria=[stopping_criteria]
+                    )
         outputs = tokenizer.decode(output_ids[0, input_ids.shape[1]:], skip_special_tokens=True)
         outputs = outputs.replace("<Number>", "")
         outputs = outputs.strip()