Spaces:
Running
Running
| import torch | |
| from torch import nn | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from peft import LoraConfig, get_peft_model, PeftModel | |
| import pytorch_lightning as pl | |
| from model import HubertXCNNEnoder | |
| from torch.quantization import quantize_dynamic | |
| import torch.jit as jit | |
| if torch.cuda.is_available(): | |
| # Set the device to CUDA | |
| device = "cuda" | |
| else: | |
| # Set the device to CPU | |
| device = "cpu" | |
| class SpeechLLMLightning(pl.LightningModule): | |
| def __init__(self, audio_enc_dim=512, llm_dim=2048, llm_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0", quantize=False): | |
| super().__init__() | |
| self.save_hyperparameters() | |
| self.audio_enc_dim = audio_enc_dim | |
| self.llm_dim = llm_dim | |
| self.llm_name = llm_name | |
| self.audio_encoder = HubertXCNNEnoder(self.audio_enc_dim, self.llm_dim) | |
| self.llm_tokenizer = AutoTokenizer.from_pretrained(self.llm_name) | |
| self.llm_tokenizer.pad_token = self.llm_tokenizer.eos_token | |
| self.llm_model = AutoModelForCausalLM.from_pretrained( | |
| self.llm_name, | |
| device_map=device, | |
| ) | |
| peft_config = LoraConfig( | |
| r=4, | |
| lora_alpha=8, | |
| target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'up_proj', 'down_proj', 'gate_proj'], | |
| lora_dropout=0.05, | |
| task_type="CAUSAL_LM", | |
| ) | |
| self.llm_model = get_peft_model(self.llm_model, peft_config) | |
| self.llm_model.print_trainable_parameters() | |
| for param in self.llm_model.parameters(): | |
| param.requires_grad = False | |
| self.audio_encoder.eval() | |
| self.llm_model.eval() | |
| if quantize: | |
| # self.llm_model = jit.script(self.llm_model) | |
| self.llm_model = quantize_dynamic( | |
| self.llm_model, {nn.Linear}, dtype=torch.qint8 | |
| ) | |
| def encode(self, mel, pre_tokenized_ids, post_tokenized_ids, output_tokenized_ids): | |
| batch_size = mel.shape[0] | |
| with torch.no_grad(): | |
| speech_embeds = self.audio_encoder(mel) | |
| embedder = self.llm_model.model.model.embed_tokens | |
| pre_prompt_embeds = embedder(pre_tokenized_ids) | |
| post_prompt_embeds = embedder(post_tokenized_ids) | |
| output_prompt_embeds = embedder(output_tokenized_ids) | |
| combined_embeds = torch.cat([pre_prompt_embeds, speech_embeds, post_prompt_embeds, output_prompt_embeds], dim=1) | |
| atts = torch.ones(combined_embeds.size()[:-1], dtype=torch.long).to(combined_embeds.device) | |
| input_token_length = pre_tokenized_ids.shape[1] + speech_embeds.shape[1] + post_tokenized_ids.shape[1] | |
| label_ids = torch.cat([ | |
| torch.ones([batch_size, input_token_length], device=combined_embeds.device)*-100, | |
| output_tokenized_ids | |
| ], 1).to(combined_embeds.device).to(torch.int64) | |
| return combined_embeds, atts, label_ids | |
| def forward(self, embeds, atts, label_ids): | |
| with torch.no_grad(): | |
| return self.llm_model( | |
| inputs_embeds=embeds, | |
| attention_mask=atts, | |
| labels=label_ids, | |
| ) |