update readme.md with one sample inferrence

README.md

@@ -43,27 +43,48 @@ class StoppingCriteriaSub(StoppingCriteria):
                 return True
         return False
 
-def answer(self, samples, stopping, max_new_tokens=300, num_beams=1, min_length=1, top_p=0.5,
-        repetition_penalty=1.0, length_penalty=1, temperature=0.1, max_length=2000):
-    audio = samples["audio"].cuda()
-    audio_embeds, atts_audio = self.encode_audio(audio)
-    if 'instruction_input' in samples:  # instruction dataset
-        #print('Instruction Batch')
-        instruction_prompt = []
-        for instruction in samples['instruction_input']:
-            prompt = '<Audio><AudioHere></Audio> ' + instruction
-            instruction_prompt.append(self.prompt_template.format(prompt))
-        audio_embeds, atts_audio = self.instruction_prompt_wrap(audio_embeds, atts_audio, instruction_prompt)
-    self.llama_tokenizer.padding_side = "right"
+
+class StoppingCriteriaSub(StoppingCriteria):
+    def __init__(self, stops=[], encounters=1):
+        super().__init__()
+        self.stops = stops
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
+        for stop in self.stops:
+            if torch.all((stop == input_ids[0][-len(stop):])).item():
+                return True
+        return False
+
+def get_musilingo_pred(model, text, audio_path, stopping, length_penalty=1, temperature=0.1,
+    max_new_tokens=300, num_beams=1, min_length=1, top_p=0.5, repetition_penalty=1.0):
+
+    # see https://huggingface.co/m-a-p/MusiLingo-musicqa-v1 for load_audio function definition
+    audio = load_audio(audio_path, target_sr=24000,
+                        is_mono=True,
+                        is_normalize=False,
+                        crop_to_length_in_sample_points=int(30*16000)+1,
+                        crop_randomly=True, 
+                        pad=False).cuda()    
+    processor = Wav2Vec2FeatureExtractor.from_pretrained("m-a-p/MERT-v1-330M",trust_remote_code=True) 
+    audio = processor(audio, 
+                    sampling_rate=24000, 
+                    return_tensors="pt")['input_values'][0].cuda() 
+        
+    audio_embeds, atts_audio = model.encode_audio(audio)
+        
+    prompt = '<Audio><AudioHere></Audio> ' + text
+    instruction_prompt = [model.prompt_template.format(prompt)]
+    audio_embeds, atts_audio = model.instruction_prompt_wrap(audio_embeds, atts_audio, instruction_prompt)
+    
+    model.llama_tokenizer.padding_side = "right"
     batch_size = audio_embeds.shape[0]
     bos = torch.ones([batch_size, 1],
                     dtype=torch.long,
-                    device=torch.device('cuda')) * self.llama_tokenizer.bos_token_id
-    bos_embeds = self.llama_model.model.embed_tokens(bos)
-    atts_bos = atts_audio[:, :1]
+                    device=torch.device('cuda')) * model.llama_tokenizer.bos_token_id
+    bos_embeds = model.llama_model.model.embed_tokens(bos)
+    # atts_bos = atts_audio[:, :1]
     inputs_embeds = torch.cat([bos_embeds, audio_embeds], dim=1)
-    attention_mask = torch.cat([atts_bos, atts_audio], dim=1)
-    outputs = self.llama_model.generate(
+    # attention_mask = torch.cat([atts_bos, atts_audio], dim=1)
+    outputs = model.llama_model.generate(
         inputs_embeds=inputs_embeds,
         max_new_tokens=max_new_tokens,
         stopping_criteria=stopping,
@@ -80,34 +101,21 @@ def answer(self, samples, stopping, max_new_tokens=300, num_beams=1, min_length=
         output_token = output_token[1:]
     if output_token[0] == 1:  # if there is a start token <s> at the beginning. remove it
         output_token = output_token[1:]
-    output_text = self.llama_tokenizer.decode(output_token, add_special_tokens=False)
+    output_text = model.llama_tokenizer.decode(output_token, add_special_tokens=False)
     output_text = output_text.split('###')[0]  # remove the stop sign '###'
     output_text = output_text.split('Assistant:')[-1].strip()
     return output_text
 
-processor = Wav2Vec2FeatureExtractor.from_pretrained("m-a-p/MERT-v1-330M",trust_remote_code=True)
-ds = CMIDataset(processor, 'path/to/MI_dataset', 'test', question_type='long')
-dl = DataLoader(
-                ds,
-                batch_size=1,
-                num_workers=0,
-                pin_memory=True,
-                shuffle=False,
-                drop_last=True,
-                collate_fn=ds.collater
-                )
+musilingo = AutoModel.from_pretrained("m-a-p/MusiLingo-long-v1", trust_remote_code=True)
+musilingo.to("cuda")
+musilingo.eval()
 
+prompt = "this is the task instruction and input question for MusiLingo model"
+audio = "/path/to/the/audio"
 stopping = StoppingCriteriaList([StoppingCriteriaSub([torch.tensor([835]).cuda(),
-                                torch.tensor([2277, 29937]).cuda()])])
-
-from transformers import AutoModel
-model_long = AutoModel.from_pretrained("m-a-p/MusiLingo-long-v1")
-
-for idx, sample in tqdm(enumerate(dl)):
-    ans = answer(Musilingo_long.model, sample, stopping, length_penalty=100, temperature=0.1)
-    txt = sample['text_input'][0]
-    print(txt)
-    print(and)
+                                  torch.tensor([2277, 29937]).cuda()])])
+response = get_musilingo_pred(musilingo.model, prompt, audio_path, stopping, length_penalty=100, temperature=0.1)
+        
 ```
 
 # Citing This Work