optimize

Alpha-VLLM · Aug 14, 2024 · c8e180a · c8e180a
1 parent 03e2292
commit c8e180a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/lumina_mgpt/data/item_processor.py b/lumina_mgpt/data/item_processor.py
@@ -64,7 +64,6 @@ def __init__(
         tokenizer="Alpha-VLLM/Lumina-mGPT-7B-768",
         conv_template=Conversation,
         target_size=512,
-        with_decoder=False,
     ):
 
         super().__init__(

diff --git a/lumina_mgpt/inference_solver.py b/lumina_mgpt/inference_solver.py
@@ -279,9 +279,11 @@ def __init__(self, model_path, precision, target_size=512):
         self.dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
 
         self.model = ChameleonForConditionalGeneration.from_pretrained(
-            model_path, torch_dtype=self.dtype, device_map="cuda"
+            model_path,
+            torch_dtype=self.dtype,
+            device_map="cuda",
         )
-        self.item_processor = FlexARItemProcessor(with_decoder=True, target_size=target_size)
+        self.item_processor = FlexARItemProcessor(target_size=target_size)
 
     def get_streamer(self):
         return TextStreamer(self.item_processor.tokenizer)