Upload TaiVisionForCausalLM

config.json

@@ -1,6 +1,10 @@
 {
+  "architectures": [
+    "TaiVisionForCausalLM"
+  ],
   "auto_map": {
-    "AutoConfig": "configuration_taivisionlm.TaiVisionLMConfig"
+    "AutoConfig": "configuration_taivisionlm.TaiVisionLMConfig",
+    "AutoModelForCausalLM": "modeling_taivisionlm.TaiVisionForCausalLM"
   },
   "hidden_size": 2048,
   "ignore_index": -100,
@@ -21,6 +25,7 @@
     "torch_dtype": "bfloat16",
     "vocab_size": 32001
   },
+  "torch_dtype": "float32",
   "transformers_version": "4.44.0",
   "vision_config": {
     "model_type": "siglip_vision_model",

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0e3c91245701f070448659cda849d90ef35ea419ad8aa53c459b20a7d516df00
+oid sha256:11d50e45bc0203fb3be9a06add95e21a024690098db67cd7b97f29ae03c2bb57
 size 4806424752

modeling_taivisionlm.py

@@ -156,18 +156,17 @@ class TaiVisionForCausalLM(TaiVisionPreTrainedModel):
         self.language_model = language_model
         self.post_init()
         
-    def load_pretrained(self):
-        '''
-        load the pretrained weights for language model and vision model
-        '''
-        import transformers
-        language_model = AutoModelForCausalLM.from_pretrained("benchang1110/Taiwan-tinyllama-v1.0-chat")
+    def load_language_model(self, model_id = "benchang1110/Taiwan-tinyllama-v1.0-chat"):
+        language_model = AutoModelForCausalLM.from_pretrained(model_id)
         if language_model.vocab_size != self.vocab_size:
             print("vocab size mismatch, resize the token embeddings for the pretained language model")
             language_model.resize_token_embeddings(self.vocab_size)
-        self.language_model = language_model
-        vision_model = transformers.SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224")
-        self.vision_tower = vision_model
+        self.language_model.load_state_dict(language_model.state_dict(),strict=True)
+        
+    def load_vision_model(self,model_id = "google/siglip-base-patch16-224"):
+        import transformers
+        vision_model = transformers.SiglipVisionModel.from_pretrained(model_id)
+        self.vision_tower.load_state_dict(vision_model.state_dict(),strict=True)
         
     # Copied from transformers.models.paligemma.modeling_paligemma.PaliGemmaForConditionalGeneration.get_input_embeddings with PaliGemma->TaiVisionLM
     def get_input_embeddings(self):
@@ -439,34 +438,4 @@ class TaiVisionForCausalLM(TaiVisionPreTrainedModel):
         if cache_position[0] == 0:
             model_inputs["pixel_values"] = pixel_values
 
-        return model_inputs
-    
-    
-    
-if __name__ == '__main__':
-    import transformers
-    config = transformers.AutoConfig.from_pretrained("benchang1110/TaiVision-base",trust_remote_code=True)
-    model = TaiVisionForCausalLM(config).to("cuda")
-    print(model)
-    model.save_pretrained
-    # Test forward
-    import torch
-    from PIL import Image
-    import requests
-    # Initialize processor
-    processor = transformers.AutoProcessor.from_pretrained("benchang1110/TaiVision-base", trust_remote_code=True)
-    
-    # Load image
-    url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg"
-    image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
-
-    # Define prompt and label
-    prompt = "What is the color of the car?"
-    label = "I am fine, thank you."
-
-    # Process inputs
-    inputs = processor(prompts=prompt,images=image, return_tensors="pt", padding=False, max_length=512).to('cuda')
-
-    outputs = model.generate(**inputs, max_length=512, do_sample=True, pad_token_id=processor.tokenizer.pad_token_id)
-    print(processor.decode(outputs[0], skip_special_tokens=True))
-    
+        return model_inputs