From fce4e1d1b48f23cd8332e60afce3df8d6209a6a7 Mon Sep 17 00:00:00 2001
From: gaochangfeng <54253717+gaochangfeng@users.noreply.github.com>
Date: 星期四, 11 四月 2024 14:59:22 +0800
Subject: [PATCH] SenseVoice对富文本解码的参数 (#1608)

---
 funasr/models/qwen_audio/model.py |   25 ++++++++++++++++++-------
 1 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/funasr/models/qwen_audio/model.py b/funasr/models/qwen_audio/model.py
index 3eba026..f981b67 100644
--- a/funasr/models/qwen_audio/model.py
+++ b/funasr/models/qwen_audio/model.py
@@ -9,8 +9,7 @@
 from torch import nn
 import whisper
 from funasr.utils.load_utils import load_audio_text_image_video, extract_fbank
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers.generation import GenerationConfig
+
 
 from funasr.register import tables
 
@@ -20,8 +19,15 @@
 @tables.register("model_classes", "QwenAudio")
 @tables.register("model_classes", "QwenAudioWarp")
 class QwenAudioWarp(nn.Module):
+    """
+    Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models
+    https://arxiv.org/abs/2311.07919
+    Modified from https://github.com/QwenLM/Qwen-Audio
+    """
     def __init__(self, *args, **kwargs):
         super().__init__()
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        from transformers.generation import GenerationConfig
 
         model_or_path = kwargs.get("model_path", "QwenAudio")
         model = AutoModelForCausalLM.from_pretrained(model_or_path, device_map="cpu",
@@ -49,14 +55,13 @@
 
         meta_data = {}
         # meta_data["batch_data_time"] = -1
-
-        sp_prompt = "<|startoftranscription|><|en|><|transcribe|><|en|><|notimestamps|><|wo_itn|>"
-        query = f"<audio>{data_in[0]}</audio>{sp_prompt}"
+        prompt = kwargs.get("prompt", "<|startoftranscription|><|en|><|transcribe|><|en|><|notimestamps|><|wo_itn|>")
+        query = f"<audio>{data_in[0]}</audio>{prompt}"
         audio_info = self.tokenizer.process_audio(query)
         inputs = self.tokenizer(query, return_tensors='pt', audio_info=audio_info)
         inputs = inputs.to(self.model.device)
         pred = self.model.generate(**inputs, audio_info=audio_info)
-        response = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False, audio_info=audio_info)
+        response = self.tokenizer.decode(pred.cpu()[0], skip_special_tokens=False, audio_info=audio_info)
 
         results = []
         result_i = {"key": key[0], "text": response}
@@ -72,8 +77,14 @@
 @tables.register("model_classes", "QwenAudioChatWarp")
 class QwenAudioChatWarp(nn.Module):
     def __init__(self, *args, **kwargs):
+        """
+        Qwen-Audio: Advancing Universal Audio Understanding via Unified Large-Scale Audio-Language Models
+        https://arxiv.org/abs/2311.07919
+        Modified from https://github.com/QwenLM/Qwen-Audio
+        """
         super().__init__()
-        
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        from transformers.generation import GenerationConfig
         model_or_path = kwargs.get("model_path", "QwenAudio")
         bf16 = kwargs.get("bf16", False)
         fp16 = kwargs.get("fp16", False)

--
Gitblit v1.9.1