Upload model

config.json

@@ -1,13 +1,18 @@
 {
+  "architectures": [
+    "SpeechLLMModel"
+  ],
   "audio_enc_dim": 1024,
   "audio_encoder_name": "microsoft/wavlm-large",
   "audio_processor_name": "microsoft/wavlm-base",
   "auto_map": {
-    "AutoConfig": "config.SpeechLLMModelConfig"
+    "AutoConfig": "config.SpeechLLMModelConfig",
+    "AutoModel": "model.SpeechLLMModel"
   },
   "llm_dim": 2048,
   "llm_model_checkpoint": "hf_repo/llm_model_checkpoint",
   "llm_model_name": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
   "model_type": "custom_model",
+  "torch_dtype": "float32",
   "transformers_version": "4.41.2"
 }

model.py

@@ -1,14 +1,15 @@
 import torch
 from torch import nn
 import torchaudio
-from transformers import PreTrainedModel, AutoModelForCausalLM, AutoTokenizer, HubertModel, AutoFeatureExtractor, AutoModel
+from transformers import PreTrainedModel, AutoModelForCausalLM, AutoTokenizer, HubertModel, AutoProcessor, AutoConfig, AutoModel, AutoFeatureExtractor
 from .config import SpeechLLMModelConfig
 from peft import LoraConfig, get_peft_model
 
 class TransformerAudioEnoder(nn.Module):
     def __init__(self, model_name='microsoft/wavlm-large', finetune=False):
         super().__init__()
-        self.encoder = AutoModel.from_pretrained(model_name)
+        config = AutoConfig.from_pretrained(model_name)
+        self.encoder =  AutoModel.from_config(config)  
 
     def forward(self, x):
         return self.encoder(x).last_hidden_state
@@ -44,11 +45,9 @@ class SpeechLLMModel(PreTrainedModel):
         self.audio_processor = AutoFeatureExtractor.from_pretrained(config.audio_processor_name)
         self.audio_encoder = TransformerAudioEnoder(config.audio_encoder_name)
         self.connector = CNNConnector(config.audio_enc_dim, config.llm_dim)
-
-        # self.llm_model = AutoModelForCausalLM.from_pretrained(config.llm_model_checkpoint)
-        # self.llm_tokenizer = AutoTokenizer.from_pretrained(config.llm_model_name)
         
-        self.llm_model = AutoModelForCausalLM.from_pretrained(config.llm_model_name)
+        llm_config = AutoConfig.from_pretrained(config.llm_model_name)
+        self.llm_model =  AutoModelForCausalLM.from_config(llm_config)  
         self.llm_tokenizer = AutoTokenizer.from_pretrained(config.llm_model_name)
 
         peft_config = LoraConfig(